diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 9c94a3d3..146db779 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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, 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..c681893e --- /dev/null +++ b/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 ) 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"