danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Simplify _resolve_title()

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-26 10:56:28 +02:00
parent d3a1031ec4
commit b5cb2852c4
No known key found for this signature in database
2 changed files with 15 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
haiku_rag_slim/haiku/rag/client.py
View file

@ -310,20 +310,13 @@ class HaikuRAG:
async def _resolve_title( async def _resolve_title(
self, self,
title: str | None,
docling_document: "DoclingDocument", docling_document: "DoclingDocument",
content: str, content: str,
) -> str | None: ) -> str | None:
"""Resolve the title for a document. """Auto-generate a title from document structure or LLM.
1. Explicit title always wins. Returns None if auto_title is disabled or generation fails.
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: if not self._config.processing.auto_title:
return None return None
@ -493,7 +486,8 @@ class HaikuRAG:
# The original content is preserved in docling_document # The original content is preserved in docling_document
stored_content = docling_document.export_to_markdown() stored_content = docling_document.export_to_markdown()
title = await self._resolve_title(title, docling_document, stored_content) if title is None:
title = await self._resolve_title(docling_document, stored_content)
# Create document model # Create document model
document = Document( document = Document(
@ -533,7 +527,8 @@ class HaikuRAG:
The created Document instance. The created Document instance.
""" """
content = docling_document.export_to_markdown() content = docling_document.export_to_markdown()
title = await self._resolve_title(title, docling_document, content) if title is None:
title = await self._resolve_title(docling_document, content)
document = Document( document = Document(
content=content, content=content,
@ -685,14 +680,15 @@ class HaikuRAG:
existing_doc.title = title existing_doc.title = title
elif existing_doc.title is None: elif existing_doc.title is None:
existing_doc.title = await self._resolve_title( existing_doc.title = await self._resolve_title(
None, docling_document, stored_content docling_document, stored_content
) )
return await self._update_document_with_chunks( return await self._update_document_with_chunks(
existing_doc, embedded_chunks existing_doc, embedded_chunks
) )
else: else:
# Create new document # Create new document
title = await self._resolve_title(title, docling_document, stored_content) if title is None:
title = await self._resolve_title(docling_document, stored_content)
document = Document( document = Document(
content=stored_content, content=stored_content,
uri=uri, uri=uri,
@ -801,16 +797,15 @@ class HaikuRAG:
existing_doc.title = title existing_doc.title = title
elif existing_doc.title is None: elif existing_doc.title is None:
existing_doc.title = await self._resolve_title( existing_doc.title = await self._resolve_title(
None, docling_document, stored_content docling_document, stored_content
) )
return await self._update_document_with_chunks( return await self._update_document_with_chunks(
existing_doc, embedded_chunks existing_doc, embedded_chunks
) )
else: else:
# Create new document # Create new document
title = await self._resolve_title( if title is None:
title, docling_document, stored_content title = await self._resolve_title(docling_document, stored_content)
)
document = Document( document = Document(
content=stored_content, content=stored_content,
uri=url, uri=url,

16
tests/test_title_generation.py
View file

@ -143,16 +143,6 @@ class TestResolveTitle:
config = AppConfig(processing=ProcessingConfig(auto_title=auto_title)) config = AppConfig(processing=ProcessingConfig(auto_title=auto_title))
return HaikuRAG(tmp_path / "test.lancedb", config=config, create=True) 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 @pytest.mark.asyncio
async def test_auto_title_disabled_returns_none(self, tmp_path): async def test_auto_title_disabled_returns_none(self, tmp_path):
"""When auto_title is False, returns None (no title generation).""" """When auto_title is False, returns None (no title generation)."""
@ -160,7 +150,7 @@ class TestResolveTitle:
doc.add_text(label=DocItemLabel.TITLE, text="Structural Title") doc.add_text(label=DocItemLabel.TITLE, text="Structural Title")
client = self._make_client(tmp_path, auto_title=False) client = self._make_client(tmp_path, auto_title=False)
result = await client._resolve_title(None, doc, "some content") result = await client._resolve_title(doc, "some content")
assert result is None assert result is None
@pytest.mark.asyncio @pytest.mark.asyncio
@ -170,7 +160,7 @@ class TestResolveTitle:
doc.add_text(label=DocItemLabel.TITLE, text="Auto Extracted Title") doc.add_text(label=DocItemLabel.TITLE, text="Auto Extracted Title")
client = self._make_client(tmp_path) client = self._make_client(tmp_path)
result = await client._resolve_title(None, doc, "some content") result = await client._resolve_title(doc, "some content")
assert result == "Auto Extracted Title" assert result == "Auto Extracted Title"
@pytest.mark.asyncio @pytest.mark.asyncio
@ -185,7 +175,7 @@ class TestResolveTitle:
raise RuntimeError("LLM is down") raise RuntimeError("LLM is down")
monkeypatch.setattr(HaikuRAG, "_generate_title_with_llm", exploding_llm) monkeypatch.setattr(HaikuRAG, "_generate_title_with_llm", exploding_llm)
result = await client._resolve_title(None, doc, "some content") result = await client._resolve_title(doc, "some content")
assert result is None assert result is None