add automatic title generation for documents

haiku_rag_slim/haiku/rag/client.py

@@ -248,6 +248,95 @@ 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
+
+        # Truncate content to limit token usage
+        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."
+            ),
+        )
+        try:
+            result = await agent.run(truncated)
+            title = result.output.strip()
+            return title if title else None
+        except Exception:
+            logger.warning("LLM title generation failed", exc_info=True)
+            return None
+
+    async def _resolve_title(
+        self,
+        title: str | None,
+        docling_document: "DoclingDocument",
+        content: str,
+    ) -> str | None:
+        """Resolve the title for a document.
+
+        1. Explicit title always wins.
+        2. If auto_title is disabled, return None.
+        3. Try structural extraction from docling metadata.
+        4. Fall back to LLM generation.
+        """
+        if title is not None:
+            return title
+
+        if not self._config.processing.auto_title:
+            return None
+
+        structural = self._extract_structural_title(docling_document)
+        if structural:
+            return structural
+
+        return await self._generate_title_with_llm(content)
+
     async def _store_document_with_chunks(
         self,
         document: Document,
@@ -383,6 +472,8 @@ class HaikuRAG:
         # The original content is preserved in docling_document
         stored_content = docling_document.export_to_markdown()
 
+        title = await self._resolve_title(title, docling_document, stored_content)
+
         # Create document model
         document = Document(
             content=stored_content,
@@ -420,8 +511,11 @@ class HaikuRAG:
         Returns:
             The created Document instance.
         """
+        content = docling_document.export_to_markdown()
+        title = await self._resolve_title(title, docling_document, content)
+
         document = Document(
-            content=docling_document.export_to_markdown(),
+            content=content,
             uri=uri,
             title=title,
             metadata=metadata or {},
@@ -556,9 +650,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 +662,18 @@ 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(
+                    None, docling_document, stored_content
+                )
             return await self._update_document_with_chunks(
                 existing_doc, embedded_chunks
             )
         else:
             # Create new document
+            title = await self._resolve_title(title, docling_document, stored_content)
             document = Document(
-                content=docling_document.export_to_markdown(),
+                content=stored_content,
                 uri=uri,
                 title=title,
                 metadata=metadata,
@@ -665,9 +766,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 +778,20 @@ 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(
+                        None, docling_document, stored_content
+                    )
                 return await self._update_document_with_chunks(
                     existing_doc, embedded_chunks
                 )
             else:
                 # Create new document
+                title = await self._resolve_title(
+                    title, docling_document, stored_content
+                )
                 document = Document(
-                    content=docling_document.export_to_markdown(),
+                    content=stored_content,
                     uri=url,
                     title=title,
                     metadata=metadata,

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):

tests/test_title_generation.py

@@ -0,0 +1,258 @@
+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_explicit_title_always_wins(self, tmp_path):
+        """Caller-supplied title is never overridden."""
+        doc = DoclingDocument(name="test")
+        doc.add_text(label=DocItemLabel.TITLE, text="Structural Title")
+
+        client = self._make_client(tmp_path)
+        result = await client._resolve_title("My Explicit Title", doc, "some content")
+        assert result == "My Explicit Title"
+
+    @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(None, 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(None, doc, "some content")
+        assert result == "Auto Extracted Title"
+
+    @pytest.mark.asyncio
+    async def test_no_structural_title_no_llm_returns_none(self, tmp_path):
+        """Returns None when no structural title and no LLM available."""
+        doc = DoclingDocument(name="test")
+        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just text")
+
+        client = self._make_client(tmp_path)
+        result = await client._resolve_title(None, doc, "some content")
+        # LLM call will fail without allow_model_requests, so we get None
+        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"