Merge pull request #102 from ggozad/feat/recursive-add-src

create_document_from_source (and add-src from CLI) will recursively add directories

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",

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)
+            if isinstance(result, list):
-            self.console.print(
+                for doc in result:
-                f"[bold green]Document {doc.id} added successfully.[/bold green]"
+                    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:

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,

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:
+    ) -> Document | list[Document]:
-        """Create or update a document from a file path or URL.
+        """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"
                         )

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)

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
 

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" },