Refresh source-backed documents in place on a FULL rebuild
FULL rebuild deleted a document before re-ingesting it from its URI, and the handler around that logged and continued. A 404, a timeout or any conversion error therefore removed the document permanently. Deleting after a successful create is not an alternative: create_document_from_source resolves the same URI to the existing document and updates it in place, so a trailing delete would remove the freshly rebuilt row. Refresh in place instead. create_document_from_source takes an internal force flag that skips the revision and MD5 short-circuits, so an unchanged source is still re-converted, re-chunked and re-embedded into the existing document, and the document id survives a rebuild. A failed refresh now falls through to the stored-content path rather than skipping the document: FULL recreates the chunks table before the loop, so skipping left the document present but unsearchable until the next rebuild. The pending-batch flush moves out of the try. A failed flush is a lost write and should abort the rebuild, not be logged and skipped.
This commit is contained in:
Yiorgis Gozadinos
parent
13bd69b908
commit
895b5f0532
6 changed files with 163 additions and 89 deletions
|
|
@ -25,6 +25,7 @@
|
|||
|
||||
### Fixed
|
||||
|
||||
- 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`.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -730,6 +735,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)
|
||||
|
|
@ -784,7 +790,7 @@ async def create_document_from_source(
|
|||
stored_revision = (
|
||||
(existing_doc.metadata or {}).get("source_revision") if existing_doc else None
|
||||
)
|
||||
if existing_doc and stored_revision:
|
||||
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(
|
||||
|
|
@ -808,7 +814,11 @@ async def create_document_from_source(
|
|||
# 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:
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue