373 lines
15 KiB
Python
373 lines
15 KiB
Python
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.client.titles import extract_structural_title, resolve_title
|
|
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 test_furniture_title(self):
|
|
"""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")
|
|
|
|
assert extract_structural_title(doc) == "Website Page Title"
|
|
|
|
def test_body_title(self):
|
|
"""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")
|
|
|
|
assert extract_structural_title(doc) == "Document Heading"
|
|
|
|
def test_section_header_fallback(self):
|
|
"""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")
|
|
|
|
assert extract_structural_title(doc) == "Introduction"
|
|
|
|
def test_no_title_or_headers(self):
|
|
"""Returns None when no TITLE or SECTION_HEADER exists."""
|
|
doc = DoclingDocument(name="test")
|
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just a paragraph")
|
|
|
|
assert extract_structural_title(doc) is None
|
|
|
|
def test_furniture_title_preferred_over_body_title(self):
|
|
"""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,
|
|
)
|
|
|
|
assert extract_structural_title(doc) == "HTML Page Title"
|
|
|
|
def test_whitespace_stripped(self):
|
|
"""Whitespace is stripped from extracted titles."""
|
|
doc = DoclingDocument(name="test")
|
|
doc.add_text(
|
|
label=DocItemLabel.TITLE,
|
|
text=" Padded Title ",
|
|
content_layer=ContentLayer.BODY,
|
|
)
|
|
|
|
assert extract_structural_title(doc) == "Padded Title"
|
|
|
|
def test_empty_title_text_skipped(self):
|
|
"""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")
|
|
|
|
assert extract_structural_title(doc) == "Actual Heading"
|
|
|
|
|
|
# =========================================================================
|
|
# resolve_title
|
|
# =========================================================================
|
|
|
|
|
|
class TestResolveTitle:
|
|
@pytest.mark.asyncio
|
|
async def test_auto_title_disabled_returns_none(self):
|
|
"""When auto_title is False, returns None (no title generation)."""
|
|
doc = DoclingDocument(name="test")
|
|
doc.add_text(label=DocItemLabel.TITLE, text="Structural Title")
|
|
|
|
config = AppConfig(processing=ProcessingConfig(auto_title=False))
|
|
result = await resolve_title(config, doc, "some content")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_structural_title_extracted(self):
|
|
"""Structural title is extracted when auto_title is enabled."""
|
|
doc = DoclingDocument(name="test")
|
|
doc.add_text(label=DocItemLabel.TITLE, text="Auto Extracted Title")
|
|
|
|
config = AppConfig(processing=ProcessingConfig(auto_title=True))
|
|
result = await resolve_title(config, doc, "some content")
|
|
assert result == "Auto Extracted Title"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_failure_returns_none(self, monkeypatch):
|
|
"""LLM failure during ingestion returns None instead of raising."""
|
|
doc = DoclingDocument(name="test")
|
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just text")
|
|
|
|
async def exploding_llm(config, content):
|
|
raise RuntimeError("LLM is down")
|
|
|
|
monkeypatch.setattr(
|
|
"haiku.rag.client.titles.generate_title_with_llm", exploding_llm
|
|
)
|
|
config = AppConfig(processing=ProcessingConfig(auto_title=True))
|
|
result = await resolve_title(config, 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
|