diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index acd229d4..a0213a6f 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -261,14 +261,18 @@ class HaikuRAGApp: # pragma: no cover for doc in documents: self._rich_print_document(doc, truncate=True) - async def add_document_from_text(self, text: str, metadata: dict | None = None): + async def add_document_from_text( + self, text: str, title: str | None = None, metadata: dict | None = None + ): async with HaikuRAG( db_path=self.db_path, config=self.config, read_only=self.read_only, before=self.before, ) as self.client: - doc = await self.client.create_document(text, metadata=metadata) + doc = await self.client.create_document( + text, title=title, metadata=metadata + ) self._rich_print_document(doc, truncate=True) self.console.print( f"[bold green]Document {doc.id} added successfully.[/bold green]" @@ -534,6 +538,7 @@ class HaikuRAGApp: # pragma: no cover RebuildMode.FULL: "full rebuild", RebuildMode.RECHUNK: "rechunk", RebuildMode.EMBED_ONLY: "embed only", + RebuildMode.TITLE_ONLY: "title only", }[mode] self.console.print( diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index f9afa8aa..3e758cd6 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -200,6 +200,11 @@ def add_document_text( # pragma: no cover text: str = typer.Argument( help="The text content of the document to add", ), + title: str | None = typer.Option( + None, + "--title", + help="Optional title for the document", + ), meta: list[str] | None = typer.Option( None, "--meta", @@ -214,7 +219,9 @@ def add_document_text( # pragma: no cover ): app = create_app(db) metadata = _parse_meta_options(meta) - asyncio.run(app.add_document_from_text(text=text, metadata=metadata or None)) + asyncio.run( + app.add_document_from_text(text=text, title=title, metadata=metadata or None) + ) @_cli.command("add-src", help="Add a document from a file path, directory, or URL") @@ -469,17 +476,27 @@ def rebuild( "--rechunk", help="Re-chunk from existing content without accessing source files", ), + title_only: bool = typer.Option( + False, + "--title-only", + help="Only generate titles for documents without one", + ), ): from haiku.rag.client import RebuildMode - if embed_only and rechunk: - typer.echo("Error: --embed-only and --rechunk are mutually exclusive") + exclusive = sum([embed_only, rechunk, title_only]) + if exclusive > 1: + typer.echo( + "Error: --embed-only, --rechunk, and --title-only are mutually exclusive" + ) raise typer.Exit(1) if embed_only: # pragma: no cover mode = RebuildMode.EMBED_ONLY elif rechunk: # pragma: no cover mode = RebuildMode.RECHUNK + elif title_only: # pragma: no cover + mode = RebuildMode.TITLE_ONLY else: # pragma: no cover mode = RebuildMode.FULL diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 146db779..20426bb4 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -46,6 +46,7 @@ class RebuildMode(Enum): FULL = "full" # Re-convert from source, re-chunk, re-embed RECHUNK = "rechunk" # Re-chunk from existing content, re-embed EMBED_ONLY = "embed_only" # Keep chunks, only regenerate embeddings + TITLE_ONLY = "title_only" # Only generate titles for untitled documents @dataclass @@ -291,7 +292,6 @@ class HaikuRAG: from haiku.rag.utils import get_model - # Truncate content to limit token usage truncated = content[:2000] model = get_model(self._config.processing.title_model, self._config) @@ -337,6 +337,25 @@ class HaikuRAG: return await self._generate_title_with_llm(content) + async def generate_title(self, document: Document) -> str | None: + """Generate a title for a document. + + Attempts structural extraction from the stored DoclingDocument, + then falls back to LLM generation. Bypasses the auto_title config + since this is an explicit call. + + Does NOT update the document — caller decides. + """ + docling_doc = document.get_docling_document() + content = document.content or "" + + if docling_doc is not None: + structural = self._extract_structural_title(docling_doc) + if structural: + return structural + + return await self._generate_title_with_llm(content) + async def _store_document_with_chunks( self, document: Document, @@ -1629,6 +1648,7 @@ class HaikuRAG: - FULL: Re-convert from source files, re-chunk, re-embed (default) - RECHUNK: Re-chunk from existing content, re-embed (no source access) - EMBED_ONLY: Keep existing chunks, only regenerate embeddings + - TITLE_ONLY: Only generate titles for untitled documents Yields: The ID of the document currently being processed. @@ -1639,7 +1659,10 @@ class HaikuRAG: documents = await self.list_documents(include_content=True) - if mode == RebuildMode.EMBED_ONLY: + if mode == RebuildMode.TITLE_ONLY: + async for doc_id in self._rebuild_title_only(documents): + yield doc_id + elif mode == RebuildMode.EMBED_ONLY: async for doc_id in self._rebuild_embed_only(documents): yield doc_id elif mode == RebuildMode.RECHUNK: @@ -1660,6 +1683,20 @@ class HaikuRAG: except Exception: pass + async def _rebuild_title_only( + self, documents: list[Document] + ) -> AsyncGenerator[str, None]: + """Generate titles for documents that don't have one.""" + for doc in documents: + if doc.title is not None: + continue + assert doc.id is not None + title = await self.generate_title(doc) + if title is not None: + doc.title = title + await self.document_repository.update(doc) + yield doc.id + async def _rebuild_embed_only( self, documents: list[Document] ) -> AsyncGenerator[str, None]: diff --git a/tests/test_title_generation.py b/tests/test_title_generation.py index c681893e..40474a34 100644 --- a/tests/test_title_generation.py +++ b/tests/test_title_generation.py @@ -256,3 +256,104 @@ class TestImportDocumentAutoTitle: title="Keep This Title", ) assert doc.title == "Keep This Title" + + +# ========================================================================= +# generate_title() public method +# ========================================================================= + + +class TestGenerateTitle: + @pytest.mark.asyncio + async def test_structural_title(self, temp_db_path): + """generate_title extracts structural title from document.""" + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc = await client.create_document( + "# Great Heading\n\nSome content.", uri="test://gen-title" + ) + title = await client.generate_title(doc) + assert title == "Great Heading" + + @pytest.mark.asyncio + async def test_no_structural_title_no_llm(self, temp_db_path): + """generate_title returns None when no structural title and LLM unavailable.""" + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.create_document( + "Just plain text without headings.", + uri="test://gen-no-title", + format="plain", + ) + title = await client.generate_title(doc) + # LLM blocked by ALLOW_MODEL_REQUESTS=False, so falls back to None + assert title is None + + @pytest.mark.asyncio + async def test_bypasses_auto_title_config(self, temp_db_path): + """generate_title works even when auto_title is False.""" + config = AppConfig(processing=ProcessingConfig(auto_title=False)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc = await client.create_document( + "# Heading Present\n\nBody text.", uri="test://gen-bypass" + ) + title = await client.generate_title(doc) + assert title == "Heading Present" + + +# ========================================================================= +# rebuild --title-only +# ========================================================================= + + +class TestRebuildTitleOnly: + @pytest.mark.asyncio + async def test_generates_titles_for_untitled_docs(self, temp_db_path): + """TITLE_ONLY mode generates titles for documents without one.""" + from haiku.rag.client import RebuildMode + + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc1 = await client.create_document( + "# Doc With Heading\n\nContent.", + uri="test://titled", + ) + assert doc1.title == "Doc With Heading" + + # Simulate an untitled doc + doc1.title = None + await client.document_repository.update(doc1) + + doc2 = await client.create_document( + "# Another Heading\n\nMore content.", + uri="test://also-titled", + ) + assert doc2.title == "Another Heading" + + processed_ids = [] + async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY): + processed_ids.append(doc_id) + + assert len(processed_ids) == 1 + assert doc1.id in processed_ids + + updated_doc1 = await client.get_document_by_id(doc1.id) + assert updated_doc1 is not None + assert updated_doc1.title == "Doc With Heading" + + @pytest.mark.asyncio + async def test_skips_already_titled_docs(self, temp_db_path): + """TITLE_ONLY mode skips documents that already have titles.""" + from haiku.rag.client import RebuildMode + + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + await client.create_document( + "# Has Title\n\nContent.", + uri="test://has-title", + ) + + processed_ids = [] + async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY): + processed_ids.append(doc_id) + + assert len(processed_ids) == 0