diff --git a/packages/moss-cli/src/moss_cli/documents.py b/packages/moss-cli/src/moss_cli/documents.py index 9c0631a3..a7f94b99 100644 --- a/packages/moss-cli/src/moss_cli/documents.py +++ b/packages/moss-cli/src/moss_cli/documents.py @@ -3,10 +3,11 @@ from __future__ import annotations import csv +import io import json import sys from pathlib import Path -from typing import Any, List +from typing import Any, List, Optional import typer from moss import DocumentInfo @@ -23,7 +24,7 @@ def load_documents(file_path: str) -> List[DocumentInfo]: raise typer.BadParameter(f"File not found: {file_path}") suffix = path.suffix.lower() - content = path.read_text() + content = path.read_text(encoding="utf-8-sig") if suffix == ".csv": return _parse_csv_docs(content) @@ -67,35 +68,49 @@ def _parse_jsonl_docs(raw: str, source: str = "input") -> List[DocumentInfo]: def _parse_csv_docs(content: str) -> List[DocumentInfo]: - reader = csv.DictReader(content.splitlines()) + reader = csv.DictReader(io.StringIO(content, newline="")) + if reader.fieldnames is None: + raise typer.BadParameter("CSV is empty: expected a header row") + + reader.fieldnames = [(name or "").strip() for name in reader.fieldnames] + + seen = set() + dupes = set() + for name in reader.fieldnames: + if not name: + continue + if name in seen: + dupes.add(name) + else: + seen.add(name) + if dupes: + raise typer.BadParameter( + f"CSV header has duplicate column name(s): {', '.join(sorted(dupes))}" + ) + + missing = [name for name in ("id", "text") if name not in reader.fieldnames] + if missing: + raise typer.BadParameter( + f"CSV header is missing required column(s): {', '.join(missing)}" + ) + docs = [] - for i, row in enumerate(reader): - if "id" not in row or "text" not in row: + for row in reader: + line_no = reader.line_num + doc_id = row.get("id") + text = row.get("text") + if not doc_id or text is None: raise typer.BadParameter( - f"CSV row {i + 1}: missing required 'id' or 'text' column" + f"CSV line {line_no}: empty value in required 'id' or 'text' column" ) - metadata = None - if "metadata" in row and row["metadata"]: - try: - metadata = json.loads(row["metadata"]) - except json.JSONDecodeError: - raise typer.BadParameter( - f"CSV row {i + 1}: invalid JSON in 'metadata' column" - ) - - embedding = None - if "embedding" in row and row["embedding"]: - try: - embedding = json.loads(row["embedding"]) - except json.JSONDecodeError: - raise typer.BadParameter( - f"CSV row {i + 1}: invalid JSON in 'embedding' column" - ) + + metadata = _parse_csv_json(row.get("metadata"), "metadata", line_no) + embedding = _parse_csv_json(row.get("embedding"), "embedding", line_no) docs.append( DocumentInfo( - id=row["id"], - text=row["text"], + id=doc_id, + text=text, metadata=metadata, embedding=embedding, ) @@ -103,6 +118,20 @@ def _parse_csv_docs(content: str) -> List[DocumentInfo]: return docs +def _parse_csv_json(value: Optional[str], column: str, line_no: int) -> Any: + if value is None: + return None + value = value.strip() + if not value: + return None + try: + return json.loads(value) + except json.JSONDecodeError: + raise typer.BadParameter( + f"CSV line {line_no}: invalid JSON in '{column}' column" + ) + + def _dict_to_doc(d: Any, index: int) -> DocumentInfo: if not isinstance(d, dict): raise typer.BadParameter(f"Document at index {index}: expected object, got {type(d).__name__}") diff --git a/packages/moss-cli/tests/test_documents.py b/packages/moss-cli/tests/test_documents.py new file mode 100644 index 00000000..1c42f6b7 --- /dev/null +++ b/packages/moss-cli/tests/test_documents.py @@ -0,0 +1,52 @@ +import pytest +import typer + +from moss_cli.documents import _parse_csv_docs + + +def test_csv_keeps_newlines_inside_quoted_text() -> None: + docs = _parse_csv_docs('id,text\ndoc1,"line one\nline two"\ndoc2,plain\n') + assert len(docs) == 2 + assert docs[0].text == "line one\nline two" + assert docs[1].text == "plain" + + +def test_csv_parses_metadata_and_embedding_columns() -> None: + docs = _parse_csv_docs( + "id,text,metadata,embedding\n" + 'doc1,hello,"{""topic"": ""ml""}","[0.1, 0.2]"\n' + "doc2,world,,\n" + ) + assert docs[0].metadata == {"topic": "ml"} + assert docs[0].embedding == pytest.approx([0.1, 0.2]) + assert docs[1].metadata is None + assert docs[1].embedding is None + + +def test_csv_header_is_trimmed() -> None: + docs = _parse_csv_docs("id , text \ndoc1,hello\n") + assert docs[0].id == "doc1" + assert docs[0].text == "hello" + + +def test_csv_missing_column_reports_header_error() -> None: + with pytest.raises(typer.BadParameter) as exc: + _parse_csv_docs("id,body\ndoc1,hello\n") + assert "text" in str(exc.value) + + +def test_csv_empty_content_is_rejected() -> None: + with pytest.raises(typer.BadParameter): + _parse_csv_docs("") + + +def test_csv_short_row_is_rejected() -> None: + with pytest.raises(typer.BadParameter) as exc: + _parse_csv_docs("id,text\ndoc1\n") + assert "line 2" in str(exc.value) + + +def test_csv_invalid_metadata_reports_file_line_number() -> None: + with pytest.raises(typer.BadParameter) as exc: + _parse_csv_docs('id,text,metadata\ndoc1,"a\nb",not-json\n') + assert "line 3" in str(exc.value)