Type writing against one database
Every implementation in documents.py and rebuild.py takes the session it writes to, so a set cannot reach one: the facade narrows once and passes the database on, rather than checking and carrying a union. Tests calling an implementation directly go through `writing()`.
This commit is contained in:
Yiorgis Gozadinos
parent
028217b7c0
commit
16b7319c48
12 changed files with 334 additions and 249 deletions
|
|
@ -139,7 +139,6 @@ class HaikuRAG:
|
|||
self._read_only = read_only
|
||||
self._requested_sources = sources
|
||||
self._clients: dict[str, HaikuRAG] = {}
|
||||
self._clients_lock = asyncio.Lock()
|
||||
self._source: str | None = None
|
||||
self._session: SingleDatabaseSession | FederatedSession | None = None
|
||||
self._owns_session = True
|
||||
|
|
@ -418,14 +417,6 @@ class HaikuRAG:
|
|||
if cached is not None:
|
||||
await aclose_quietly(cached, name)
|
||||
|
||||
async def _await_vacuum_tasks(self) -> None:
|
||||
if isinstance(self._session, SingleDatabaseSession):
|
||||
await self._session.drain_vacuum()
|
||||
|
||||
def _schedule_vacuum(self) -> None:
|
||||
if isinstance(self._session, SingleDatabaseSession):
|
||||
self._session.schedule_vacuum()
|
||||
|
||||
# =========================================================================
|
||||
# Processing Primitives
|
||||
# =========================================================================
|
||||
|
|
@ -487,9 +478,9 @@ class HaikuRAG:
|
|||
) -> Document:
|
||||
from haiku.rag.client.documents import create_document
|
||||
|
||||
self._single_session("create_document")
|
||||
session = self._single_session("create_document")
|
||||
|
||||
return await create_document(self, content, uri, title, metadata, format)
|
||||
return await create_document(session, content, uri, title, metadata, format)
|
||||
|
||||
async def import_document(
|
||||
self,
|
||||
|
|
@ -501,10 +492,10 @@ class HaikuRAG:
|
|||
) -> Document:
|
||||
from haiku.rag.client.documents import import_document
|
||||
|
||||
self._single_session("import_document")
|
||||
session = self._single_session("import_document")
|
||||
|
||||
return await import_document(
|
||||
self, docling_document, chunks, uri, title, metadata
|
||||
session, docling_document, chunks, uri, title, metadata
|
||||
)
|
||||
|
||||
async def import_documents(
|
||||
|
|
@ -513,9 +504,9 @@ class HaikuRAG:
|
|||
) -> list[Document]:
|
||||
from haiku.rag.client.documents import import_documents
|
||||
|
||||
self._single_session("import_documents")
|
||||
session = self._single_session("import_documents")
|
||||
|
||||
return await import_documents(self, imports)
|
||||
return await import_documents(session, imports)
|
||||
|
||||
async def create_document_from_source(
|
||||
self,
|
||||
|
|
@ -530,10 +521,10 @@ class HaikuRAG:
|
|||
) -> Document | list[Document]:
|
||||
from haiku.rag.client.documents import create_document_from_source
|
||||
|
||||
self._single_session("create_document_from_source")
|
||||
session = self._single_session("create_document_from_source")
|
||||
|
||||
return await create_document_from_source(
|
||||
self,
|
||||
session,
|
||||
source,
|
||||
title,
|
||||
metadata,
|
||||
|
|
@ -556,10 +547,10 @@ class HaikuRAG:
|
|||
) -> Document:
|
||||
from haiku.rag.client.documents import update_document
|
||||
|
||||
self._single_session("update_document")
|
||||
session = self._single_session("update_document")
|
||||
|
||||
return await update_document(
|
||||
self,
|
||||
session,
|
||||
document_id,
|
||||
content,
|
||||
metadata,
|
||||
|
|
@ -841,9 +832,9 @@ class HaikuRAG:
|
|||
) -> AsyncGenerator[str, None]:
|
||||
from haiku.rag.client.rebuild import rebuild_database
|
||||
|
||||
self._single_session("rebuild_database")
|
||||
session = self._single_session("rebuild_database")
|
||||
|
||||
async for doc_id in rebuild_database(self, mode):
|
||||
async for doc_id in rebuild_database(session, mode):
|
||||
yield doc_id
|
||||
|
||||
async def vacuum(self) -> None:
|
||||
|
|
|
|||
|
|
@ -11,9 +11,12 @@ from urllib.parse import quote, urlparse
|
|||
from haiku.rag.client.exceptions import UnsupportedSourceError
|
||||
from haiku.rag.client.processing import (
|
||||
_write_fetch_body,
|
||||
chunk,
|
||||
convert,
|
||||
ensure_chunks_embedded,
|
||||
get_extension_from_content_type_or_url,
|
||||
)
|
||||
from haiku.rag.client.session import SingleDatabaseSession
|
||||
from haiku.rag.client.titles import resolve_title
|
||||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
|
|
@ -25,7 +28,6 @@ from haiku.rag.uri import is_local_uri, uri_to_path
|
|||
if TYPE_CHECKING:
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.ingester.metadata import MetadataProvider
|
||||
from haiku.rag.sources.base import FetchResult, Source
|
||||
|
||||
|
|
@ -86,7 +88,9 @@ async def _prepare_document_from_docling(
|
|||
|
||||
|
||||
async def _prepare_and_title(
|
||||
client: "HaikuRAG", document: Document, docling_document: "DoclingDocument"
|
||||
session: SingleDatabaseSession,
|
||||
document: Document,
|
||||
docling_document: "DoclingDocument",
|
||||
) -> None:
|
||||
"""Fill the document from its converted form and title it if it has none.
|
||||
|
||||
|
|
@ -97,10 +101,27 @@ async def _prepare_and_title(
|
|||
stored_content = await _prepare_document_from_docling(document, docling_document)
|
||||
if document.title is None:
|
||||
document.title = await resolve_title(
|
||||
client._config, docling_document, stored_content
|
||||
session.config, docling_document, stored_content
|
||||
)
|
||||
|
||||
|
||||
async def chunk_document(
|
||||
session: SingleDatabaseSession,
|
||||
docling_document: "DoclingDocument",
|
||||
*,
|
||||
existing_picture_data: dict[str, bytes] | None = None,
|
||||
document_id: str | None = None,
|
||||
) -> list[Chunk]:
|
||||
"""Chunk and embed through the database this write belongs to."""
|
||||
return await chunk(
|
||||
session.config,
|
||||
docling_document,
|
||||
embedder=session.store.embedder,
|
||||
existing_picture_data=existing_picture_data,
|
||||
document_id=document_id,
|
||||
)
|
||||
|
||||
|
||||
def parent_uri_filter(parent_uri: str) -> str:
|
||||
"""SQL `WHERE` clause matching documents whose ``metadata.parent_uri``
|
||||
equals ``parent_uri``. ``metadata`` is stored as a JSON string produced by
|
||||
|
|
@ -112,7 +133,7 @@ def parent_uri_filter(parent_uri: str) -> str:
|
|||
|
||||
|
||||
async def _store_document_with_chunks(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
document: Document,
|
||||
chunks: list[Chunk],
|
||||
docling_document: "DoclingDocument",
|
||||
|
|
@ -121,16 +142,18 @@ async def _store_document_with_chunks(
|
|||
|
||||
Handles versioning/rollback on failure.
|
||||
"""
|
||||
chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
|
||||
chunks = await ensure_chunks_embedded(
|
||||
session.config, chunks, session.store.embedder
|
||||
)
|
||||
items = await asyncio.to_thread(extract_items, "", docling_document)
|
||||
|
||||
async with client.store.write_transaction():
|
||||
async with session.store.write_transaction():
|
||||
# A concurrent ingestion of the same URI may have created the document
|
||||
# while this one was converting/embedding outside the lock. LanceDB has
|
||||
# no unique constraint on `uri`, so re-check under the lock and update in
|
||||
# place rather than inserting a duplicate.
|
||||
existing = (
|
||||
await client.get_document_by_uri(document.uri)
|
||||
await session.get_document_by_uri(document.uri)
|
||||
if document.uri is not None
|
||||
else None
|
||||
)
|
||||
|
|
@ -138,9 +161,9 @@ async def _store_document_with_chunks(
|
|||
if existing is not None:
|
||||
document.id = existing.id
|
||||
document.created_at = existing.created_at
|
||||
stored_doc = await client.document_repository.update(document)
|
||||
stored_doc = await session.document_repository.update(document)
|
||||
else:
|
||||
stored_doc = await client.document_repository.create(document)
|
||||
stored_doc = await session.document_repository.create(document)
|
||||
|
||||
assert stored_doc.id is not None, "Document ID should not be None after storing"
|
||||
for order, chunk in enumerate(chunks):
|
||||
|
|
@ -150,22 +173,22 @@ async def _store_document_with_chunks(
|
|||
item.document_id = stored_doc.id
|
||||
|
||||
if existing is not None:
|
||||
await client.chunk_repository.replace_for_document(stored_doc.id, chunks)
|
||||
await client.document_item_repository.replace_for_document(
|
||||
await session.chunk_repository.replace_for_document(stored_doc.id, chunks)
|
||||
await session.document_item_repository.replace_for_document(
|
||||
stored_doc.id, items
|
||||
)
|
||||
else:
|
||||
await client.chunk_repository.create(chunks)
|
||||
await client.document_item_repository.create_items(stored_doc.id, items)
|
||||
await session.chunk_repository.create(chunks)
|
||||
await session.document_item_repository.create_items(stored_doc.id, items)
|
||||
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
if session.config.storage.auto_vacuum:
|
||||
session.schedule_vacuum()
|
||||
|
||||
return stored_doc
|
||||
|
||||
|
||||
async def _update_document_with_chunks(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
document: Document,
|
||||
chunks: list[Chunk],
|
||||
docling_document: "DoclingDocument | None" = None,
|
||||
|
|
@ -183,10 +206,12 @@ async def _update_document_with_chunks(
|
|||
existing_picture_data: dict[str, bytes] | None = None
|
||||
if docling_document is not None:
|
||||
existing_picture_data = (
|
||||
await client.document_item_repository.get_all_picture_data(document.id)
|
||||
await session.document_item_repository.get_all_picture_data(document.id)
|
||||
)
|
||||
|
||||
chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
|
||||
chunks = await ensure_chunks_embedded(
|
||||
session.config, chunks, session.store.embedder
|
||||
)
|
||||
|
||||
items: list[DocumentItem] | None = None
|
||||
if docling_document is not None:
|
||||
|
|
@ -194,29 +219,29 @@ async def _update_document_with_chunks(
|
|||
extract_items, document.id, docling_document, existing_picture_data
|
||||
)
|
||||
|
||||
async with client.store.write_transaction():
|
||||
updated_doc = await client.document_repository.update(document)
|
||||
async with session.store.write_transaction():
|
||||
updated_doc = await session.document_repository.update(document)
|
||||
|
||||
assert updated_doc.id is not None
|
||||
for order, chunk in enumerate(chunks):
|
||||
chunk.document_id = updated_doc.id
|
||||
chunk.order = order
|
||||
|
||||
await client.chunk_repository.replace_for_document(updated_doc.id, chunks)
|
||||
await session.chunk_repository.replace_for_document(updated_doc.id, chunks)
|
||||
|
||||
if items is not None:
|
||||
await client.document_item_repository.replace_for_document(
|
||||
await session.document_item_repository.replace_for_document(
|
||||
updated_doc.id, items
|
||||
)
|
||||
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
if session.config.storage.auto_vacuum:
|
||||
session.schedule_vacuum()
|
||||
|
||||
return updated_doc
|
||||
|
||||
|
||||
async def create_document(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
content: str,
|
||||
uri: str | None = None,
|
||||
title: str | None = None,
|
||||
|
|
@ -227,9 +252,9 @@ async def create_document(
|
|||
|
||||
Converts the content, chunks it, and generates embeddings.
|
||||
"""
|
||||
converter = get_converter(client._config)
|
||||
converter = get_converter(session.config)
|
||||
docling_document = await converter.convert_text(content, format=format)
|
||||
chunks = await client.chunk(docling_document)
|
||||
chunks = await chunk_document(session, docling_document)
|
||||
|
||||
document = Document(
|
||||
content="",
|
||||
|
|
@ -237,13 +262,15 @@ async def create_document(
|
|||
title=title,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
await _prepare_and_title(client, document, docling_document)
|
||||
await _prepare_and_title(session, document, docling_document)
|
||||
|
||||
return await _store_document_with_chunks(client, document, chunks, docling_document)
|
||||
return await _store_document_with_chunks(
|
||||
session, document, chunks, docling_document
|
||||
)
|
||||
|
||||
|
||||
async def import_document(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
docling_document: "DoclingDocument",
|
||||
chunks: list[Chunk],
|
||||
uri: str | None = None,
|
||||
|
|
@ -261,13 +288,15 @@ async def import_document(
|
|||
title=title,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
await _prepare_and_title(client, document, docling_document)
|
||||
await _prepare_and_title(session, document, docling_document)
|
||||
|
||||
return await _store_document_with_chunks(client, document, chunks, docling_document)
|
||||
return await _store_document_with_chunks(
|
||||
session, document, chunks, docling_document
|
||||
)
|
||||
|
||||
|
||||
async def _store_documents_with_chunks(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
prepared: list[tuple[Document, list[Chunk], "DoclingDocument"]],
|
||||
) -> list[Document]:
|
||||
"""Store many documents with their chunks in a single table version each.
|
||||
|
|
@ -276,9 +305,9 @@ async def _store_documents_with_chunks(
|
|||
and document_items tables once apiece. Restores all tables on any failure.
|
||||
"""
|
||||
flat = await ensure_chunks_embedded(
|
||||
client._config,
|
||||
session.config,
|
||||
[chunk for _, chunks, _ in prepared for chunk in chunks],
|
||||
client.embedder,
|
||||
session.store.embedder,
|
||||
)
|
||||
embedded: list[list[Chunk]] = []
|
||||
position = 0
|
||||
|
|
@ -291,8 +320,8 @@ async def _store_documents_with_chunks(
|
|||
|
||||
all_item_lists = await asyncio.to_thread(_extract_all_items)
|
||||
|
||||
async with client.store.write_transaction():
|
||||
created = await client.document_repository.create(
|
||||
async with session.store.write_transaction():
|
||||
created = await session.document_repository.create(
|
||||
[doc for doc, _, _ in prepared]
|
||||
)
|
||||
|
||||
|
|
@ -308,17 +337,17 @@ async def _store_documents_with_chunks(
|
|||
item.document_id = doc.id
|
||||
all_items.extend(item_list)
|
||||
|
||||
await client.chunk_repository.create(all_chunks)
|
||||
await client.document_item_repository.create_all(all_items)
|
||||
await session.chunk_repository.create(all_chunks)
|
||||
await session.document_item_repository.create_all(all_items)
|
||||
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
if session.config.storage.auto_vacuum:
|
||||
session.schedule_vacuum()
|
||||
|
||||
return created
|
||||
|
||||
|
||||
async def import_documents(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
imports: list[DocumentImport],
|
||||
) -> list[Document]:
|
||||
"""Batch-import pre-processed documents with their chunks.
|
||||
|
|
@ -338,14 +367,14 @@ async def import_documents(
|
|||
title=item.title,
|
||||
metadata=item.metadata or {},
|
||||
)
|
||||
await _prepare_and_title(client, document, item.docling_document)
|
||||
await _prepare_and_title(session, document, item.docling_document)
|
||||
prepared.append((document, item.chunks, item.docling_document))
|
||||
|
||||
return await _store_documents_with_chunks(client, prepared)
|
||||
return await _store_documents_with_chunks(session, prepared)
|
||||
|
||||
|
||||
async def _refresh_doc_metadata(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
doc: Document,
|
||||
*,
|
||||
title: str | None,
|
||||
|
|
@ -367,12 +396,12 @@ async def _refresh_doc_metadata(
|
|||
updated = True
|
||||
|
||||
if updated:
|
||||
async with client.store._write_lock:
|
||||
result = await client.document_repository.update_meta(doc)
|
||||
async with session.store._write_lock:
|
||||
result = await session.document_repository.update_meta(doc)
|
||||
# Reclaim the document_meta churn from rolling source_revision sweeps.
|
||||
# The vacuum is debounced, and document_meta is tiny, so this is cheap.
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
if session.config.storage.auto_vacuum:
|
||||
session.schedule_vacuum()
|
||||
return result
|
||||
return doc
|
||||
|
||||
|
|
@ -398,7 +427,7 @@ async def _provider_metadata(
|
|||
|
||||
|
||||
async def _ingest_fetch_result(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
result: "FetchResult",
|
||||
*,
|
||||
title: str | None,
|
||||
|
|
@ -417,7 +446,7 @@ async def _ingest_fetch_result(
|
|||
fallback. Callers pass it when ``result.uri`` cannot yield the right
|
||||
extension, e.g. embedded attachments whose name lives in a URI fragment."""
|
||||
|
||||
converter = get_converter(client._config)
|
||||
converter = get_converter(session.config)
|
||||
if filename is not None:
|
||||
file_extension = Path(filename).suffix.lower()
|
||||
else:
|
||||
|
|
@ -446,9 +475,11 @@ async def _ingest_fetch_result(
|
|||
|
||||
try:
|
||||
with logfire.span("document.convert", uri=result.uri):
|
||||
docling_document = await client.convert(target_path, source_uri=result.uri)
|
||||
docling_document = await convert(
|
||||
session.config, target_path, source_uri=result.uri
|
||||
)
|
||||
with logfire.span("document.chunk", uri=result.uri) as chunk_span:
|
||||
chunks = await client.chunk(docling_document)
|
||||
chunks = await chunk_document(session, docling_document)
|
||||
chunk_span.set_attribute("chunks_created", len(chunks))
|
||||
finally:
|
||||
if cleanup_path is not None:
|
||||
|
|
@ -460,13 +491,13 @@ async def _ingest_fetch_result(
|
|||
existing_doc.metadata = final_metadata
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
await _prepare_and_title(client, existing_doc, docling_document)
|
||||
await _prepare_and_title(session, existing_doc, docling_document)
|
||||
with logfire.span("document.store", uri=result.uri, op="update") as store_span:
|
||||
updated = await _update_document_with_chunks(
|
||||
client, existing_doc, chunks, docling_document
|
||||
session, existing_doc, chunks, docling_document
|
||||
)
|
||||
store_span.set_attribute("document_id", updated.id)
|
||||
await _reconcile_pdf_attachments(client, updated, result.body, depth=depth)
|
||||
await _reconcile_pdf_attachments(session, updated, result.body, depth=depth)
|
||||
return updated
|
||||
|
||||
document = Document(
|
||||
|
|
@ -475,13 +506,13 @@ async def _ingest_fetch_result(
|
|||
title=title,
|
||||
metadata=final_metadata,
|
||||
)
|
||||
await _prepare_and_title(client, document, docling_document)
|
||||
await _prepare_and_title(session, document, docling_document)
|
||||
with logfire.span("document.store", uri=result.uri, op="create") as store_span:
|
||||
created = await _store_document_with_chunks(
|
||||
client, document, chunks, docling_document
|
||||
session, document, chunks, docling_document
|
||||
)
|
||||
store_span.set_attribute("document_id", created.id)
|
||||
await _reconcile_pdf_attachments(client, created, result.body, depth=depth)
|
||||
await _reconcile_pdf_attachments(session, created, result.body, depth=depth)
|
||||
return created
|
||||
|
||||
|
||||
|
|
@ -545,7 +576,7 @@ def _extract_pdf_attachments(
|
|||
|
||||
|
||||
async def _reconcile_pdf_attachments(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
parent_doc: Document,
|
||||
parent_body: bytes,
|
||||
*,
|
||||
|
|
@ -559,7 +590,7 @@ async def _reconcile_pdf_attachments(
|
|||
path runs uniformly — child PDFs recurse into this helper one level deeper,
|
||||
bounded by ``MAX_ATTACHMENT_DEPTH``.
|
||||
"""
|
||||
if not client._config.processing.extract_pdf_attachments:
|
||||
if not session.config.processing.extract_pdf_attachments:
|
||||
return
|
||||
if not parent_doc.uri:
|
||||
return
|
||||
|
|
@ -572,7 +603,7 @@ async def _reconcile_pdf_attachments(
|
|||
if new_attachments is None:
|
||||
return
|
||||
|
||||
existing = await client.list_documents(filter=parent_uri_filter(parent_doc.uri))
|
||||
existing = await session.list_documents(filter=parent_uri_filter(parent_doc.uri))
|
||||
existing_by_uri: dict[str, Document] = {d.uri: d for d in existing if d.uri}
|
||||
|
||||
for child_uri, (name, data, content_type, content_hash) in new_attachments.items():
|
||||
|
|
@ -594,7 +625,7 @@ async def _reconcile_pdf_attachments(
|
|||
)
|
||||
try:
|
||||
await _ingest_fetch_result(
|
||||
client,
|
||||
session,
|
||||
child_fr,
|
||||
title=None,
|
||||
user_metadata={},
|
||||
|
|
@ -615,11 +646,11 @@ async def _reconcile_pdf_attachments(
|
|||
|
||||
for child_uri, child in existing_by_uri.items():
|
||||
if child_uri not in new_attachments and child.id:
|
||||
await client.delete_document(child.id)
|
||||
await session.delete_document(child.id)
|
||||
|
||||
|
||||
async def create_document_from_source(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
source: str | Path,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
|
|
@ -674,7 +705,7 @@ async def create_document_from_source(
|
|||
for child in walk_files(local_path):
|
||||
if child.is_file() and filter.include_file(str(child)):
|
||||
doc = await create_document_from_source(
|
||||
client,
|
||||
session,
|
||||
child,
|
||||
title=None,
|
||||
metadata=metadata,
|
||||
|
|
@ -692,7 +723,7 @@ async def create_document_from_source(
|
|||
|
||||
# Match the old _create_document_from_file behaviour: fail fast on
|
||||
# unsupported extension before reading any bytes.
|
||||
converter = get_converter(client._config)
|
||||
converter = get_converter(session.config)
|
||||
if local_path.suffix.lower() not in converter.supported_extensions:
|
||||
raise UnsupportedSourceError(
|
||||
f"Unsupported file extension: {local_path.suffix}"
|
||||
|
|
@ -733,7 +764,7 @@ async def create_document_from_source(
|
|||
else:
|
||||
stored_uri = source_str
|
||||
|
||||
existing_doc = await client.get_document_by_uri(stored_uri)
|
||||
existing_doc = await session.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
|
||||
|
|
@ -748,7 +779,7 @@ async def create_document_from_source(
|
|||
current_revision = await fetcher.head(source_str)
|
||||
if current_revision == stored_revision:
|
||||
return await _refresh_doc_metadata(
|
||||
client,
|
||||
session,
|
||||
existing_doc,
|
||||
title=title,
|
||||
user_metadata=metadata,
|
||||
|
|
@ -781,7 +812,7 @@ async def create_document_from_source(
|
|||
if result.revision is not None:
|
||||
source_meta["source_revision"] = result.revision
|
||||
return await _refresh_doc_metadata(
|
||||
client,
|
||||
session,
|
||||
existing_doc,
|
||||
title=title,
|
||||
user_metadata=user_metadata,
|
||||
|
|
@ -789,7 +820,7 @@ async def create_document_from_source(
|
|||
)
|
||||
|
||||
return await _ingest_fetch_result(
|
||||
client,
|
||||
session,
|
||||
result,
|
||||
title=title,
|
||||
user_metadata=user_metadata,
|
||||
|
|
@ -802,7 +833,7 @@ async def create_document_from_source(
|
|||
|
||||
|
||||
async def update_document(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
document_id: str,
|
||||
content: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
|
|
@ -830,7 +861,7 @@ async def update_document(
|
|||
|
||||
# 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(
|
||||
existing_doc = await session.document_repository.get_by_id(
|
||||
document_id, include_blobs=chunks is not None and docling_document is None
|
||||
)
|
||||
if existing_doc is None:
|
||||
|
|
@ -844,10 +875,10 @@ async def update_document(
|
|||
existing_doc.uri = uri
|
||||
|
||||
if content is None and chunks is None and docling_document is None:
|
||||
async with client.store._write_lock:
|
||||
updated = await client.document_repository.update_meta(existing_doc)
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
async with session.store._write_lock:
|
||||
updated = await session.document_repository.update_meta(existing_doc)
|
||||
if session.config.storage.auto_vacuum:
|
||||
session.schedule_vacuum()
|
||||
return updated
|
||||
|
||||
if chunks is not None:
|
||||
|
|
@ -857,26 +888,26 @@ async def update_document(
|
|||
existing_doc.content = content
|
||||
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, chunks, docling_document
|
||||
session, existing_doc, chunks, docling_document
|
||||
)
|
||||
|
||||
if docling_document is not None:
|
||||
await _prepare_document_from_docling(existing_doc, docling_document)
|
||||
|
||||
new_chunks = await client.chunk(docling_document)
|
||||
new_chunks = await chunk_document(session, docling_document)
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, new_chunks, docling_document
|
||||
session, existing_doc, new_chunks, docling_document
|
||||
)
|
||||
|
||||
assert content is not None
|
||||
existing_doc.content = content
|
||||
converter = get_converter(client._config)
|
||||
converter = get_converter(session.config)
|
||||
converted_docling = await converter.convert_text(existing_doc.content, format="md")
|
||||
await _prepare_document_from_docling(existing_doc, converted_docling)
|
||||
|
||||
new_chunks = await client.chunk(converted_docling)
|
||||
new_chunks = await chunk_document(session, converted_docling)
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, new_chunks, converted_docling
|
||||
session, existing_doc, new_chunks, converted_docling
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,8 +10,11 @@ from lancedb.pydantic import LanceModel
|
|||
|
||||
from haiku.rag.client.documents import (
|
||||
check_source_accessible,
|
||||
chunk_document,
|
||||
create_document_from_source,
|
||||
)
|
||||
from haiku.rag.client.session import SingleDatabaseSession
|
||||
from haiku.rag.client.titles import generate_title
|
||||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.store.compression import compress_docling_split
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
|
|
@ -23,7 +26,7 @@ from haiku.rag.store.schema import ChunkRecordBase
|
|||
if TYPE_CHECKING:
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -69,7 +72,7 @@ class _StagingMarkerRecord(LanceModel):
|
|||
|
||||
|
||||
async def rebuild_database(
|
||||
client: "HaikuRAG", mode: "RebuildMode"
|
||||
session: SingleDatabaseSession, mode: "RebuildMode"
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Rebuild the database with the specified mode.
|
||||
|
||||
|
|
@ -82,17 +85,17 @@ async def rebuild_database(
|
|||
"""
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
async with client.store._rebuild_lock:
|
||||
async with session.store._rebuild_lock:
|
||||
if mode == RebuildMode.SET_EMBEDDER:
|
||||
await _set_embedder(client)
|
||||
await _set_embedder(session)
|
||||
return
|
||||
|
||||
async for doc_id in _rebuild_locked(client, mode):
|
||||
async for doc_id in _rebuild_locked(session, mode):
|
||||
yield doc_id
|
||||
|
||||
|
||||
async def _rebuild_locked(
|
||||
client: "HaikuRAG", mode: "RebuildMode"
|
||||
session: SingleDatabaseSession, mode: "RebuildMode"
|
||||
) -> AsyncGenerator[str, None]:
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
|
|
@ -100,7 +103,7 @@ async def _rebuild_locked(
|
|||
# interrupted rebuild. Returns True only when phase 1 was already
|
||||
# complete and the current mode is EMBED_ONLY, in which case we resume
|
||||
# phase 2 from the existing staging table instead of recopying.
|
||||
resume_from_staging = await _resolve_rebuild_recovery(client, mode)
|
||||
resume_from_staging = await _resolve_rebuild_recovery(session, mode)
|
||||
|
||||
# Wait for any already-scheduled background vacuum before the destructive
|
||||
# table operations at the top of RECHUNK / FULL. Rebuild drops and
|
||||
|
|
@ -109,52 +112,52 @@ async def _rebuild_locked(
|
|||
# lance. Note: FULL calls create_document_from_source inside its loop,
|
||||
# which may schedule *new* background vacuums — those run after the
|
||||
# destructive phase and are fine.
|
||||
await client._await_vacuum_tasks()
|
||||
await session.drain_vacuum()
|
||||
|
||||
settings_repo = SettingsRepository(client.store)
|
||||
settings_repo = SettingsRepository(session.store)
|
||||
await settings_repo.save_current_settings()
|
||||
|
||||
# Light listing — id/uri/title/metadata only. Each rebuild function
|
||||
# fetches full content (including the multi-MB docling_pages blob) one
|
||||
# document at a time so a 1000-doc database doesn't pull ~15 GB of
|
||||
# blobs into memory before the loop starts.
|
||||
documents = await client.list_documents(include_content=False)
|
||||
documents = await session.list_documents(include_content=False)
|
||||
|
||||
if mode == RebuildMode.TITLE_ONLY:
|
||||
async for doc_id in _rebuild_title_only(client, documents):
|
||||
async for doc_id in _rebuild_title_only(session, documents):
|
||||
yield doc_id
|
||||
elif mode == RebuildMode.EMBED_ONLY:
|
||||
async for doc_id in _rebuild_embed_only(
|
||||
client, documents, resume_from_staging=resume_from_staging
|
||||
session, documents, resume_from_staging=resume_from_staging
|
||||
):
|
||||
yield doc_id
|
||||
elif mode == RebuildMode.RECHUNK:
|
||||
await client.chunk_repository.delete_all()
|
||||
await client.store.recreate_embeddings_table()
|
||||
async for doc_id in _rebuild_rechunk(client, documents):
|
||||
await session.chunk_repository.delete_all()
|
||||
await session.store.recreate_embeddings_table()
|
||||
async for doc_id in _rebuild_rechunk(session, documents):
|
||||
yield doc_id
|
||||
elif mode == RebuildMode.DESCRIPTIONS:
|
||||
await client.chunk_repository.delete_all()
|
||||
await client.store.recreate_embeddings_table()
|
||||
async for doc_id in _rebuild_descriptions(client, documents):
|
||||
await session.chunk_repository.delete_all()
|
||||
await session.store.recreate_embeddings_table()
|
||||
async for doc_id in _rebuild_descriptions(session, documents):
|
||||
yield doc_id
|
||||
else: # FULL
|
||||
await client.chunk_repository.delete_all()
|
||||
await client.store.recreate_embeddings_table()
|
||||
async for doc_id in _rebuild_full(client, documents):
|
||||
await session.chunk_repository.delete_all()
|
||||
await session.store.recreate_embeddings_table()
|
||||
async for doc_id in _rebuild_full(session, documents):
|
||||
yield doc_id
|
||||
|
||||
# Final maintenance if auto_vacuum enabled. Swallowing only so that a
|
||||
# failed post-rebuild optimize doesn't mask a successful rebuild — but
|
||||
# log it so the failure is visible in the output.
|
||||
if client._config.storage.auto_vacuum:
|
||||
if session.config.storage.auto_vacuum:
|
||||
try:
|
||||
await client.store.vacuum()
|
||||
await session.store.vacuum()
|
||||
except Exception:
|
||||
logger.warning("Post-rebuild vacuum failed", exc_info=True)
|
||||
|
||||
|
||||
async def _set_embedder(client: "HaikuRAG") -> None:
|
||||
async def _set_embedder(session: SingleDatabaseSession) -> None:
|
||||
"""Adopt the current embedder identity without re-embedding.
|
||||
|
||||
Only valid when the vector dimension is unchanged — the stored vectors stay
|
||||
|
|
@ -163,10 +166,10 @@ async def _set_embedder(client: "HaikuRAG") -> None:
|
|||
"""
|
||||
from haiku.rag.store.exceptions import ConfigMismatchError
|
||||
|
||||
settings_repo = SettingsRepository(client.store)
|
||||
settings_repo = SettingsRepository(session.store)
|
||||
stored = await settings_repo.get_current_settings()
|
||||
stored_dim = stored.get("embeddings", {}).get("model", {}).get("vector_dim")
|
||||
current_dim = client._config.embeddings.model.vector_dim
|
||||
current_dim = session.config.embeddings.model.vector_dim
|
||||
|
||||
if stored_dim is not None and current_dim != stored_dim:
|
||||
raise ConfigMismatchError(
|
||||
|
|
@ -178,7 +181,9 @@ async def _set_embedder(client: "HaikuRAG") -> None:
|
|||
|
||||
|
||||
async def _hydrate(
|
||||
client: "HaikuRAG", light_docs: list[Document], include_blobs: bool = True
|
||||
session: SingleDatabaseSession,
|
||||
light_docs: list[Document],
|
||||
include_blobs: bool = True,
|
||||
) -> AsyncGenerator[Document, None]:
|
||||
"""Yield fully-loaded documents one at a time from a light listing.
|
||||
|
||||
|
|
@ -189,7 +194,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(
|
||||
doc = await session.document_repository.get_by_id(
|
||||
light_doc.id, include_blobs=include_blobs
|
||||
)
|
||||
if doc is None:
|
||||
|
|
@ -199,22 +204,22 @@ async def _hydrate(
|
|||
|
||||
|
||||
async def _rebuild_title_only(
|
||||
client: "HaikuRAG", documents: list[Document]
|
||||
session: SingleDatabaseSession, documents: list[Document]
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""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
|
||||
repo = session.document_repository
|
||||
untitled = [d for d in documents if d.title is None]
|
||||
async for doc in _hydrate(client, untitled, include_blobs=False):
|
||||
async for doc in _hydrate(session, 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)
|
||||
title = await generate_title(session.config, doc)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to generate title for document %s", doc.id, exc_info=True
|
||||
|
|
@ -226,7 +231,9 @@ async def _rebuild_title_only(
|
|||
yield doc.id
|
||||
|
||||
|
||||
async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") -> bool:
|
||||
async def _resolve_rebuild_recovery(
|
||||
session: SingleDatabaseSession, mode: "RebuildMode"
|
||||
) -> bool:
|
||||
"""Resolve any partially-completed rebuild state from a previous crash.
|
||||
|
||||
Returns ``True`` if ``_rebuild_embed_only`` should resume from the
|
||||
|
|
@ -243,7 +250,7 @@ async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") ->
|
|||
"""
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
db = client.store.db
|
||||
db = session.store.db
|
||||
tables = (await db.list_tables()).tables
|
||||
has_staging = _STAGING_TABLE_NAME in tables
|
||||
has_marker = _STAGING_MARKER_TABLE_NAME in tables
|
||||
|
|
@ -286,7 +293,7 @@ async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") ->
|
|||
return False
|
||||
|
||||
|
||||
async def _populate_staging_table(client: "HaikuRAG") -> None:
|
||||
async def _populate_staging_table(session: SingleDatabaseSession) -> None:
|
||||
"""Stream the non-vector columns of the chunks table into staging.
|
||||
|
||||
Uses ``to_batches`` for a single streaming read (no offset/limit
|
||||
|
|
@ -297,7 +304,7 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
|
|||
Requires ``_resolve_rebuild_recovery`` to have cleared any leftover
|
||||
staging table first: ``create_table`` raises if the name is already taken.
|
||||
"""
|
||||
db = client.store.db
|
||||
db = session.store.db
|
||||
tables = (await db.list_tables()).tables
|
||||
|
||||
staging = await db.create_table(_STAGING_TABLE_NAME, schema=_StagingChunkRecord)
|
||||
|
|
@ -305,7 +312,7 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
|
|||
return
|
||||
|
||||
stream = (
|
||||
await client.store.chunks_table.query()
|
||||
await session.store.chunks_table.query()
|
||||
.select(["id", "document_id", "content", "metadata", "order"])
|
||||
.to_batches(max_batch_length=_STAGING_COPY_BATCH_SIZE)
|
||||
)
|
||||
|
|
@ -324,13 +331,13 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
|
|||
await staging.add(records)
|
||||
|
||||
|
||||
async def _mark_phase1_complete(client: "HaikuRAG") -> None:
|
||||
async def _mark_phase1_complete(session: SingleDatabaseSession) -> None:
|
||||
"""Create the marker table that designates staging as authoritative.
|
||||
|
||||
Called after ``_populate_staging_table`` finishes. On crash recovery the
|
||||
marker's presence flips ``_rebuild_embed_only`` into resume mode.
|
||||
"""
|
||||
db = client.store.db
|
||||
db = session.store.db
|
||||
if _STAGING_MARKER_TABLE_NAME in (await db.list_tables()).tables:
|
||||
return
|
||||
marker = await db.create_table(
|
||||
|
|
@ -339,7 +346,7 @@ async def _mark_phase1_complete(client: "HaikuRAG") -> None:
|
|||
await marker.add([_StagingMarkerRecord(id="phase1_complete")])
|
||||
|
||||
|
||||
async def _drop_staging_tables(client: "HaikuRAG") -> None:
|
||||
async def _drop_staging_tables(session: SingleDatabaseSession) -> None:
|
||||
"""Drop the marker first, then the staging table.
|
||||
|
||||
Ordering matters: if a crash interrupts cleanup between the two drops,
|
||||
|
|
@ -347,7 +354,7 @@ async def _drop_staging_tables(client: "HaikuRAG") -> None:
|
|||
partial phase 1 → drops staging harmlessly. The reverse order would
|
||||
leak a marker pointing at nothing.
|
||||
"""
|
||||
db = client.store.db
|
||||
db = session.store.db
|
||||
tables = (await db.list_tables()).tables
|
||||
if _STAGING_MARKER_TABLE_NAME in tables:
|
||||
await db.drop_table(_STAGING_MARKER_TABLE_NAME)
|
||||
|
|
@ -384,7 +391,7 @@ async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chu
|
|||
|
||||
|
||||
async def _rebuild_embed_only(
|
||||
client: "HaikuRAG",
|
||||
session: SingleDatabaseSession,
|
||||
documents: list[Document],
|
||||
*,
|
||||
resume_from_staging: bool = False,
|
||||
|
|
@ -414,19 +421,19 @@ async def _rebuild_embed_only(
|
|||
"""
|
||||
from haiku.rag.embeddings import contextualize, embed_chunks
|
||||
|
||||
db = client.store.db
|
||||
embedder = client.chunk_repository.embedder
|
||||
db = session.store.db
|
||||
embedder = session.chunk_repository.embedder
|
||||
|
||||
if not resume_from_staging:
|
||||
# Phase 1: copy chunks into staging, then mark it complete. After the
|
||||
# marker exists, a crash will resume phase 2 from staging.
|
||||
await _populate_staging_table(client)
|
||||
await _mark_phase1_complete(client)
|
||||
await _populate_staging_table(session)
|
||||
await _mark_phase1_complete(session)
|
||||
|
||||
# Recreate the chunks table fresh (idempotent; handles vector-dim
|
||||
# changes and discards any partial new chunks from a prior crashed
|
||||
# phase 2).
|
||||
await client.store.recreate_embeddings_table()
|
||||
await session.store.recreate_embeddings_table()
|
||||
|
||||
staging_table = await db.open_table(_STAGING_TABLE_NAME)
|
||||
|
||||
|
|
@ -443,7 +450,7 @@ async def _rebuild_embed_only(
|
|||
# chunks route through embed_image rather than being text-embedded.
|
||||
# Bytes live in document_items (embed-only never touches that table).
|
||||
if embedder.supports_images:
|
||||
picture_data = await client.document_item_repository.get_all_picture_data(
|
||||
picture_data = await session.document_item_repository.get_all_picture_data(
|
||||
doc.id
|
||||
)
|
||||
for chunk in chunks:
|
||||
|
|
@ -462,7 +469,7 @@ async def _rebuild_embed_only(
|
|||
)
|
||||
|
||||
content_fts_list = contextualize(chunks)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, client._config)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, session.config)
|
||||
|
||||
for chunk, content_fts, embedded in zip(
|
||||
chunks, content_fts_list, embedded_chunks
|
||||
|
|
@ -471,7 +478,7 @@ async def _rebuild_embed_only(
|
|||
assert chunk.document_id is not None
|
||||
assert embedded.embedding is not None
|
||||
pending_records.append(
|
||||
client.store.ChunkRecord(
|
||||
session.store.ChunkRecord(
|
||||
id=chunk.id,
|
||||
document_id=chunk.document_id,
|
||||
content=chunk.content,
|
||||
|
|
@ -491,16 +498,16 @@ async def _rebuild_embed_only(
|
|||
yield doc.id
|
||||
|
||||
if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
|
||||
await client.store.chunks_table.add(pending_records)
|
||||
await session.store.chunks_table.add(pending_records)
|
||||
pending_records = []
|
||||
|
||||
if pending_records:
|
||||
await client.store.chunks_table.add(pending_records)
|
||||
await session.store.chunks_table.add(pending_records)
|
||||
|
||||
# Phase 2 finished. Drop the recovery state — marker first so a crash
|
||||
# between the two drops leaves only staging behind, which the next
|
||||
# rebuild discards harmlessly.
|
||||
await _drop_staging_tables(client)
|
||||
await _drop_staging_tables(session)
|
||||
|
||||
for doc in documents:
|
||||
if doc.id and doc.id not in yielded_docs:
|
||||
|
|
@ -508,7 +515,7 @@ async def _rebuild_embed_only(
|
|||
|
||||
|
||||
async def _flush_rebuild_batch(
|
||||
client: "HaikuRAG", documents: list[Document], chunks: list[Chunk]
|
||||
session: SingleDatabaseSession, documents: list[Document], chunks: list[Chunk]
|
||||
) -> None:
|
||||
"""Batch write documents and chunks during rebuild.
|
||||
|
||||
|
|
@ -552,12 +559,12 @@ async def _flush_rebuild_batch(
|
|||
)
|
||||
|
||||
await (
|
||||
client.store.documents_table.merge_insert("id")
|
||||
session.store.documents_table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.execute(doc_records)
|
||||
)
|
||||
await (
|
||||
client.store.document_meta_table.merge_insert("id")
|
||||
session.store.document_meta_table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(meta_records)
|
||||
|
|
@ -565,7 +572,7 @@ async def _flush_rebuild_batch(
|
|||
|
||||
# Batch create all chunks (single LanceDB version)
|
||||
if chunks:
|
||||
await client.chunk_repository.create(chunks)
|
||||
await session.chunk_repository.create(chunks)
|
||||
|
||||
# Repopulate document items from stored docling data. The stored docling
|
||||
# blob has had its picture URIs stripped (compress_docling_split), so
|
||||
|
|
@ -576,28 +583,28 @@ async def _flush_rebuild_batch(
|
|||
docling_doc = doc.get_docling_document()
|
||||
if docling_doc is not None:
|
||||
existing_picture_data = (
|
||||
await client.document_item_repository.get_all_picture_data(doc.id)
|
||||
await session.document_item_repository.get_all_picture_data(doc.id)
|
||||
)
|
||||
await client.document_item_repository.delete_by_document_id(doc.id)
|
||||
await session.document_item_repository.delete_by_document_id(doc.id)
|
||||
items = extract_items(
|
||||
doc.id,
|
||||
docling_doc,
|
||||
existing_picture_data=existing_picture_data,
|
||||
)
|
||||
await client.document_item_repository.create_items(doc.id, items)
|
||||
await session.document_item_repository.create_items(doc.id, items)
|
||||
|
||||
|
||||
async def _rebuild_rechunk(
|
||||
client: "HaikuRAG", documents: list[Document]
|
||||
session: SingleDatabaseSession, documents: list[Document]
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Re-chunk and re-embed each document from its stored docling blob."""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
embedder = client.embedder
|
||||
embedder = session.store.embedder
|
||||
|
||||
async for doc in _hydrate(client, documents):
|
||||
async for doc in _hydrate(session, documents):
|
||||
assert doc.id is not None
|
||||
docling_document = doc.get_docling_document()
|
||||
if docling_document is None:
|
||||
|
|
@ -609,16 +616,17 @@ async def _rebuild_rechunk(
|
|||
# Stored blob has stripped picture URIs; pass the snapshot so
|
||||
# build_picture_chunks (inside chunk()) can recover the bytes.
|
||||
existing_picture_data = (
|
||||
await client.document_item_repository.get_all_picture_data(doc.id)
|
||||
await session.document_item_repository.get_all_picture_data(doc.id)
|
||||
if embedder.supports_images
|
||||
else None
|
||||
)
|
||||
chunks = await client.chunk(
|
||||
chunks = await chunk_document(
|
||||
session,
|
||||
docling_document,
|
||||
existing_picture_data=existing_picture_data,
|
||||
document_id=doc.id,
|
||||
)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, client._config)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, session.config)
|
||||
|
||||
for order, chunk in enumerate(embedded_chunks):
|
||||
chunk.document_id = doc.id
|
||||
|
|
@ -634,12 +642,12 @@ async def _rebuild_rechunk(
|
|||
yield doc.id
|
||||
|
||||
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
pending_chunks = []
|
||||
pending_docs = []
|
||||
|
||||
if pending_docs:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
|
||||
|
||||
def _apply_descriptions_sync(
|
||||
|
|
@ -668,7 +676,9 @@ def _apply_descriptions_sync(
|
|||
return len(descriptions)
|
||||
|
||||
|
||||
async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
|
||||
async def _patch_picture_descriptions(
|
||||
session: SingleDatabaseSession, doc: Document
|
||||
) -> int:
|
||||
"""Run the VLM against pictures lacking a description, patch the docling
|
||||
blob in-place. Returns the number of newly described pictures.
|
||||
Pictures that already carry ``meta.description.text`` are skipped, so the
|
||||
|
|
@ -692,7 +702,7 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
|
|||
if not needs_description:
|
||||
return 0
|
||||
|
||||
bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk(
|
||||
bytes_by_ref = await session.document_item_repository.get_pictures_for_chunk(
|
||||
doc.id, needs_description
|
||||
)
|
||||
if not bytes_by_ref:
|
||||
|
|
@ -705,7 +715,7 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
|
|||
)
|
||||
return 0
|
||||
|
||||
descriptions = await describe_pictures(bytes_by_ref, config=client._config)
|
||||
descriptions = await describe_pictures(bytes_by_ref, config=session.config)
|
||||
|
||||
if not descriptions:
|
||||
return 0
|
||||
|
|
@ -716,7 +726,7 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
|
|||
|
||||
|
||||
async def _rebuild_descriptions(
|
||||
client: "HaikuRAG", documents: list[Document]
|
||||
session: SingleDatabaseSession, documents: list[Document]
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Run the VLM over already-stored picture bytes, patch descriptions into
|
||||
the docling blob, then re-chunk + re-embed.
|
||||
|
|
@ -727,7 +737,7 @@ async def _rebuild_descriptions(
|
|||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
if client._config.processing.pictures != "description":
|
||||
if session.config.processing.pictures != "description":
|
||||
raise ValueError(
|
||||
"rebuild --descriptions requires processing.pictures = 'description' "
|
||||
"in your config."
|
||||
|
|
@ -735,10 +745,10 @@ async def _rebuild_descriptions(
|
|||
|
||||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
embedder = client.embedder
|
||||
embedder = session.store.embedder
|
||||
|
||||
described_total = 0
|
||||
async for doc in _hydrate(client, documents):
|
||||
async for doc in _hydrate(session, documents):
|
||||
assert doc.id is not None
|
||||
docling_document = doc.get_docling_document()
|
||||
if docling_document is None:
|
||||
|
|
@ -747,23 +757,24 @@ async def _rebuild_descriptions(
|
|||
"rebuild --descriptions requires it. Run a full rebuild instead."
|
||||
)
|
||||
|
||||
n = await _patch_picture_descriptions(client, doc)
|
||||
n = await _patch_picture_descriptions(session, doc)
|
||||
described_total += n
|
||||
# Use the (possibly patched) docling document for chunking.
|
||||
docling_document = doc.get_docling_document()
|
||||
assert docling_document is not None
|
||||
|
||||
existing_picture_data = (
|
||||
await client.document_item_repository.get_all_picture_data(doc.id)
|
||||
await session.document_item_repository.get_all_picture_data(doc.id)
|
||||
if embedder.supports_images
|
||||
else None
|
||||
)
|
||||
chunks = await client.chunk(
|
||||
chunks = await chunk_document(
|
||||
session,
|
||||
docling_document,
|
||||
existing_picture_data=existing_picture_data,
|
||||
document_id=doc.id,
|
||||
)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, client._config)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, session.config)
|
||||
|
||||
for order, chunk in enumerate(embedded_chunks):
|
||||
chunk.document_id = doc.id
|
||||
|
|
@ -774,12 +785,12 @@ async def _rebuild_descriptions(
|
|||
yield doc.id
|
||||
|
||||
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
pending_chunks = []
|
||||
pending_docs = []
|
||||
|
||||
if pending_docs:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
|
||||
logger.info(
|
||||
"rebuild --descriptions: %d new picture descriptions added across %d documents",
|
||||
|
|
@ -789,15 +800,15 @@ async def _rebuild_descriptions(
|
|||
|
||||
|
||||
async def _rebuild_full(
|
||||
client: "HaikuRAG", documents: list[Document]
|
||||
session: SingleDatabaseSession, documents: list[Document]
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Full rebuild: re-convert from source, re-chunk, re-embed."""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
converter = get_converter(client._config)
|
||||
embedder = client.embedder
|
||||
converter = get_converter(session.config)
|
||||
embedder = session.store.embedder
|
||||
|
||||
for light_doc in documents:
|
||||
assert light_doc.id is not None
|
||||
|
|
@ -805,10 +816,10 @@ 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):
|
||||
# The refresh writes through the client, not the batch buffer, so
|
||||
# The refresh writes through the database, not the batch buffer, so
|
||||
# anything pending has to land first.
|
||||
if pending_docs:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
pending_chunks = []
|
||||
pending_docs = []
|
||||
|
||||
|
|
@ -817,7 +828,7 @@ async def _rebuild_full(
|
|||
# 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,
|
||||
session,
|
||||
source=light_doc.uri,
|
||||
metadata=light_doc.metadata or {},
|
||||
force=True,
|
||||
|
|
@ -840,7 +851,7 @@ async def _rebuild_full(
|
|||
|
||||
# Fallback: rebuild from stored content. Now we need the full
|
||||
# record (content + docling_pages for the round-trip write).
|
||||
doc = await client.document_repository.get_by_id(
|
||||
doc = await session.document_repository.get_by_id(
|
||||
light_doc.id, include_blobs=True
|
||||
)
|
||||
if doc is None:
|
||||
|
|
@ -848,8 +859,8 @@ async def _rebuild_full(
|
|||
assert doc.id is not None
|
||||
|
||||
docling_document = await converter.convert_text(doc.content, format="md")
|
||||
chunks = await client.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, client._config)
|
||||
chunks = await chunk_document(session, docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, embedder, session.config)
|
||||
|
||||
doc.set_docling(docling_document)
|
||||
|
||||
|
|
@ -862,9 +873,9 @@ async def _rebuild_full(
|
|||
yield doc.id
|
||||
|
||||
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
pending_chunks = []
|
||||
pending_docs = []
|
||||
|
||||
if pending_docs:
|
||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||
await _flush_rebuild_batch(session, pending_docs, pending_chunks)
|
||||
|
|
|
|||
|
|
@ -30,6 +30,9 @@ from .services import reachable # noqa: E402
|
|||
if TYPE_CHECKING:
|
||||
from vcr import VCR
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.client.session import SingleDatabaseSession
|
||||
|
||||
setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
|
||||
logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
|
||||
|
||||
|
|
@ -216,3 +219,14 @@ def docling_serve_url() -> str:
|
|||
if not reachable("localhost", 5001):
|
||||
pytest.skip(f"docling-serve not reachable on localhost:5001 — {_COMPOSE_HINT}")
|
||||
return "http://localhost:5001"
|
||||
|
||||
|
||||
def writing(client: "HaikuRAG") -> "SingleDatabaseSession":
|
||||
"""The database a write implementation works on, from a client holding one.
|
||||
|
||||
Write implementations take a session rather than a client, so a set can
|
||||
never reach them. Tests that call one directly go through here."""
|
||||
from haiku.rag.client.session import SingleDatabaseSession
|
||||
|
||||
assert isinstance(client._session, SingleDatabaseSession)
|
||||
return client._session
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from haiku.rag.store.models.document_item import (
|
|||
extract_items,
|
||||
)
|
||||
from haiku.rag.store.repositories.document_item import DocumentItemRepository
|
||||
from tests.conftest import writing
|
||||
|
||||
|
||||
def _make_docling_doc():
|
||||
|
|
@ -416,7 +417,9 @@ class TestDocumentItemPopulation:
|
|||
|
||||
# Use _store_document_with_chunks directly with empty chunks
|
||||
# to avoid needing embeddings
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
count = await rag.document_item_repository.get_item_count(created.id)
|
||||
|
|
@ -446,7 +449,9 @@ class TestDocumentItemPopulation:
|
|||
uri="test://doc",
|
||||
)
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
assert await rag.document_item_repository.get_item_count(created.id) == 6
|
||||
|
||||
|
|
@ -455,7 +460,7 @@ class TestDocumentItemPopulation:
|
|||
new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
|
||||
created.set_docling(new_doc)
|
||||
|
||||
await _update_document_with_chunks(rag, created, [], new_doc)
|
||||
await _update_document_with_chunks(writing(rag), created, [], new_doc)
|
||||
assert await rag.document_item_repository.get_item_count(created.id) == 1
|
||||
|
||||
async def test_delete_document_cascades_items(self, temp_db_path):
|
||||
|
|
@ -470,7 +475,9 @@ class TestDocumentItemPopulation:
|
|||
uri="test://doc",
|
||||
)
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
assert await rag.document_item_repository.get_item_count(created.id) == 6
|
||||
|
||||
|
|
@ -846,7 +853,9 @@ class TestPictureDataPreservedThroughRoundTrip:
|
|||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||
document = Document(content="Hello world", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
original = await rag.document_item_repository.get_all_picture_data(
|
||||
|
|
@ -859,7 +868,7 @@ class TestPictureDataPreservedThroughRoundTrip:
|
|||
assert from_blob is not None
|
||||
assert all(p.image is None for p in from_blob.pictures)
|
||||
|
||||
await _update_document_with_chunks(rag, created, [], from_blob)
|
||||
await _update_document_with_chunks(writing(rag), created, [], from_blob)
|
||||
|
||||
after = await rag.document_item_repository.get_all_picture_data(created.id)
|
||||
assert after.get("#/pictures/0") == original.get("#/pictures/0")
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from haiku.rag.context import (
|
|||
)
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.document_item import DocumentItem
|
||||
from tests.conftest import writing
|
||||
|
||||
|
||||
def _item(
|
||||
|
|
@ -352,7 +353,7 @@ class TestExpandWithItems:
|
|||
|
||||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||
doc = await _store_document_with_chunks(
|
||||
rag,
|
||||
writing(rag),
|
||||
Document(content="test"),
|
||||
[],
|
||||
__import__(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from haiku.rag.client.documents import _store_document_with_chunks
|
|||
from haiku.rag.client.processing import ensure_chunks_embedded
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models import SearchResult
|
||||
from tests.conftest import writing
|
||||
|
||||
|
||||
async def create_document_with_docling(
|
||||
|
|
@ -318,7 +319,9 @@ async def test_expand_context_single_item_document(temp_db_path):
|
|||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
document = Document(content="Simple test content")
|
||||
document.set_docling(docling_doc)
|
||||
doc = await _store_document_with_chunks(client, document, [], docling_doc)
|
||||
doc = await _store_document_with_chunks(
|
||||
writing(client), document, [], docling_doc
|
||||
)
|
||||
assert doc.id is not None
|
||||
|
||||
# Create a search result with a doc_item_ref pointing to the item
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from haiku.rag.client.documents import (
|
|||
)
|
||||
from haiku.rag.sources import FetchResult
|
||||
from haiku.rag.store.models.document import Document
|
||||
from tests.conftest import writing
|
||||
|
||||
|
||||
def build_pdf(attachments: list[tuple[str, bytes]]) -> bytes:
|
||||
|
|
@ -27,7 +28,7 @@ def build_pdf(attachments: list[tuple[str, bytes]]) -> bytes:
|
|||
|
||||
|
||||
async def fake_ingest_fetch_result(
|
||||
client,
|
||||
session,
|
||||
result: FetchResult,
|
||||
*,
|
||||
title,
|
||||
|
|
@ -55,9 +56,9 @@ async def fake_ingest_fetch_result(
|
|||
existing_doc.metadata = final_metadata
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
doc = await client.document_repository.update(existing_doc)
|
||||
doc = await session.document_repository.update(existing_doc)
|
||||
else:
|
||||
doc = await client.document_repository.create(
|
||||
doc = await session.document_repository.create(
|
||||
Document(
|
||||
content="",
|
||||
uri=stored_uri,
|
||||
|
|
@ -65,7 +66,7 @@ async def fake_ingest_fetch_result(
|
|||
metadata=final_metadata,
|
||||
)
|
||||
)
|
||||
await _reconcile_pdf_attachments(client, doc, result.body, depth=depth)
|
||||
await _reconcile_pdf_attachments(session, doc, result.body, depth=depth)
|
||||
return doc
|
||||
|
||||
|
||||
|
|
@ -101,7 +102,7 @@ async def test_first_ingest_creates_one_doc_per_attachment(temp_db_path, monkeyp
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
|
||||
children = await client.list_documents(filter=parent_uri_filter(parent_uri))
|
||||
assert len(children) == 2
|
||||
|
|
@ -129,7 +130,7 @@ async def test_attachment_with_spaces_in_name_is_percent_encoded(
|
|||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
|
||||
children = await client.list_documents(filter=parent_uri_filter(parent_uri))
|
||||
assert len(children) == 1
|
||||
|
|
@ -145,13 +146,13 @@ async def test_reingest_removes_dropped_attachment(temp_db_path, monkeypatch):
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
first = build_pdf([("a.txt", b"A"), ("b.txt", b"B")])
|
||||
parent = await _make_parent(client, parent_uri, first)
|
||||
await _reconcile_pdf_attachments(client, parent, first, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, first, depth=0)
|
||||
assert (
|
||||
len(await client.list_documents(filter=parent_uri_filter(parent_uri))) == 2
|
||||
)
|
||||
|
||||
second = build_pdf([("a.txt", b"A")])
|
||||
await _reconcile_pdf_attachments(client, parent, second, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, second, depth=0)
|
||||
remaining = await client.list_documents(filter=parent_uri_filter(parent_uri))
|
||||
assert len(remaining) == 1
|
||||
assert remaining[0].uri == f"{parent_uri}#attachment=a.txt"
|
||||
|
|
@ -166,11 +167,11 @@ async def test_reingest_updates_changed_attachment_in_place(temp_db_path, monkey
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
first = build_pdf([("a.txt", b"original")])
|
||||
parent = await _make_parent(client, parent_uri, first)
|
||||
await _reconcile_pdf_attachments(client, parent, first, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, first, depth=0)
|
||||
before = (await client.list_documents(filter=parent_uri_filter(parent_uri)))[0]
|
||||
|
||||
second = build_pdf([("a.txt", b"different")])
|
||||
await _reconcile_pdf_attachments(client, parent, second, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, second, depth=0)
|
||||
after = (await client.list_documents(filter=parent_uri_filter(parent_uri)))[0]
|
||||
|
||||
assert after.id == before.id
|
||||
|
|
@ -186,10 +187,10 @@ async def test_reingest_adds_new_attachment(temp_db_path, monkeypatch):
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
first = build_pdf([("a.txt", b"A")])
|
||||
parent = await _make_parent(client, parent_uri, first)
|
||||
await _reconcile_pdf_attachments(client, parent, first, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, first, depth=0)
|
||||
|
||||
second = build_pdf([("a.txt", b"A"), ("c.txt", b"C")])
|
||||
await _reconcile_pdf_attachments(client, parent, second, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, second, depth=0)
|
||||
children = await client.list_documents(filter=parent_uri_filter(parent_uri))
|
||||
assert len(children) == 2
|
||||
names = {c.uri for c in children}
|
||||
|
|
@ -213,7 +214,7 @@ async def test_nested_pdf_attachments_recurse_up_to_cap(temp_db_path, monkeypatc
|
|||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
root_uri = "file:///fixtures/root.pdf"
|
||||
parent = await _make_parent(client, root_uri, root)
|
||||
await _reconcile_pdf_attachments(client, parent, root, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, root, depth=0)
|
||||
|
||||
l1_uri = f"{root_uri}#attachment=l1.pdf"
|
||||
l2_uri = f"{l1_uri}#attachment=l2.pdf"
|
||||
|
|
@ -234,7 +235,7 @@ async def test_config_off_skips_extraction(temp_db_path, monkeypatch):
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
pdf_bytes = build_pdf([("a.txt", b"A")])
|
||||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
|
||||
assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
|
||||
|
||||
|
|
@ -256,7 +257,7 @@ async def test_non_pdf_parent_is_ignored(temp_db_path, monkeypatch):
|
|||
)
|
||||
# Even with PDF bytes, content_type=text/plain blocks extraction.
|
||||
pdf_bytes = build_pdf([("a.txt", b"A")])
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
|
||||
|
||||
|
||||
|
|
@ -272,7 +273,7 @@ async def test_parent_without_uri_is_skipped(temp_db_path, monkeypatch):
|
|||
metadata={"content_type": "application/pdf", "md5": "abc"},
|
||||
)
|
||||
pdf_bytes = build_pdf([("a.txt", b"A")])
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
assert await client.list_documents() == []
|
||||
|
||||
|
||||
|
|
@ -290,7 +291,7 @@ async def test_malformed_pdf_logs_warning_and_skips(temp_db_path, monkeypatch, c
|
|||
garbage = b"this is not a pdf at all"
|
||||
parent = await _make_parent(client, parent_uri, garbage)
|
||||
with caplog.at_level(logging.WARNING, logger="haiku.rag.client.documents"):
|
||||
await _reconcile_pdf_attachments(client, parent, garbage, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, garbage, depth=0)
|
||||
assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
|
||||
|
||||
|
||||
|
|
@ -328,7 +329,7 @@ async def test_unsupported_attachment_continues_loop(temp_db_path, monkeypatch):
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
pdf_bytes = build_pdf([("ok.txt", b"keep me"), ("unsupported.xyz", b"data")])
|
||||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
children = await client.list_documents(filter=parent_uri_filter(parent_uri))
|
||||
assert {c.uri for c in children} == {f"{parent_uri}#attachment=ok.txt"}
|
||||
|
||||
|
|
@ -346,7 +347,9 @@ async def test_joboptions_attachment_skipped_not_routed_as_pdf(temp_db_path, cap
|
|||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="haiku.rag.client.documents"):
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(
|
||||
writing(client), parent, pdf_bytes, depth=0
|
||||
)
|
||||
|
||||
assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
|
||||
|
||||
|
|
@ -360,7 +363,7 @@ async def test_cascade_delete_removes_reconciled_children(temp_db_path, monkeypa
|
|||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
pdf_bytes = build_pdf([("a.txt", b"A"), ("b.txt", b"B")])
|
||||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
assert len(await client.list_documents()) == 3
|
||||
|
||||
await client.delete_document(parent.id)
|
||||
|
|
@ -512,7 +515,7 @@ async def test_extract_pdf_attachments_called_off_event_loop_thread(
|
|||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
parent_uri = "file:///fixtures/parent.pdf"
|
||||
parent = await _make_parent(client, parent_uri, pdf_bytes)
|
||||
await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
|
||||
await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
|
||||
|
||||
assert called_from, "_extract_pdf_attachments was never called"
|
||||
assert called_from[0] is not event_loop_thread, (
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from haiku.rag.config import AppConfig, get_config
|
|||
from haiku.rag.store.models.chunk import Chunk, SearchResult
|
||||
from haiku.rag.store.models.document_item import DocumentItem
|
||||
from haiku.rag.tools.search import create_search_toolset
|
||||
from tests.conftest import writing
|
||||
from tests.test_context import _fetch_and_expand
|
||||
|
||||
|
||||
|
|
@ -244,7 +245,9 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
|
|||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
before = await rag.document_item_repository.get_all_picture_data(created.id)
|
||||
assert before.get("#/pictures/0") is not None
|
||||
|
|
@ -309,7 +312,7 @@ async def test_embed_only_preserves_picture_vectors(temp_db_path, monkeypatch):
|
|||
embedded = await embed_chunks(chunks, rag.embedder, rag._config)
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
await _store_document_with_chunks(rag, document, embedded, docling_doc)
|
||||
await _store_document_with_chunks(writing(rag), document, embedded, docling_doc)
|
||||
|
||||
before = await _picture_chunk_row(rag)
|
||||
assert list(before["vector"]) == pytest.approx(IMAGE_VEC)
|
||||
|
|
@ -819,7 +822,7 @@ async def test_ingest_emits_picture_chunks_with_multimodal_embedder(
|
|||
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
await _store_document_with_chunks(rag, document, embedded, docling_doc)
|
||||
await _store_document_with_chunks(writing(rag), document, embedded, docling_doc)
|
||||
|
||||
all_db_chunks = await rag.chunk_repository.store.chunks_table.query().to_list()
|
||||
picture_db_chunks = [
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import pytest
|
|||
|
||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||
from haiku.rag.config import get_config
|
||||
from tests.conftest import capture_logs
|
||||
from tests.conftest import capture_logs, writing
|
||||
|
||||
|
||||
class ChunkData(TypedDict):
|
||||
|
|
@ -603,12 +603,14 @@ async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch)
|
|||
)
|
||||
assert doc1.id is not None and doc2.id is not None
|
||||
|
||||
async def fake_generate_title(doc):
|
||||
async def fake_generate_title(config, doc):
|
||||
if doc.id == doc1.id:
|
||||
raise RuntimeError("simulated LLM failure")
|
||||
return "Second Title"
|
||||
|
||||
monkeypatch.setattr(client, "generate_title", fake_generate_title)
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.client.rebuild.generate_title", fake_generate_title
|
||||
)
|
||||
|
||||
processed_ids = [
|
||||
doc_id
|
||||
|
|
@ -743,7 +745,9 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
# _docling_doc_with_picture has no PageItems, so set_docling leaves
|
||||
|
|
@ -827,7 +831,9 @@ async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkey
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
called_with: list[dict[str, bytes]] = []
|
||||
|
|
@ -873,7 +879,7 @@ async def test_patch_picture_descriptions_returns_zero_for_doc_without_pictures(
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
doc = await rag.create_document(content="Just text, no pictures.")
|
||||
assert doc.id is not None
|
||||
n = await _patch_picture_descriptions(rag, doc)
|
||||
n = await _patch_picture_descriptions(writing(rag), doc)
|
||||
assert n == 0
|
||||
|
||||
|
||||
|
|
@ -897,7 +903,9 @@ async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path):
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
# Wipe the stored picture bytes to simulate a doc that knows about
|
||||
|
|
@ -910,7 +918,7 @@ async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path):
|
|||
from haiku.rag.client import rebuild as rebuild_module
|
||||
|
||||
with capture_logs(rebuild_module.logger, logging.WARNING) as records:
|
||||
n = await _patch_picture_descriptions(rag, created)
|
||||
n = await _patch_picture_descriptions(writing(rag), created)
|
||||
|
||||
assert n == 0
|
||||
assert any("no stored picture bytes" in r.getMessage() for r in records)
|
||||
|
|
@ -941,7 +949,9 @@ async def test_patch_picture_descriptions_skips_when_all_already_described(
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
called = False
|
||||
|
|
@ -956,7 +966,7 @@ async def test_patch_picture_descriptions_skips_when_all_already_described(
|
|||
fake_describe,
|
||||
)
|
||||
|
||||
n = await _patch_picture_descriptions(rag, created)
|
||||
n = await _patch_picture_descriptions(writing(rag), created)
|
||||
assert n == 0
|
||||
assert called is False
|
||||
|
||||
|
|
@ -979,7 +989,9 @@ async def test_rebuild_descriptions_raises_when_blob_is_missing(
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
# Force the stored doc to come back without a docling blob.
|
||||
|
|
@ -1251,9 +1263,11 @@ async def test_patch_picture_descriptions_returns_zero_without_descriptions(
|
|||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
|
||||
assert await _patch_picture_descriptions(rag, created) == 0
|
||||
assert await _patch_picture_descriptions(writing(rag), created) == 0
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
|
|
@ -1462,7 +1476,9 @@ async def test_rebuild_descriptions_flushes_in_batches(temp_db_path, monkeypatch
|
|||
docling_doc = _docling_doc_with_picture()
|
||||
document = Document(content=f"picture doc {i}", uri=f"test://doc-{i}")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(
|
||||
writing(rag), document, [], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
ids.append(created.id)
|
||||
|
||||
|
|
@ -1543,7 +1559,7 @@ async def test_rebuild_embed_only_recovers_picture_bytes(
|
|||
embedding=[0.1] * rag.embedder.vector_dim,
|
||||
)
|
||||
created = await _store_document_with_chunks(
|
||||
rag, document, [picture_chunk, text_chunk], docling_doc
|
||||
writing(rag), document, [picture_chunk, text_chunk], docling_doc
|
||||
)
|
||||
assert created.id is not None
|
||||
|
||||
|
|
|
|||
|
|
@ -353,17 +353,19 @@ class TestRebuildTitleOnly:
|
|||
await client.document_repository.update(doc2)
|
||||
|
||||
# Make generate_title fail for the first doc, succeed for the second
|
||||
original = HaikuRAG.generate_title
|
||||
import haiku.rag.client.rebuild as rebuild_mod
|
||||
|
||||
original = rebuild_mod.generate_title
|
||||
call_count = 0
|
||||
|
||||
async def flaky_generate(self, document):
|
||||
async def flaky_generate(config, document):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise RuntimeError("LLM failed")
|
||||
return await original(self, document)
|
||||
return await original(config, document)
|
||||
|
||||
monkeypatch.setattr(HaikuRAG, "generate_title", flaky_generate)
|
||||
monkeypatch.setattr(rebuild_mod, "generate_title", flaky_generate)
|
||||
|
||||
processed_ids = []
|
||||
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from haiku.rag.client import HaikuRAG
|
|||
from haiku.rag.client.documents import _refresh_doc_metadata
|
||||
from haiku.rag.config import get_config
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from tests.conftest import writing
|
||||
|
||||
|
||||
def _docling_doc(name: str, text: str):
|
||||
|
|
@ -85,7 +86,7 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path):
|
|||
client._session._vacuum_dirty = False
|
||||
|
||||
await _refresh_doc_metadata(
|
||||
client,
|
||||
writing(client),
|
||||
doc,
|
||||
title=None,
|
||||
user_metadata={},
|
||||
|
|
@ -111,7 +112,7 @@ async def test_metadata_refresh_waits_for_write_lock(temp_db_path):
|
|||
async with client.store._write_lock:
|
||||
task = asyncio.create_task(
|
||||
_refresh_doc_metadata(
|
||||
client,
|
||||
writing(client),
|
||||
doc,
|
||||
title=None,
|
||||
user_metadata={},
|
||||
|
|
|
|||
Loading…
Reference in a new issue