diff --git a/CHANGELOG.md b/CHANGELOG.md index 17dfdc98..965e7802 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Changed + +- `list_documents(include_content=True)` returns `content` only; the docling structure and page-image blobs are no longer loaded by a listing. + ## [0.72.0] - 2026-07-30 ### Added diff --git a/docs/python.md b/docs/python.md index 764f1717..abaecc1e 100644 --- a/docs/python.md +++ b/docs/python.md @@ -98,7 +98,8 @@ List all documents: ```python docs = await client.list_documents(limit=10, offset=0) -# Include full content and docling document (not loaded by default) +# Include the text content (not loaded by default). A listing never loads the +# docling blobs. docs = await client.list_documents(include_content=True) ``` diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index ecfd4b01..4dd365a9 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -439,8 +439,8 @@ class HaikuRAG: limit: Maximum number of documents to return. offset: Number of documents to skip. filter: Optional SQL WHERE clause to filter documents. - include_content: Whether to load content and docling_document. - Defaults to False to avoid loading large blobs. + include_content: Whether to load the text content. Defaults to + False. A listing never loads the docling blobs. Returns: List of Document instances matching the criteria. diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 76e0b7fe..8f332f0b 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -15,6 +15,10 @@ from haiku.rag.store.engine import ( from haiku.rag.store.models.document import Document from haiku.rag.utils import escape_sql_string +# Ids per `id IN (...)` content lookup. Keeps the filter string bounded on an +# unpaginated listing of a large database. +_CONTENT_BATCH = 512 + class DocumentRepository: """Repository for Document operations. @@ -320,13 +324,16 @@ class DocumentRepository: Listing reads `document_meta` (uri/title/metadata/timestamps); the SQL `filter` is evaluated against those columns. When `include_content` - is set, the content+blob row is loaded from `documents` and merged in. + is set, `content` is projected out of `documents` and merged in. The + docling blobs are left out: a single document's page rasters run to + hundreds of MB, so reading whole rows stalls a listing. Load them per + document with `get_docling_data` / `get_pages_data`. Args: limit: Maximum number of documents to return. offset: Number of documents to skip. filter: Optional SQL WHERE clause over document_meta columns. - include_content: Whether to also load content and docling blobs. + include_content: Whether to also load the text content. Returns: List of Document instances matching the criteria. @@ -342,26 +349,25 @@ class DocumentRepository: meta_records = await query_to_pydantic(query, DocumentMetaRecord) - if not include_content: - return [ - self._merge_to_document(DocumentRecord(id=m.id, content=""), m) - for m in meta_records - ] + content_by_id: dict[str, str] = {} + if include_content: + for start in range(0, len(meta_records), _CONTENT_BATCH): + batch = meta_records[start : start + _CONTENT_BATCH] + ids = ", ".join(f"'{escape_sql_string(m.id)}'" for m in batch) + rows = await ( + self.store.documents_table.query() + .select(["id", "content"]) + .where(f"id IN ({ids})") + .to_list() + ) + content_by_id.update((row["id"], row["content"]) for row in rows) - documents: list[Document] = [] - for meta in meta_records: - safe_id = escape_sql_string(meta.id) - doc_results = await query_to_pydantic( - self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1), - DocumentRecord, + return [ + self._merge_to_document( + DocumentRecord(id=m.id, content=content_by_id.get(m.id, "")), m ) - doc_record = ( - doc_results[0] - if doc_results - else DocumentRecord(id=meta.id, content="") - ) - documents.append(self._merge_to_document(doc_record, meta)) - return documents + for m in meta_records + ] async def count(self, filter: str | None = None) -> int: """Count documents with optional filtering (over document_meta columns).""" diff --git a/tests/test_document.py b/tests/test_document.py index a1c2ccb0..a45fd03d 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -3,6 +3,7 @@ import pytest from haiku.rag.store.engine import Store from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import DocumentItem +from haiku.rag.store.repositories import document as document_repository from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document_item import DocumentItemRepository @@ -35,6 +36,44 @@ async def test_document_list_all( assert docs[0].docling_document is None +@pytest.mark.asyncio +async def test_document_list_all_content_skips_docling_blobs(temp_db_path): + """include_content loads content only; the docling blobs (page rasters run + to hundreds of MB per document) are never materialized.""" + async with Store(temp_db_path, create=True) as store: + doc_repo = DocumentRepository(store) + + for content in ("first", "second"): + await doc_repo.create( + Document( + content=content, + docling_document=b"structure-blob", + docling_pages=b"page-raster-blob", + ) + ) + + docs = await doc_repo.list_all(include_content=True) + + assert {d.content for d in docs} == {"first", "second"} + assert all(d.docling_document is None for d in docs) + assert all(d.docling_pages is None for d in docs) + + +@pytest.mark.asyncio +async def test_document_list_all_content_spans_batches(temp_db_path, monkeypatch): + """Every document gets its content when the lookup spans several batches.""" + monkeypatch.setattr(document_repository, "_CONTENT_BATCH", 2) + async with Store(temp_db_path, create=True) as store: + doc_repo = DocumentRepository(store) + contents = {f"document {i}" for i in range(5)} + for content in contents: + await doc_repo.create(Document(content=content)) + + docs = await doc_repo.list_all(include_content=True) + + assert {d.content for d in docs} == contents + + @pytest.mark.asyncio async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path): """Test listing documents with filter clause."""