diff --git a/CHANGELOG.md b/CHANGELOG.md index 9eec119d..f9f05db4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,13 @@ # Changelog ## [Unreleased] +### Added + +- **Automatic title generation**: Documents can now have titles auto-generated during ingestion via `processing.auto_title: true`. Uses two-tier extraction: structural metadata from DoclingDocument (HTML ``, h1, section headers) first, with LLM fallback via configurable `processing.title_model` +- **`generate_title()`**: Public method on `HaikuRAG` to generate a title for an existing document on demand +- **`rebuild --title-only`**: New rebuild mode that generates titles only for untitled documents without re-chunking or re-embedding +- **`add --title`**: CLI option to set a title when adding text documents + ## [0.32.0] - 2026-02-24 ### Changed diff --git a/docs/cli.md b/docs/cli.md index 4622be47..38f0f390 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -50,6 +50,9 @@ From text: ```bash haiku-rag add "Your document content here" +# Set a title +haiku-rag add "Your document content here" --title "My Document" + # Attach metadata (repeat --meta for multiple entries) haiku-rag add "Your document content here" --meta author=alice --meta topic=notes ``` @@ -391,6 +394,9 @@ haiku-rag rebuild --rechunk # Only regenerate embeddings (fastest, keeps existing chunks) haiku-rag rebuild --embed-only + +# Only generate titles for untitled documents +haiku-rag rebuild --title-only ``` **Rebuild modes:** @@ -400,6 +406,7 @@ haiku-rag rebuild --embed-only | Full | (default) | Changed converter, source files updated | | Rechunk | `--rechunk` | Changed chunking strategy or chunk size | | Embed only | `--embed-only` | Changed embedding model or vector dimensions | +| Title only | `--title-only` | Generate titles for documents without one | ### Download Models diff --git a/docs/configuration/index.md b/docs/configuration/index.md index 3b313216..1754349f 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -117,6 +117,11 @@ processing: chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B" chunking_merge_peers: true chunking_use_markdown_tables: false + auto_title: false # Auto-generate titles on ingestion + title_model: + provider: ollama + name: gpt-oss + enable_thinking: false conversion_options: do_ocr: true force_ocr: false diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index 938e1b7e..0dae70ba 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -21,6 +21,13 @@ processing: chunking_merge_peers: true # Merge undersized successive chunks chunking_use_markdown_tables: false # Use markdown tables vs narrative format + # Automatic title generation + auto_title: false # Auto-generate titles on ingestion + title_model: # LLM for title generation (fallback) + provider: ollama + name: gpt-oss + enable_thinking: false + # Conversion options (works with both local and remote converters) conversion_options: # OCR settings @@ -201,6 +208,30 @@ picture_description: See [VLM Picture Description with docling-serve](../remote-processing.md#vlm-picture-description-with-docling-serve) for a complete example. +### Automatic Title Generation + +Enable automatic title generation during document ingestion: + +```yaml +processing: + auto_title: true + title_model: + provider: ollama + name: gpt-oss + enable_thinking: false +``` + +When `auto_title` is enabled, haiku.rag attempts to extract a title for each document during ingestion using a two-tier approach: + +1. **Structural extraction** (free, no model calls): Scans the DoclingDocument for semantic labels — HTML `<title>` tags, `<h1>` headings, PDF title blocks, and section headers +2. **LLM fallback**: When no structural title is found (e.g., plain text), generates a title using the configured `title_model` + +Priority order: HTML `<title>` (furniture layer) → h1/PDF title (body layer) → first section header → LLM generation. + +Explicit titles passed via `title=` parameter always take precedence and are never overridden. When updating documents, existing titles are preserved — auto-generation only applies to untitled documents. + +To generate titles for existing untitled documents, use [`rebuild --title-only`](../cli.md#rebuild-database). + ### Local vs Remote Processing **Local processing** (default): diff --git a/docs/python.md b/docs/python.md index 9a21508e..3035fb73 100644 --- a/docs/python.md +++ b/docs/python.md @@ -229,6 +229,28 @@ async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY): - `RebuildMode.FULL` - Re-convert from source files, re-chunk, re-embed (default) - `RebuildMode.RECHUNK` - Re-chunk from existing document content, re-embed - `RebuildMode.EMBED_ONLY` - Keep existing chunks, only regenerate embeddings +- `RebuildMode.TITLE_ONLY` - Generate titles for untitled documents (no re-chunking or re-embedding) + +### Generating Titles + +Generate a title for an existing document on demand: + +```python +title = await client.generate_title(doc) +if title: + await client.update_document(document_id=doc.id, title=title) +``` + +Uses the same two-tier approach as automatic ingestion: structural extraction from DoclingDocument metadata first, with LLM fallback via `processing.title_model`. Unlike ingestion, this method does not catch exceptions — if the LLM call fails, the error propagates. + +To batch-generate titles for all untitled documents, use `RebuildMode.TITLE_ONLY`: + +```python +async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY): + print(f"Generated title for {doc_id}") +``` + +See [Automatic Title Generation](configuration/processing.md#automatic-title-generation) for configuration details. ## Maintenance 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 9c94a3d3..15db857e 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 @@ -248,6 +249,108 @@ class HaikuRAG: return result + # ========================================================================= + # Title Generation + # ========================================================================= + + def _extract_structural_title( + self, docling_document: "DoclingDocument" + ) -> str | None: + """Extract a title from DoclingDocument structural metadata. + + Priority: FURNITURE TITLE > BODY TITLE > first SECTION_HEADER. + """ + from docling_core.types.doc.document import ContentLayer + from docling_core.types.doc.labels import DocItemLabel + + furniture_title = None + body_title = None + first_section_header = None + + for item in docling_document.texts: + if item.label == DocItemLabel.TITLE: + text = item.text.strip() + if not text: + continue + if item.content_layer == ContentLayer.FURNITURE: + furniture_title = text + elif body_title is None: + body_title = text + elif ( + item.label == DocItemLabel.SECTION_HEADER + and first_section_header is None + ): + text = item.text.strip() + if text: + first_section_header = text + + return furniture_title or body_title or first_section_header + + async def _generate_title_with_llm(self, content: str) -> str | None: + """Generate a title using LLM from document content.""" + from pydantic_ai import Agent + + from haiku.rag.utils import get_model + + truncated = content[:2000] + + model = get_model(self._config.processing.title_model, self._config) + agent: Agent[None, str] = Agent( + model=model, + output_type=str, + instructions=( + "Generate a concise, descriptive title for the following document. " + "The title should be at most 10 words. " + "Return ONLY the title text, nothing else." + ), + ) + result = await agent.run(truncated) + title = result.output.strip() + return title if title else None + + async def _resolve_title( + self, + docling_document: "DoclingDocument", + content: str, + ) -> str | None: + """Auto-generate a title from document structure or LLM. + + Returns None if auto_title is disabled or generation fails. + """ + if not self._config.processing.auto_title: + return None + + structural = self._extract_structural_title(docling_document) + if structural: + return structural + + try: + return await self._generate_title_with_llm(content) + except Exception: + logger.warning( + "LLM title generation failed during ingestion", exc_info=True + ) + return None + + 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, @@ -383,6 +486,9 @@ class HaikuRAG: # The original content is preserved in docling_document stored_content = docling_document.export_to_markdown() + if title is None: + title = await self._resolve_title(docling_document, stored_content) + # Create document model document = Document( content=stored_content, @@ -420,8 +526,12 @@ class HaikuRAG: Returns: The created Document instance. """ + content = docling_document.export_to_markdown() + if title is None: + title = await self._resolve_title(docling_document, content) + document = Document( - content=docling_document.export_to_markdown(), + content=content, uri=uri, title=title, metadata=metadata or {}, @@ -556,9 +666,11 @@ class HaikuRAG: chunks = await self.chunk(docling_document) embedded_chunks = await embed_chunks(chunks, self._config) + stored_content = docling_document.export_to_markdown() + if existing_doc: # Update existing document and rechunk - existing_doc.content = docling_document.export_to_markdown() + existing_doc.content = stored_content existing_doc.metadata = metadata existing_doc.docling_document = compress_json( docling_document.model_dump_json() @@ -566,13 +678,19 @@ class HaikuRAG: existing_doc.docling_version = docling_document.version if title is not None: existing_doc.title = title + elif existing_doc.title is None: + existing_doc.title = await self._resolve_title( + docling_document, stored_content + ) return await self._update_document_with_chunks( existing_doc, embedded_chunks ) else: # Create new document + if title is None: + title = await self._resolve_title(docling_document, stored_content) document = Document( - content=docling_document.export_to_markdown(), + content=stored_content, uri=uri, title=title, metadata=metadata, @@ -665,9 +783,11 @@ class HaikuRAG: # Merge metadata with contentType and md5 metadata.update({"contentType": content_type, "md5": md5_hash}) + stored_content = docling_document.export_to_markdown() + if existing_doc: # Update existing document and rechunk - existing_doc.content = docling_document.export_to_markdown() + existing_doc.content = stored_content existing_doc.metadata = metadata existing_doc.docling_document = compress_json( docling_document.model_dump_json() @@ -675,13 +795,19 @@ class HaikuRAG: existing_doc.docling_version = docling_document.version if title is not None: existing_doc.title = title + elif existing_doc.title is None: + existing_doc.title = await self._resolve_title( + docling_document, stored_content + ) return await self._update_document_with_chunks( existing_doc, embedded_chunks ) else: # Create new document + if title is None: + title = await self._resolve_title(docling_document, stored_content) document = Document( - content=docling_document.export_to_markdown(), + content=stored_content, uri=url, title=title, metadata=metadata, @@ -1519,6 +1645,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. @@ -1529,7 +1656,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: @@ -1550,6 +1680,26 @@ 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 + try: + title = await self.generate_title(doc) + except Exception: + logger.warning( + "Failed to generate title for document %s", doc.id, exc_info=True + ) + continue + 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]: @@ -1884,6 +2034,11 @@ class HaikuRAG: pic_desc = self._config.processing.conversion_options.picture_description if pic_desc.enabled and pic_desc.model.provider == "ollama": required_models.add(pic_desc.model.name) + if ( + self._config.processing.auto_title + and self._config.processing.title_model.provider == "ollama" + ): + required_models.add(self._config.processing.title_model.name) if not required_models: return diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 9e78d142..a07c857f 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -156,6 +156,14 @@ class ProcessingConfig(BaseModel): chunking_merge_peers: bool = True chunking_use_markdown_tables: bool = False conversion_options: ConversionOptions = Field(default_factory=ConversionOptions) + auto_title: bool = False + title_model: ModelConfig = Field( + default_factory=lambda: ModelConfig( + provider="ollama", + name="gpt-oss", + enable_thinking=False, + ) + ) class SearchConfig(BaseModel): diff --git a/tests/test_title_generation.py b/tests/test_title_generation.py new file mode 100644 index 00000000..7411c674 --- /dev/null +++ b/tests/test_title_generation.py @@ -0,0 +1,393 @@ +import random + +import pytest +from docling_core.types.doc.document import ContentLayer, DoclingDocument +from docling_core.types.doc.labels import DocItemLabel + +from haiku.rag.client import HaikuRAG +from haiku.rag.config import AppConfig +from haiku.rag.config.models import ProcessingConfig +from haiku.rag.embeddings import EmbedderWrapper + + +@pytest.fixture(autouse=True) +def mock_embedder(monkeypatch): + """Monkeypatch the embedder to return deterministic vectors.""" + + async def fake_embed_query(self, text): + random.seed(hash(text) % (2**32)) + return [random.random() for _ in range(2560)] + + async def fake_embed_documents(self, texts): + result = [] + for t in texts: + random.seed(hash(t) % (2**32)) + result.append([random.random() for _ in range(2560)]) + return result + + monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query) + monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents) + + +# ========================================================================= +# Structural title extraction +# ========================================================================= + + +class TestExtractStructuralTitle: + def _make_client(self, tmp_path): + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + return HaikuRAG(tmp_path / "test.lancedb", config=config, create=True) + + def test_furniture_title(self, tmp_path): + """TITLE on FURNITURE layer (HTML <title>) is extracted.""" + doc = DoclingDocument(name="test") + doc.add_text( + label=DocItemLabel.TITLE, + text="Website Page Title", + content_layer=ContentLayer.FURNITURE, + ) + doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text") + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result == "Website Page Title" + + def test_body_title(self, tmp_path): + """TITLE on BODY layer (h1, PDF title) is extracted.""" + doc = DoclingDocument(name="test") + doc.add_text( + label=DocItemLabel.TITLE, + text="Document Heading", + content_layer=ContentLayer.BODY, + ) + doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text") + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result == "Document Heading" + + def test_section_header_fallback(self, tmp_path): + """First SECTION_HEADER is used when no TITLE exists.""" + doc = DoclingDocument(name="test") + doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Introduction") + doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Background") + doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text") + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result == "Introduction" + + def test_no_title_or_headers(self, tmp_path): + """Returns None when no TITLE or SECTION_HEADER exists.""" + doc = DoclingDocument(name="test") + doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just a paragraph") + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result is None + + def test_furniture_title_preferred_over_body_title(self, tmp_path): + """FURNITURE TITLE takes priority over BODY TITLE.""" + doc = DoclingDocument(name="test") + doc.add_text( + label=DocItemLabel.TITLE, + text="Body H1 Title", + content_layer=ContentLayer.BODY, + ) + doc.add_text( + label=DocItemLabel.TITLE, + text="HTML Page Title", + content_layer=ContentLayer.FURNITURE, + ) + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result == "HTML Page Title" + + def test_whitespace_stripped(self, tmp_path): + """Whitespace is stripped from extracted titles.""" + doc = DoclingDocument(name="test") + doc.add_text( + label=DocItemLabel.TITLE, + text=" Padded Title ", + content_layer=ContentLayer.BODY, + ) + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result == "Padded Title" + + def test_empty_title_text_skipped(self, tmp_path): + """Empty or whitespace-only TITLE text is skipped.""" + doc = DoclingDocument(name="test") + doc.add_text( + label=DocItemLabel.TITLE, + text=" ", + content_layer=ContentLayer.BODY, + ) + doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Actual Heading") + + client = self._make_client(tmp_path) + result = client._extract_structural_title(doc) + assert result == "Actual Heading" + + +# ========================================================================= +# _resolve_title +# ========================================================================= + + +class TestResolveTitle: + def _make_client(self, tmp_path, auto_title=True): + config = AppConfig(processing=ProcessingConfig(auto_title=auto_title)) + return HaikuRAG(tmp_path / "test.lancedb", config=config, create=True) + + @pytest.mark.asyncio + async def test_auto_title_disabled_returns_none(self, tmp_path): + """When auto_title is False, returns None (no title generation).""" + doc = DoclingDocument(name="test") + doc.add_text(label=DocItemLabel.TITLE, text="Structural Title") + + client = self._make_client(tmp_path, auto_title=False) + result = await client._resolve_title(doc, "some content") + assert result is None + + @pytest.mark.asyncio + async def test_structural_title_extracted(self, tmp_path): + """Structural title is extracted when auto_title is enabled.""" + doc = DoclingDocument(name="test") + doc.add_text(label=DocItemLabel.TITLE, text="Auto Extracted Title") + + client = self._make_client(tmp_path) + result = await client._resolve_title(doc, "some content") + assert result == "Auto Extracted Title" + + @pytest.mark.asyncio + async def test_llm_failure_returns_none(self, tmp_path, monkeypatch): + """LLM failure during ingestion returns None instead of raising.""" + doc = DoclingDocument(name="test") + doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just text") + + client = self._make_client(tmp_path) + + async def exploding_llm(self, content): + raise RuntimeError("LLM is down") + + monkeypatch.setattr(HaikuRAG, "_generate_title_with_llm", exploding_llm) + result = await client._resolve_title(doc, "some content") + assert result is None + + +# ========================================================================= +# Integration: create_document with auto_title +# ========================================================================= + + +class TestCreateDocumentAutoTitle: + @pytest.mark.asyncio + async def test_auto_title_from_structural(self, temp_db_path): + """create_document with auto_title=True extracts title from docling.""" + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc = await client.create_document( + "# My Document\n\nSome content here.", uri="test://auto-title" + ) + assert doc.title == "My Document" + + @pytest.mark.asyncio + async def test_auto_title_disabled(self, temp_db_path): + """create_document with auto_title=False leaves title as None.""" + config = AppConfig(processing=ProcessingConfig(auto_title=False)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc = await client.create_document( + "# My Document\n\nSome content here.", uri="test://no-auto-title" + ) + assert doc.title is None + + @pytest.mark.asyncio + async def test_explicit_title_not_overridden(self, temp_db_path): + """Explicit title is never overridden by auto-generation.""" + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + doc = await client.create_document( + "# Auto Title\n\nSome content here.", + uri="test://explicit-title", + title="My Explicit Title", + ) + assert doc.title == "My Explicit Title" + + +# ========================================================================= +# Integration: import_document with auto_title +# ========================================================================= + + +class TestImportDocumentAutoTitle: + @pytest.mark.asyncio + async def test_auto_title_from_structural(self, temp_db_path): + """import_document with auto_title=True extracts title from docling.""" + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + docling_doc = await client.convert("# Imported Doc\n\nContent here.") + chunks = await client.chunk(docling_doc) + doc = await client.import_document( + docling_doc, chunks, uri="test://import-auto-title" + ) + assert doc.title == "Imported Doc" + + @pytest.mark.asyncio + async def test_explicit_title_preserved(self, temp_db_path): + """import_document explicit title is not overridden.""" + config = AppConfig(processing=ProcessingConfig(auto_title=True)) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + docling_doc = await client.convert("# Auto Title\n\nContent here.") + chunks = await client.chunk(docling_doc) + doc = await client.import_document( + docling_doc, + chunks, + uri="test://import-explicit", + 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 raises 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", + ) + with pytest.raises(RuntimeError): + await client.generate_title(doc) + + @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 + + @pytest.mark.asyncio + async def test_continues_on_per_document_failure(self, temp_db_path, monkeypatch): + """TITLE_ONLY mode continues when generate_title fails for a document.""" + 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( + "# First Heading\n\nContent.", + uri="test://first", + ) + doc2 = await client.create_document( + "# Second Heading\n\nContent.", + uri="test://second", + ) + # Clear both titles + doc1.title = None + await client.document_repository.update(doc1) + doc2.title = None + await client.document_repository.update(doc2) + + # Make generate_title fail for the first doc, succeed for the second + original = HaikuRAG.generate_title + call_count = 0 + + async def flaky_generate(self, document): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("LLM failed") + return await original(self, document) + + monkeypatch.setattr(HaikuRAG, "generate_title", flaky_generate) + + processed_ids = [] + async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY): + processed_ids.append(doc_id) + + # Only the second doc should have been processed + assert len(processed_ids) == 1