add automatic title generation for documents
This commit is contained in:
Yiorgis Gozadinos
parent
2acb0292e4
commit
ae3469b282
3 changed files with 381 additions and 5 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
258
tests/test_title_generation.py
Normal file
258
tests/test_title_generation.py
Normal file
|
|
@ -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"
|
||||
Loading…
Reference in a new issue