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

Merge pull request #558 from ggozad/fix/ingest-rebuild-correctness

Browse source 
correctness defects in source ingestion and rebuild
This commit is contained in:
Yiorgis Gozadinos 2026-08-19 14:02:01 +03:00 committed by GitHub
commit 489b8a65f0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 537 additions and 146 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -5,6 +5,7 @@
### Added
- `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout.
- `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata.
- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes.
- Raw chunk metadata is now exposed to search and citation results, through `SearchResult.chunk_meta` and `Citation.chunk_meta`. For context-expanded results, the metadata is that of the anchor chunk.
@ -27,6 +28,9 @@
### Fixed
- `create_document_from_source` closes the source adapter it builds for the call; adapters passed in through `sources` are left to their owner.
- Directory ingestion skips symlinked files resolving outside the given directory, matching `FSSource.discover`.
- A FULL rebuild no longer deletes a source-backed document before re-ingesting it: it refreshes the document in place, so the document id is preserved and a failed fetch or conversion falls back to rebuilding from stored content instead of losing the document.
- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`.
- `server.json` runtime arguments are `mcp --stdio`, was `serve --mcp`.

1
docs/remote-processing.md
View file

@ -64,6 +64,7 @@ providers:
docling_serve:
base_url: http://localhost:5001
api_key: "" # Optional API key for authentication
timeout: 300 # Per-request timeout in seconds
```
For converter / chunker config options (chunking strategy, tokenizer,

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

@ -677,6 +677,7 @@ async def create_document_from_source(
sources: "list[Source] | None" = None,
source_id: str | None = None,
metadata_provider: "MetadataProvider | None" = None,
force: bool = False,
) -> Document | list[Document]:
"""Create or update document(s) from a file path, directory, or URL.
@ -685,6 +686,10 @@ async def create_document_from_source(
- If MD5 changed, updates the document
- If no document exists, creates a new one
``force`` skips both freshness checks so an unchanged source is re-converted,
re-chunked and re-embedded into the existing document. Internal: rebuild uses
it to refresh a document in place instead of deleting and recreating it.
If ``uri`` is provided, it overrides the URI auto-derived from the source
(which is normally ``file://`` for local files or the URL for remote
sources). Not supported for directory sources, which produce one document
@ -721,6 +726,14 @@ async def create_document_from_source(
documents: list[Document] = []
filter = FileFilter()
for child in local_path.rglob("*"):
# rglob does not recurse into symlinked directories, but it does
# yield symlinked files. Skip the ones resolving outside the
# directory the caller named, as FSSource.discover does.
if child.is_symlink():
resolved = child.resolve(strict=False)
if not resolved.is_relative_to(local_path.resolve()):
continue
child = resolved
if child.is_file() and filter.include_file(str(child)):
doc = await create_document_from_source(
client,
@ -730,6 +743,7 @@ async def create_document_from_source(
sources=sources,
source_id=source_id,
metadata_provider=metadata_provider,
force=force,
)
assert isinstance(doc, Document)
documents.append(doc)
@ -762,76 +776,91 @@ async def create_document_from_source(
source_str, sources=sources, storage_options=storage_options
)
# The stored URI is what we look up + persist by. For an explicit uri
# override, use it as-is. For a file:// input the source string is
# already canonical (URL-encoded); round-tripping via Path.as_uri()
# would double-encode any escapes like %5B. For bare paths, canonicalize.
if uri is not None:
stored_uri = uri
elif parsed_url.scheme == "file":
stored_uri = source_str
elif parsed_url.scheme == "":
stored_uri = Path(source_str).absolute().as_uri()
else:
stored_uri = source_str
# A fetcher built for this call holds its own httpx pool (HTTP, WebDAV) and
# has to be closed here. One handed in through `sources` belongs to the
# caller: the ingester keeps its sources open across jobs.
owns_fetcher = all(fetcher is not configured for configured in sources or ())
existing_doc = await client.get_document_by_uri(stored_uri)
try:
# The stored URI is what we look up + persist by. For an explicit uri
# override, use it as-is. For a file:// input the source string is
# already canonical (URL-encoded); round-tripping via Path.as_uri()
# would double-encode any escapes like %5B. For bare paths, canonicalize.
if uri is not None:
stored_uri = uri
elif parsed_url.scheme == "file":
stored_uri = source_str
elif parsed_url.scheme == "":
stored_uri = Path(source_str).absolute().as_uri()
else:
stored_uri = source_str
# Cheap revision-based short-circuit: only worth a HEAD when we have a
# stored revision to compare against. All sources persist their native
# revision (mtime_ns for FS, ETag for S3, ETag/Last-Modified for HTTP)
# under the canonical "source_revision" metadata key.
stored_revision = (
(existing_doc.metadata or {}).get("source_revision") if existing_doc else None
)
if existing_doc and stored_revision:
current_revision = await fetcher.head(source_str)
if current_revision == stored_revision:
existing_doc = await client.get_document_by_uri(stored_uri)
# Cheap revision-based short-circuit: only worth a HEAD when we have a
# stored revision to compare against. All sources persist their native
# revision (mtime_ns for FS, ETag for S3, ETag/Last-Modified for HTTP)
# under the canonical "source_revision" metadata key.
stored_revision = (
(existing_doc.metadata or {}).get("source_revision")
if existing_doc
else None
)
if existing_doc and stored_revision and not force:
current_revision = await fetcher.head(source_str)
if current_revision == stored_revision:
return await _refresh_doc_metadata(
client,
existing_doc,
title=title,
user_metadata=metadata,
source_metadata=None,
)
with logfire.span("document.fetch", uri=source_str) as fetch_span:
result = await fetcher.fetch(source_str)
fetch_span.set_attribute("bytes", len(result.body))
fetch_span.set_attribute("content_hash", result.content_hash)
provider_metadata = await _provider_metadata(
metadata_provider, source_id or fetcher.source_id, source_str, result
)
user_metadata = {**metadata, **provider_metadata}
# MD5 short-circuit: the bytes are unchanged even if the revision wasn't.
# Refresh the source-derived metadata (revision may have rolled) but skip
# convert/embed/store entirely.
if (
existing_doc
and not force
and existing_doc.metadata.get("md5") == result.content_hash
):
source_meta: dict = {
"content_type": result.content_type,
"md5": result.content_hash,
**result.extra_metadata,
}
if result.revision is not None:
source_meta["source_revision"] = result.revision
return await _refresh_doc_metadata(
client,
existing_doc,
title=title,
user_metadata=metadata,
source_metadata=None,
user_metadata=user_metadata,
source_metadata=source_meta,
)
with logfire.span("document.fetch", uri=source_str) as fetch_span:
result = await fetcher.fetch(source_str)
fetch_span.set_attribute("bytes", len(result.body))
fetch_span.set_attribute("content_hash", result.content_hash)
provider_metadata = await _provider_metadata(
metadata_provider, source_id or fetcher.source_id, source_str, result
)
user_metadata = {**metadata, **provider_metadata}
# MD5 short-circuit: the bytes are unchanged even if the revision wasn't.
# Refresh the source-derived metadata (revision may have rolled) but skip
# convert/embed/store entirely.
if existing_doc and existing_doc.metadata.get("md5") == result.content_hash:
source_meta: dict = {
"content_type": result.content_type,
"md5": result.content_hash,
**result.extra_metadata,
}
if result.revision is not None:
source_meta["source_revision"] = result.revision
return await _refresh_doc_metadata(
return await _ingest_fetch_result(
client,
existing_doc,
result,
title=title,
user_metadata=user_metadata,
source_metadata=source_meta,
stored_uri=stored_uri,
existing_doc=existing_doc,
)
return await _ingest_fetch_result(
client,
result,
title=title,
user_metadata=user_metadata,
stored_uri=stored_uri,
existing_doc=existing_doc,
)
finally:
if owns_fetcher:
await fetcher.aclose()
async def update_document(

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

@ -8,7 +8,10 @@ from typing import TYPE_CHECKING
from docling_core.types.doc.document import DescriptionMetaField, PictureMeta
from lancedb.pydantic import LanceModel
from haiku.rag.client.documents import check_source_accessible
from haiku.rag.client.documents import (
check_source_accessible,
create_document_from_source,
)
from haiku.rag.converters import get_converter
from haiku.rag.store.compression import compress_docling_split
from haiku.rag.store.engine import ChunkRecordBase
@ -807,28 +810,38 @@ async def _rebuild_full(
# Try to rebuild from source if available — uses the light listing
# directly, no need to load the stored content/blobs first.
if light_doc.uri and check_source_accessible(light_doc.uri):
try:
# Flush pending batch before source rebuild (creates new doc)
if pending_docs:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
pending_chunks = []
pending_docs = []
# The refresh writes through the client, not the batch buffer, so
# anything pending has to land first.
if pending_docs:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
pending_chunks = []
pending_docs = []
await client.delete_document(light_doc.id)
new_doc = await client.create_document_from_source(
source=light_doc.uri, metadata=light_doc.metadata or {}
try:
# force=True: the source bytes are usually unchanged, and the
# point of a FULL rebuild is to re-convert them anyway. Updates
# in place, so a failure here cannot cost the document.
refreshed = await create_document_from_source(
client,
source=light_doc.uri,
metadata=light_doc.metadata or {},
force=True,
)
assert isinstance(new_doc, Document)
assert new_doc.id is not None
yield new_doc.id
assert isinstance(refreshed, Document)
assert refreshed.id is not None
yield refreshed.id
continue
except Exception as e:
logger.error(
"Error recreating document from source %s: %s",
logger.warning(
"Rebuilding %s from source failed (%s), "
"falling back to stored content",
light_doc.uri,
e,
)
continue
elif light_doc.uri:
logger.warning(
"Source missing for %s, re-embedding from content", light_doc.uri
)
# Fallback: rebuild from stored content. Now we need the full
# record (content + docling_pages for the round-trip write).
@ -838,8 +851,6 @@ async def _rebuild_full(
if doc is None:
continue
assert doc.id is not None
if doc.uri:
logger.warning("Source missing for %s, re-embedding from content", doc.uri)
docling_document = await converter.convert_text(doc.content, format="md")
chunks = await client.chunk(docling_document)

5
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -295,6 +295,11 @@ class DoclingServeConfig(BaseModel):
description="Max attempts per request across the fleet before giving up; "
"each retry fails over to another instance.",
)
timeout: float = Field(
default=300,
gt=0,
description="Per-request timeout in seconds for submit, poll and result calls.",
)
circuit_breaker: CircuitBreakerConfig = Field(
default_factory=lambda: CircuitBreakerConfig(
failure_threshold=3, cooldown_s=30.0

1
haiku_rag_slim/haiku/rag/providers/docling_serve.py
View file

@ -97,6 +97,7 @@ class DoclingServeClient:
return cls(
base_urls=config.base_urls,
api_key=config.api_key,
timeout=config.timeout,
circuit_breaker=config.circuit_breaker,
max_attempts=config.max_attempts,
)

3
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -60,7 +60,8 @@ class ChunkRepository:
"""Create one or more chunks in the database.
Chunks must have embeddings set before calling this method.
Use client._ensure_chunks_embedded() to embed chunks if needed.
Use haiku.rag.client.processing.ensure_chunks_embedded() to embed
chunks if needed.
"""
self.store._assert_writable()
# Handle single chunk

44
tests/cassettes/test_client/test_adhoc_source_is_closed_after_ingest.yaml Normal file
View file

File diff suppressed because one or more lines are too long

44
tests/cassettes/test_client/test_configured_source_is_not_closed_after_ingest.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_client/test_directory_ingest_skips_symlinks_escaping_the_tree.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_rebuild/test_rebuild_full_source_failure_falls_back_to_stored_content.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_rebuild_full_source_failure_is_logged_and_skipped.yaml
View file

File diff suppressed because one or more lines are too long

95
tests/test_client.py
View file

@ -19,6 +19,7 @@ from haiku.rag.client.documents import (
)
from haiku.rag.config import Config
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.ingester.sources.base import FetchResult
from haiku.rag.store.compression import decompress_json
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
@ -444,6 +445,32 @@ async def test_client_create_document_from_directory(temp_db_path):
assert not any("unsupported.xyz" in uri for uri in uris)
@pytest.mark.vcr()
async def test_directory_ingest_skips_symlinks_escaping_the_tree(temp_db_path):
"""A symlinked file resolving outside the named directory is not ingested;
one resolving inside it is. Matches FSSource.discover."""
import os
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir) / "tree"
outside = Path(temp_dir) / "outside"
root.mkdir()
outside.mkdir()
(root / "real.txt").write_text("Inside the tree.")
(outside / "secret.txt").write_text("Outside the tree.")
os.symlink(outside / "secret.txt", root / "escape.txt")
os.symlink(root / "real.txt", root / "inside_link.txt")
result = await client.create_document_from_source(root)
assert isinstance(result, list)
uris = sorted(uri for doc in result if (uri := doc.uri))
assert not any("secret" in uri or "escape" in uri for uri in uris)
assert any("real.txt" in uri for uri in uris)
@pytest.mark.vcr()
async def test_client_create_document_from_url(temp_db_path):
"""Test creating a document from a URL."""
@ -2408,6 +2435,74 @@ def test_check_source_accessible_file_uri(tmp_path):
assert check_source_accessible((tmp_path / "gone.txt").as_uri()) is False
class _CountingSource:
"""A real Source over one in-memory document that counts its closes."""
def __init__(self, uri: str, body: bytes) -> None:
self.source_id = "counting"
self.supported_extensions = None
self.max_file_size = None
self._uri = uri
self._body = body
self.closes = 0
def supports(self, uri: str) -> bool:
return uri == self._uri
async def head(self, uri: str) -> str | None:
return "v1"
async def aclose(self) -> None:
self.closes += 1
async def fetch(self, uri: str) -> "FetchResult":
import hashlib
return FetchResult(
uri=uri,
body=self._body,
content_type="text/markdown",
content_hash=hashlib.md5(self._body).hexdigest(),
revision="v1",
)
def discover(self, since=None, *, known_uris=None):
raise NotImplementedError
@pytest.mark.vcr()
async def test_adhoc_source_is_closed_after_ingest(temp_db_path, monkeypatch):
"""An ad-hoc fetcher is built for this one call, so this call has to close
it — HTTP and WebDAV adapters hold an httpx connection pool."""
from haiku.rag.ingester import sources as sources_module
uri = "https://example.com/counting.md"
fetcher = _CountingSource(uri, b"# Counting\n\nAd-hoc fetched body.")
monkeypatch.setattr(
sources_module, "resolve_adhoc_fetcher", lambda *a, **kw: fetcher
)
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document_from_source(uri)
assert not isinstance(doc, list)
assert fetcher.closes == 1
@pytest.mark.vcr()
async def test_configured_source_is_not_closed_after_ingest(temp_db_path):
"""A source handed in by the caller (the ingester's long-lived pool) is not
ours to close: closing it would tear down the pool mid-run."""
uri = "https://example.com/configured.md"
fetcher = _CountingSource(uri, b"# Configured\n\nCaller-owned body.")
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document_from_source(uri, sources=[fetcher])
assert not isinstance(doc, list)
assert fetcher.closes == 0
def _bbox_doc(*, with_page_image: bool, pages: tuple[int, ...] = (1,)):
"""DoclingDocument with one paragraph per page, each carrying a bbox.

20
tests/test_config.py
View file

@ -1,3 +1,5 @@
from pathlib import Path
import pytest
import yaml
@ -557,3 +559,21 @@ def test_get_config_initialises_lazily_then_reuses(monkeypatch):
first = config_module.get_config()
assert isinstance(first, AppConfig)
assert config_module.get_config() is first
# Every YAML config shipped in the repo has to load. These are what users copy.
_EXAMPLE_CONFIGS = sorted(
(Path(__file__).resolve().parent.parent).glob("**/*.yaml.example")
)
def test_example_configs_are_present():
"""Guard against the glob silently matching nothing."""
assert _EXAMPLE_CONFIGS
@pytest.mark.parametrize(
"path", _EXAMPLE_CONFIGS, ids=lambda p: p.parent.name + "/" + p.name
)
def test_example_config_validates(path: Path):
AppConfig.model_validate(yaml.safe_load(path.read_text()) or {})

2
tests/test_docling_serve_client.py
View file

@ -474,11 +474,13 @@ def test_from_config_wires_retry_and_breaker():
ds = config.providers.docling_serve
ds.base_url = "http://cfg-n:5001"
ds.max_attempts = 7
ds.timeout = 42.0
ds.circuit_breaker = CircuitBreakerConfig(failure_threshold=9, cooldown_s=90.0)
for component in (DoclingServeConverter(config), DoclingServeChunker(config)):
client = component.client
assert client._max_attempts == 7
assert client.timeout == 42.0
assert client._breaker_config.failure_threshold == 9
assert client._breaker_config.cooldown_s == 90.0

66
tests/test_rebuild.py
View file

@ -523,7 +523,10 @@ async def test_rebuild_full_with_accessible_source(temp_db_path):
"""FULL rebuild re-ingests from source when the URI is accessible.
Covers the main path in _rebuild_full (source-accessible branch): the
document is deleted and re-created from its URI, producing a new ID.
document is refreshed in place, keeping its ID. The source bytes are
unchanged since ingestion, so this also pins that the refresh bypasses the
revision and MD5 short-circuits instead of returning the document
untouched.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
@ -540,15 +543,17 @@ async def test_rebuild_full_with_accessible_source(temp_db_path):
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
# Original doc was deleted and a new one created; the old ID
# must not appear, and exactly one new ID must have been yielded.
assert original_id not in processed_ids
assert len(processed_ids) == 1
assert processed_ids == [original_id]
new_doc = await client.get_document_by_id(processed_ids[0])
assert new_doc is not None
assert new_doc.uri == source_path.as_uri()
assert "Fresh content" in new_doc.content
refreshed = await client.get_document_by_id(original_id)
assert refreshed is not None
assert refreshed.uri == source_path.as_uri()
assert "Fresh content" in refreshed.content
# The chunks table is recreated at the top of FULL, so the refresh
# must have written new chunks for the document to stay searchable.
chunks = await client.chunk_repository.get_by_document_id(original_id)
assert chunks
async def test_rebuild_title_only_reads_structural_title(temp_db_path):
@ -622,14 +627,14 @@ async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch)
@pytest.mark.vcr()
async def test_rebuild_full_source_failure_is_logged_and_skipped(
async def test_rebuild_full_source_failure_falls_back_to_stored_content(
temp_db_path, monkeypatch
):
"""FULL rebuild logs-and-continues when re-ingesting from source fails.
"""A failed source refresh must never cost the document.
Covers _rebuild_full's `except Exception` branch: when
create_document_from_source raises, the doc is skipped (no yield) and
the error is logged. Regression guard against silent failures.
Covers _rebuild_full's `except Exception` branch: when the refresh raises,
the document keeps its stored row and is rebuilt from stored content, so it
is still readable and still searchable afterwards.
"""
import logging
@ -644,22 +649,31 @@ async def test_rebuild_full_source_failure_is_logged_and_skipped(
assert not isinstance(original, list)
assert original.id is not None
# Force the source rebuild branch to raise.
async def failing_create(*args, **kwargs):
# Force the source refresh to raise.
async def failing_refresh(*args, **kwargs):
raise RuntimeError("simulated ingestion failure")
monkeypatch.setattr(client, "create_document_from_source", failing_create)
monkeypatch.setattr(
rebuild_module, "create_document_from_source", failing_refresh
)
with capture_logs(rebuild_module.logger, logging.ERROR) as records:
with capture_logs(rebuild_module.logger, logging.WARNING) as records:
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
assert processed_ids == []
assert processed_ids == [original.id]
survivor = await client.get_document_by_id(original.id)
assert survivor is not None
assert "Content that will vanish" in survivor.content
chunks = await client.chunk_repository.get_by_document_id(original.id)
assert chunks
assert any(
"Error recreating document from source" in rec.getMessage()
for rec in records
"falling back to stored content" in rec.getMessage() for rec in records
)
@ -1412,12 +1426,10 @@ async def test_rebuild_full_flushes_pending_before_source_rebuild(temp_db_path):
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
# The content-path document keeps its id and must survive the
# flush that precedes the source re-ingest; the source document
# is replaced by a freshly ingested one with a new id.
assert content_doc.id in processed
assert source_doc.id not in processed
assert len(processed) == 2
# Both documents keep their ids: the content-path one must survive
# the flush that precedes the source refresh, and the source one is
# refreshed in place.
assert sorted(processed) == sorted([content_doc.id, source_doc.id])
assert await client.store.documents_table.count_rows() == 2