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

Stop loading page rasters on the title and update paths

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-08-06 12:28:18 +03:00
parent ac9b2cbf81
commit eb11a165b8
No known key found for this signature in database
6 changed files with 79 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
CHANGELOG.md
View file

@ -3,7 +3,7 @@
### Changed
- `get_document_by_id` / `get_document_by_uri` return content and the mutable attributes only; the docling structure and page-image blobs are no longer loaded. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or pass `include_blobs=True` to the repository method.
- `get_document_by_id` / `get_document_by_uri` no longer load the docling structure and page-image blobs. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or `get_by_id(..., include_blobs=True)`.
### Fixed

35
evaluations/tests/test_benchmark.py
View file

@ -15,6 +15,20 @@ from evaluations.config import DatasetSpec
from haiku.rag.config.models import AppConfig, ModelConfig
def _stub_spec(**overrides) -> DatasetSpec:
"""A DatasetSpec whose loaders/mappers are inert, for tests that only
exercise the surrounding plumbing."""
return DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
**overrides,
)
class TestBuildExperimentMetadata:
def test_basic_metadata(self) -> None:
config = AppConfig()
@ -460,16 +474,8 @@ class TestRetrievalTarget:
from evaluations.config import RetrievalSample
from evaluations.evaluators import MAPEvaluator
return DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
retrieval_loader=lambda: [ # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
{"q": "What is X?", "uris": ("uri-x",)},
],
return _stub_spec(
retrieval_loader=lambda: [{"q": "What is X?", "uris": ("uri-x",)}],
retrieval_mapper=lambda d: RetrievalSample(
question=d["q"], expected_uris=d["uris"]
),
@ -545,14 +551,7 @@ class TestRetrievalTarget:
class TestEvaluateDatasetCaseIds:
def _spec(self) -> DatasetSpec:
return DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
return _stub_spec()
@pytest.mark.asyncio
async def test_threads_case_ids_to_qa_benchmark(self) -> None:

6
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -856,10 +856,10 @@ async def update_document(
"Provide one or the other, not both."
)
# An update that only replaces content writes the record back as-is, so
# the blobs have to make the round trip.
# Caller-supplied chunks without a docling document replace neither blob,
# and the row is written back whole, so they have to make the round trip.
existing_doc = await client.document_repository.get_by_id(
document_id, include_blobs=True
document_id, include_blobs=chunks is not None and docling_document is None
)
if existing_doc is None:
raise ValueError(f"Document with ID {document_id} not found")

20
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -176,7 +176,7 @@ async def _set_embedder(client: "HaikuRAG") -> None:
async def _hydrate(
client: "HaikuRAG", light_docs: list[Document]
client: "HaikuRAG", light_docs: list[Document], include_blobs: bool = True
) -> AsyncGenerator[Document, None]:
"""Yield fully-loaded documents one at a time from a light listing.
@ -188,7 +188,7 @@ async def _hydrate(
for light_doc in light_docs:
assert light_doc.id is not None
doc = await client.document_repository.get_by_id(
light_doc.id, include_blobs=True
light_doc.id, include_blobs=include_blobs
)
if doc is None:
continue
@ -199,9 +199,18 @@ async def _hydrate(
async def _rebuild_title_only(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Generate titles for documents that don't have one."""
"""Generate titles for documents that don't have one.
A title comes from the content or the docling structure, never the page
rasters, so those are left out of the per-document load.
"""
repo = client.document_repository
untitled = [d for d in documents if d.title is None]
async for doc in _hydrate(client, untitled):
async for doc in _hydrate(client, untitled, include_blobs=False):
assert doc.id is not None
structure = await repo.get_docling_data(doc.id)
if structure is not None:
doc.docling_document = structure.docling_document
try:
title = await client.generate_title(doc)
except Exception:
@ -211,8 +220,7 @@ async def _rebuild_title_only(
continue
if title is not None:
doc.title = title
await client.document_repository.update_meta(doc)
assert doc.id is not None
await repo.update_meta(doc)
yield doc.id

33
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -184,22 +184,26 @@ class DocumentRepository:
_LIGHT_COLUMNS = ["id", "content"]
async def _record_by_id(
self, doc_id: str, include_blobs: bool
) -> DocumentRecord | None:
safe_id = escape_sql_string(doc_id)
query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
if not include_blobs:
query = query.select(self._LIGHT_COLUMNS)
results = await query_to_pydantic(query, DocumentRecord)
return results[0] if results else None
async def get_by_id(
self, entity_id: str, include_blobs: bool = False
) -> Document | None:
"""Get a document by its ID. `include_blobs` adds the docling blobs."""
safe_id = escape_sql_string(entity_id)
query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
results = await query_to_pydantic(
query if include_blobs else query.select(self._LIGHT_COLUMNS),
DocumentRecord,
)
if not results:
record = await self._record_by_id(entity_id, include_blobs)
if record is None:
return None
meta = await self._meta_by_id(entity_id)
return self._merge_to_document(results[0], meta)
return self._merge_to_document(record, meta)
async def get_content(self, entity_id: str) -> str | None:
"""Get only the text content of a document (skips docling blobs)."""
@ -394,16 +398,11 @@ class DocumentRepository:
return None
meta = meta_results[0]
safe_id = escape_sql_string(meta.id)
query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
doc_results = await query_to_pydantic(
query if include_blobs else query.select(self._LIGHT_COLUMNS),
DocumentRecord,
)
if not doc_results:
record = await self._record_by_id(meta.id, include_blobs)
if record is None:
return None
return self._merge_to_document(doc_results[0], meta)
return self._merge_to_document(record, meta)
async def delete_all(self) -> None:
"""Delete all documents from the database."""

28
tests/test_rebuild.py
View file

@ -551,6 +551,34 @@ async def test_rebuild_full_with_accessible_source(temp_db_path):
assert "Fresh content" in new_doc.content
async def test_rebuild_title_only_reads_structural_title(temp_db_path):
"""TITLE_ONLY takes the title from the stored docling structure, so it never
reaches the LLM for a document that carries one."""
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.store.models.document import Document
docling_doc = DoclingDocument(name="structured")
docling_doc.add_text(label=DocItemLabel.TITLE, text="The Stored Title")
async with HaikuRAG(temp_db_path, create=True) as client:
doc = Document(content="body text", metadata={})
doc.set_docling(docling_doc)
created = await client.document_repository.create(doc)
assert created.id is not None
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY)
]
assert processed_ids == [created.id]
refreshed = await client.get_document_by_id(created.id)
assert refreshed is not None
assert refreshed.title == "The Stored Title"
async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch):
"""TITLE_ONLY: a failure on one document does not abort the generator.