From 1cab0339ad2f8840d0a936e798c821612d943520 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 16 Oct 2025 09:55:50 +0300 Subject: [PATCH] add-src in cli & create_document_from_source will recursively add directories --- pyproject.toml | 2 +- src/haiku/rag/app.py | 17 ++++++++++---- src/haiku/rag/cli.py | 4 ++-- src/haiku/rag/client.py | 51 ++++++++++++++++++++++++++++++++++++---- tests/test_cli.py | 23 ++++++++++++++++++ tests/test_client.py | 52 +++++++++++++++++++++++++++++++++++++++++ uv.lock | 2 +- 7 files changed, 138 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f04b8c82..81d3e761 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ dev = [ "mkdocs-material>=9.6.14", "pydantic-evals>=1.0.8", "pre-commit>=4.2.0", - "pyright>=1.1.405", + "pyright>=1.1.406", "pytest>=8.4.2", "pytest-asyncio>=1.2.0", "pytest-cov>=7.0.0", diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 8e7c0310..f701eeec 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -160,13 +160,20 @@ class HaikuRAGApp: self, source: str, title: str | None = None, metadata: dict | None = None ): async with HaikuRAG(db_path=self.db_path) as self.client: - doc = await self.client.create_document_from_source( + result = await self.client.create_document_from_source( source, title=title, metadata=metadata ) - self._rich_print_document(doc, truncate=True) - self.console.print( - f"[bold green]Document {doc.id} added successfully.[/bold green]" - ) + if isinstance(result, list): + for doc in result: + self._rich_print_document(doc, truncate=True) + self.console.print( + f"[bold green]{len(result)} documents added successfully.[/bold green]" + ) + else: + self._rich_print_document(result, truncate=True) + self.console.print( + f"[bold green]Document {result.id} added successfully.[/bold green]" + ) async def get_document(self, doc_id: str): async with HaikuRAG(db_path=self.db_path) as self.client: diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 8ad370e5..7394585a 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -128,10 +128,10 @@ def add_document_text( asyncio.run(app.add_document_from_text(text=text, metadata=metadata or None)) -@cli.command("add-src", help="Add a document from a file path or URL") +@cli.command("add-src", help="Add a document from a file path, directory, or URL") def add_document_src( source: str = typer.Argument( - help="The file path or URL of the document to add", + help="The file path, directory, or URL of the document(s) to add", ), title: str | None = typer.Option( None, diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 193025b9..686d7414 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -106,8 +106,8 @@ class HaikuRAG: async def create_document_from_source( self, source: str | Path, title: str | None = None, metadata: dict | None = None - ) -> Document: - """Create or update a document from a file path or URL. + ) -> Document | list[Document]: + """Create or update document(s) from a file path, directory, or URL. Checks if a document with the same URI already exists: - If MD5 is unchanged, returns existing document @@ -115,11 +115,13 @@ class HaikuRAG: - If no document exists, creates a new one Args: - source: File path (as string or Path) or URL to parse + source: File path, directory (as string or Path), or URL to parse + title: Optional title (only used for single files, not directories) metadata: Optional metadata dictionary Returns: - Document instance (created, updated, or existing) + Document instance (created, updated, or existing) for single files/URLs + List of Document instances for directories Raises: ValueError: If the file/URL cannot be parsed or doesn't exist @@ -142,6 +144,45 @@ class HaikuRAG: else: # Handle as regular file path source_path = Path(source) if isinstance(source, str) else source + + # Handle directories + if source_path.is_dir(): + documents = [] + supported_extensions = set(FileReader.extensions) + for file_path in source_path.rglob("*"): + if ( + file_path.is_file() + and file_path.suffix.lower() in supported_extensions + ): + doc = await self._create_document_from_file( + file_path, title=None, metadata=metadata + ) + documents.append(doc) + return documents + + # Handle single file + return await self._create_document_from_file( + source_path, title=title, metadata=metadata + ) + + async def _create_document_from_file( + self, source_path: Path, title: str | None = None, metadata: dict | None = None + ) -> Document: + """Create or update a document from a single file path. + + Args: + source_path: Path to the file + title: Optional title + metadata: Optional metadata dictionary + + Returns: + Document instance (created, updated, or existing) + + Raises: + ValueError: If the file cannot be parsed or doesn't exist + """ + metadata = metadata or {} + if source_path.suffix.lower() not in FileReader.extensions: raise ValueError(f"Unsupported file extension: {source_path.suffix}") @@ -592,6 +633,8 @@ class HaikuRAG: new_doc = await self.create_document_from_source( source=doc.uri, metadata=doc.metadata or {} ) + # URIs always point to single files/URLs, never directories + assert isinstance(new_doc, Document) assert new_doc.id is not None, ( "New document ID should not be None" ) diff --git a/tests/test_cli.py b/tests/test_cli.py index a428b5ad..de173f4a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -372,3 +372,26 @@ def test_info(): assert result.exit_code == 0 mock_app_instance.info.assert_called_once() + + +def test_add_document_src_directory(tmp_path): + """Test adding documents from a directory recursively.""" + with patch("haiku.rag.app.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.add_document_from_source = AsyncMock() + mock_app.return_value = mock_app_instance + + test_dir = tmp_path / "test_docs" + test_dir.mkdir() + (test_dir / "doc1.txt").write_text("doc1") + (test_dir / "doc2.md").write_text("doc2") + subdir = test_dir / "subdir" + subdir.mkdir() + (subdir / "doc3.pdf").write_text("doc3") + + result = runner.invoke(cli, ["add-src", str(test_dir)]) + + assert result.exit_code == 0 + mock_app_instance.add_document_from_source.assert_called_once() + call_args = mock_app_instance.add_document_from_source.call_args + assert call_args[1]["source"] == str(test_dir) diff --git a/tests/test_client.py b/tests/test_client.py index 6d33ac40..25bd1249 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -9,6 +9,7 @@ from datasets import Dataset from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.store.models.chunk import Chunk +from haiku.rag.store.models.document import Document @pytest.mark.asyncio @@ -88,6 +89,7 @@ async def test_client_create_document_from_source(temp_db_path): # Test create_document_from_source with Path doc = await client.create_document_from_source(source=temp_path) + assert isinstance(doc, Document) assert doc.id is not None assert doc.content == test_content @@ -98,6 +100,7 @@ async def test_client_create_document_from_source(temp_db_path): # Test create_document_from_source with string path doc2 = await client.create_document_from_source(source=str(temp_path)) + assert isinstance(doc2, Document) assert doc2.id is not None assert doc2.content == test_content @@ -118,6 +121,7 @@ async def test_client_create_document_from_source_with_title(temp_db_path): doc = await client.create_document_from_source( source=temp_path, title="My Doc" ) + assert isinstance(doc, Document) assert doc.id is not None assert doc.title == "My Doc" @@ -131,10 +135,12 @@ async def test_client_update_title_noop_behavior(temp_db_path): temp_path.write_text("Original content") doc1 = await client.create_document_from_source(temp_path, title="Title A") + assert isinstance(doc1, Document) assert doc1.id is not None # Re-add with same content but new title doc2 = await client.create_document_from_source(temp_path, title="Title B") + assert isinstance(doc2, Document) assert doc2.id == doc1.id # Fetch and verify title updated got = await client.get_document_by_id(doc1.id) @@ -169,6 +175,41 @@ async def test_client_create_document_from_source_nonexistent(temp_db_path): await client.create_document_from_source(non_existent_path) +@pytest.mark.asyncio +async def test_client_create_document_from_directory(temp_db_path): + """Test creating documents from a directory recursively.""" + async with HaikuRAG(temp_db_path) as client: + with tempfile.TemporaryDirectory() as temp_dir: + test_dir = Path(temp_dir) / "test_docs" + test_dir.mkdir() + + (test_dir / "doc1.txt").write_text("Content of doc1") + (test_dir / "doc2.md").write_text("# Content of doc2") + + subdir = test_dir / "subdir" + subdir.mkdir() + (subdir / "doc3.py").write_text("print('hello')") + + (test_dir / "unsupported.xyz").write_text("unsupported file") + + result = await client.create_document_from_source(test_dir) + + assert isinstance(result, list) + assert len(result) == 3 + + for doc in result: + assert doc.id is not None + assert doc.uri is not None + assert "md5" in doc.metadata + assert "contentType" in doc.metadata + + uris = [doc.uri for doc in result if doc.uri] + assert any("doc1.txt" in uri for uri in uris) + assert any("doc2.md" in uri for uri in uris) + assert any("doc3.py" in uri for uri in uris) + assert not any("unsupported.xyz" in uri for uri in uris) + + @pytest.mark.asyncio async def test_client_create_document_from_url(temp_db_path): """Test creating a document from a URL.""" @@ -183,6 +224,7 @@ async def test_client_create_document_from_url(temp_db_path): doc = await client.create_document_from_source( source="https://example.com/test.html", metadata={"source_type": "web"} ) + assert isinstance(doc, Document) assert doc.id is not None assert "Test Page" in doc.content @@ -212,6 +254,7 @@ async def test_client_create_document_from_url_with_different_content_types( doc = await client.create_document_from_source( "https://api.example.com/data.json" ) + assert isinstance(doc, Document) assert doc.id is not None assert "Test JSON" in doc.content @@ -230,6 +273,7 @@ async def test_client_create_document_from_url_with_different_content_types( doc = await client.create_document_from_source( "https://example.com/readme.txt" ) + assert isinstance(doc, Document) assert doc.id is not None assert doc.content == "This is plain text content from a URL." @@ -333,6 +377,7 @@ async def test_client_metadata_content_type_and_md5(temp_db_path): temp_path.write_text(test_content) doc = await client.create_document_from_source(temp_path) + assert isinstance(doc, Document) assert doc.metadata["contentType"] == "text/plain" assert doc.metadata["md5"] == expected_md5 @@ -346,6 +391,7 @@ async def test_client_metadata_content_type_and_md5(temp_db_path): url_doc = await client.create_document_from_source( "https://example.com/test.txt" ) + assert isinstance(url_doc, Document) assert url_doc.metadata["contentType"] == "text/plain" assert url_doc.metadata["md5"] == expected_md5 @@ -363,12 +409,14 @@ async def test_client_create_update_no_op_behavior(temp_db_path): # First call - should create new document doc1 = await client.create_document_from_source(temp_path) + assert isinstance(doc1, Document) assert doc1.id is not None assert doc1.content == test_content original_id = doc1.id # Second call with same content - should return existing document (no-op) doc2 = await client.create_document_from_source(temp_path) + assert isinstance(doc2, Document) assert doc2.id == original_id # Same document assert doc2.content == test_content @@ -378,6 +426,7 @@ async def test_client_create_update_no_op_behavior(temp_db_path): # Third call with changed content - should update existing document doc3 = await client.create_document_from_source(temp_path) + assert isinstance(doc3, Document) assert doc3.id == original_id # Same document ID assert doc3.content == updated_content # Updated content @@ -404,11 +453,13 @@ async def test_client_url_create_update_no_op_behavior(temp_db_path): with patch("httpx.AsyncClient.get", return_value=mock_response1): # First call - should create new document doc1 = await client.create_document_from_source(url) + assert isinstance(doc1, Document) assert doc1.id is not None original_id = doc1.id # Second call with same content - should return existing document (no-op) doc2 = await client.create_document_from_source(url) + assert isinstance(doc2, Document) assert doc2.id == original_id # Same document mock_response2 = AsyncMock() @@ -419,6 +470,7 @@ async def test_client_url_create_update_no_op_behavior(temp_db_path): with patch("httpx.AsyncClient.get", return_value=mock_response2): # Third call with changed content - should update existing document doc3 = await client.create_document_from_source(url) + assert isinstance(doc3, Document) assert doc3.id == original_id # Same document ID assert doc3.content == updated_content.decode() # Updated content diff --git a/uv.lock b/uv.lock index cbbddc89..1121e4c4 100644 --- a/uv.lock +++ b/uv.lock @@ -1194,7 +1194,7 @@ dev = [ { name = "mkdocs-material", specifier = ">=9.6.14" }, { name = "pre-commit", specifier = ">=4.2.0" }, { name = "pydantic-evals", specifier = ">=1.0.8" }, - { name = "pyright", specifier = ">=1.1.405" }, + { name = "pyright", specifier = ">=1.1.406" }, { name = "pytest", specifier = ">=8.4.2" }, { name = "pytest-asyncio", specifier = ">=1.2.0" }, { name = "pytest-cov", specifier = ">=7.0.0" },