From 4077e3ecb28ff60beb3985d919a1cf3bc3551631 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 22 Sep 2025 11:08:08 +0300 Subject: [PATCH] Handle title in client & CLI --- src/haiku/rag/app.py | 4 +- src/haiku/rag/cli.py | 7 +++- src/haiku/rag/client.py | 83 +++++++++++++++++++++++++++++++---------- src/haiku/rag/mcp.py | 24 +++++++++--- tests/test_app.py | 4 +- tests/test_cli.py | 15 ++++++++ tests/test_client.py | 36 ++++++++++++++++++ 7 files changed, 143 insertions(+), 30 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index c41d6267..b8818f75 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -39,9 +39,9 @@ class HaikuRAGApp: f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]" ) - async def add_document_from_source(self, source: str): + async def add_document_from_source(self, source: str, title: str | None = None): async with HaikuRAG(db_path=self.db_path) as self.client: - doc = await self.client.create_document_from_source(source) + doc = await self.client.create_document_from_source(source, title=title) self._rich_print_document(doc, truncate=True) self.console.print( f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]" diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 4fc5b144..66aa733c 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -160,6 +160,11 @@ def add_document_src( help="The file path or URL of the document to add", autocompletion=complete_local_paths, ), + title: str | None = typer.Option( + None, + "--title", + help="Optional human-readable title to store with the document", + ), db: Path = typer.Option( Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", "--db", @@ -169,7 +174,7 @@ def add_document_src( from haiku.rag.app import HaikuRAGApp app = HaikuRAGApp(db_path=db) - asyncio.run(app.add_document_from_source(source=source)) + asyncio.run(app.add_document_from_source(source=source, title=title)) @cli.command("get", help="Get and display a document by its ID") diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 9a6b89ec..c9b99b57 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -50,6 +50,7 @@ class HaikuRAG: self, docling_document, uri: str | None = None, + title: str | None = None, metadata: dict | None = None, chunks: list[Chunk] | None = None, ) -> Document: @@ -58,6 +59,7 @@ class HaikuRAG: document = Document( content=content, uri=uri, + title=title, metadata=metadata or {}, ) return await self.document_repository._create_with_docling( @@ -68,6 +70,7 @@ class HaikuRAG: self, content: str, uri: str | None = None, + title: str | None = None, metadata: dict | None = None, chunks: list[Chunk] | None = None, ) -> Document: @@ -88,6 +91,7 @@ class HaikuRAG: document = Document( content=content, uri=uri, + title=title, metadata=metadata or {}, ) return await self.document_repository._create_with_docling( @@ -95,7 +99,7 @@ class HaikuRAG: ) async def create_document_from_source( - self, source: str | Path, metadata: dict = {} + self, source: str | Path, title: str | None = None, metadata: dict | None = None ) -> Document: """Create or update a document from a file path or URL. @@ -116,11 +120,16 @@ class HaikuRAG: httpx.RequestError: If URL request fails """ + # Normalize metadata + metadata = metadata or {} + # Check if it's a URL source_str = str(source) parsed_url = urlparse(source_str) if parsed_url.scheme in ("http", "https"): - return await self._create_or_update_document_from_url(source_str, metadata) + return await self._create_or_update_document_from_url( + source_str, title=title, metadata=metadata + ) elif parsed_url.scheme == "file": # Handle file:// URI by converting to path source_path = Path(parsed_url.path) @@ -136,37 +145,51 @@ class HaikuRAG: uri = source_path.absolute().as_uri() md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest() - # Check if document already exists - existing_doc = await self.get_document_by_uri(uri) - if existing_doc and existing_doc.metadata.get("md5") == md5_hash: - # MD5 unchanged, return existing document - return existing_doc - - docling_document = FileReader.parse_file(source_path) - - # Get content type from file extension + # Get content type from file extension (do before early return) content_type, _ = mimetypes.guess_type(str(source_path)) if not content_type: content_type = "application/octet-stream" - # Merge metadata with contentType and md5 metadata.update({"contentType": content_type, "md5": md5_hash}) + # Check if document already exists + existing_doc = await self.get_document_by_uri(uri) + if existing_doc and existing_doc.metadata.get("md5") == md5_hash: + # MD5 unchanged; update title/metadata if provided + updated = False + if title is not None and title != existing_doc.title: + existing_doc.title = title + updated = True + if metadata: + existing_doc.metadata = {**(existing_doc.metadata or {}), **metadata} + updated = True + if updated: + return await self.document_repository.update(existing_doc) + return existing_doc + + # Parse file only when content changed or new document + docling_document = FileReader.parse_file(source_path) + if existing_doc: # Update existing document existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata + if title is not None: + existing_doc.title = title return await self.document_repository._update_with_docling( existing_doc, docling_document ) else: # Create new document using DoclingDocument return await self._create_document_with_docling( - docling_document=docling_document, uri=uri, metadata=metadata + docling_document=docling_document, + uri=uri, + title=title, + metadata=metadata, ) async def _create_or_update_document_from_url( - self, url: str, metadata: dict = {} + self, url: str, title: str | None = None, metadata: dict | None = None ) -> Document: """Create or update a document from a URL by downloading and parsing the content. @@ -186,20 +209,35 @@ class HaikuRAG: ValueError: If the content cannot be parsed httpx.RequestError: If URL request fails """ + metadata = metadata or {} + async with httpx.AsyncClient() as client: response = await client.get(url) response.raise_for_status() md5_hash = hashlib.md5(response.content).hexdigest() + # Get content type early (used for potential no-op update) + content_type = response.headers.get("content-type", "").lower() + # Check if document already exists existing_doc = await self.get_document_by_uri(url) if existing_doc and existing_doc.metadata.get("md5") == md5_hash: - # MD5 unchanged, return existing document + # MD5 unchanged; update title/metadata if provided + updated = False + if title is not None and title != existing_doc.title: + existing_doc.title = title + updated = True + metadata.update({"contentType": content_type, "md5": md5_hash}) + if metadata: + existing_doc.metadata = { + **(existing_doc.metadata or {}), + **metadata, + } + updated = True + if updated: + return await self.document_repository.update(existing_doc) return existing_doc - - # Get content type to determine file extension - content_type = response.headers.get("content-type", "").lower() file_extension = self._get_extension_from_content_type_or_url( url, content_type ) @@ -226,12 +264,17 @@ class HaikuRAG: if existing_doc: existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata + if title is not None: + existing_doc.title = title return await self.document_repository._update_with_docling( existing_doc, docling_document ) else: return await self._create_document_with_docling( - docling_document=docling_document, uri=url, metadata=metadata + docling_document=docling_document, + uri=url, + title=title, + metadata=metadata, ) def _get_extension_from_content_type_or_url( @@ -522,7 +565,7 @@ class HaikuRAG: # Try to re-create from source (this creates the document with chunks) new_doc = await self.create_document_from_source( - doc.uri, doc.metadata or {} + source=doc.uri, metadata=doc.metadata or {} ) assert new_doc.id is not None, "New document ID should not be None" diff --git a/src/haiku/rag/mcp.py b/src/haiku/rag/mcp.py index 16a4d4a4..654cf8e8 100644 --- a/src/haiku/rag/mcp.py +++ b/src/haiku/rag/mcp.py @@ -17,6 +17,7 @@ class DocumentResult(BaseModel): id: str | None content: str uri: str | None = None + title: str | None = None metadata: dict[str, Any] = {} created_at: str updated_at: str @@ -28,13 +29,15 @@ def create_mcp_server(db_path: Path) -> FastMCP: @mcp.tool() async def add_document_from_file( - file_path: str, metadata: dict[str, Any] | None = None + file_path: str, + metadata: dict[str, Any] | None = None, + title: str | None = None, ) -> str | None: """Add a document to the RAG system from a file path.""" try: async with HaikuRAG(db_path) as rag: document = await rag.create_document_from_source( - Path(file_path), metadata or {} + Path(file_path), title=title, metadata=metadata or {} ) return document.id except Exception: @@ -42,24 +45,31 @@ def create_mcp_server(db_path: Path) -> FastMCP: @mcp.tool() async def add_document_from_url( - url: str, metadata: dict[str, Any] | None = None + url: str, metadata: dict[str, Any] | None = None, title: str | None = None ) -> str | None: """Add a document to the RAG system from a URL.""" try: async with HaikuRAG(db_path) as rag: - document = await rag.create_document_from_source(url, metadata or {}) + document = await rag.create_document_from_source( + url, title=title, metadata=metadata or {} + ) return document.id except Exception: return None @mcp.tool() async def add_document_from_text( - content: str, uri: str | None = None, metadata: dict[str, Any] | None = None + content: str, + uri: str | None = None, + metadata: dict[str, Any] | None = None, + title: str | None = None, ) -> str | None: """Add a document to the RAG system from text content.""" try: async with HaikuRAG(db_path) as rag: - document = await rag.create_document(content, uri, metadata or {}) + document = await rag.create_document( + content, uri, title=title, metadata=metadata or {} + ) return document.id except Exception: return None @@ -102,6 +112,7 @@ def create_mcp_server(db_path: Path) -> FastMCP: id=document.id, content=document.content, uri=document.uri, + title=document.title, metadata=document.metadata, created_at=str(document.created_at), updated_at=str(document.updated_at), @@ -123,6 +134,7 @@ def create_mcp_server(db_path: Path) -> FastMCP: id=doc.id, content=doc.content, uri=doc.uri, + title=doc.title, metadata=doc.metadata, created_at=str(doc.created_at), updated_at=str(doc.updated_at), diff --git a/tests/test_app.py b/tests/test_app.py index 5e7f80d1..ce0db165 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -78,7 +78,9 @@ async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.add_document_from_source(file_path) - mock_client.create_document_from_source.assert_called_once_with(file_path) + mock_client.create_document_from_source.assert_called_once_with( + file_path, title=None + ) mock_rich_print.assert_called_once_with(mock_doc, truncate=True) mock_print.assert_called_once_with( "[b]Document with id [cyan]1[/cyan] added successfully.[/b]" diff --git a/tests/test_cli.py b/tests/test_cli.py index d4e5a455..f1ed1fa6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -45,6 +45,21 @@ def test_add_document_src(): mock_app_instance.add_document_from_source.assert_called_once() +def test_add_document_src_with_title(): + with patch("haiku.rag.cli.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.add_document_from_source = AsyncMock() + mock_app.return_value = mock_app_instance + + result = runner.invoke(cli, ["add-src", "test.txt", "--title", "Nice Name"]) + + assert result.exit_code == 0 + mock_app_instance.add_document_from_source.assert_called_once() + # Verify title is forwarded + await_args = mock_app_instance.add_document_from_source.await_args + assert await_args.kwargs.get("title") == "Nice Name" + + def test_get_document(): with patch("haiku.rag.app.HaikuRAGApp") as mock_app: mock_app_instance = MagicMock() diff --git a/tests/test_client.py b/tests/test_client.py index 60038067..da593b93 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -105,6 +105,42 @@ async def test_client_create_document_from_source(temp_db_path): assert "md5" in doc2.metadata +@pytest.mark.asyncio +async def test_client_create_document_from_source_with_title(temp_db_path): + """Test creating a document from a file source with a title.""" + async with HaikuRAG(temp_db_path) as client: + with tempfile.TemporaryDirectory() as temp_dir: + test_content = "This is test content from a file." + temp_path = Path(temp_dir) / "test_title.txt" + temp_path.write_text(test_content) + + doc = await client.create_document_from_source( + source=temp_path, title="My Doc" + ) + assert doc.id is not None + assert doc.title == "My Doc" + + +@pytest.mark.asyncio +async def test_client_update_title_noop_behavior(temp_db_path): + """When content is unchanged, updating title should update document without re-chunking.""" + async with HaikuRAG(temp_db_path) as client: + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) / "test_update_title.txt" + temp_path.write_text("Original content") + + doc1 = await client.create_document_from_source(temp_path, title="Title A") + 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 doc2.id == doc1.id + # Fetch and verify title updated + got = await client.get_document_by_id(doc1.id) + assert got is not None + assert got.title == "Title B" + + @pytest.mark.asyncio async def test_client_create_document_from_source_unsupported(temp_db_path): """Test creating a document from an unsupported file type."""