Load only content when listing documents with include_content
This commit is contained in:
parent
adae9eff2e
commit
8b45f81464
5 changed files with 73 additions and 23 deletions
|
|
@ -1,6 +1,10 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
## [0.72.0] - 2026-07-30
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|
|
||||||
|
|
@ -98,7 +98,8 @@ List all documents:
|
||||||
```python
|
```python
|
||||||
docs = await client.list_documents(limit=10, offset=0)
|
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)
|
docs = await client.list_documents(include_content=True)
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -439,8 +439,8 @@ class HaikuRAG:
|
||||||
limit: Maximum number of documents to return.
|
limit: Maximum number of documents to return.
|
||||||
offset: Number of documents to skip.
|
offset: Number of documents to skip.
|
||||||
filter: Optional SQL WHERE clause to filter documents.
|
filter: Optional SQL WHERE clause to filter documents.
|
||||||
include_content: Whether to load content and docling_document.
|
include_content: Whether to load the text content. Defaults to
|
||||||
Defaults to False to avoid loading large blobs.
|
False. A listing never loads the docling blobs.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List of Document instances matching the criteria.
|
List of Document instances matching the criteria.
|
||||||
|
|
|
||||||
|
|
@ -15,6 +15,10 @@ from haiku.rag.store.engine import (
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
from haiku.rag.utils import escape_sql_string
|
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:
|
class DocumentRepository:
|
||||||
"""Repository for Document operations.
|
"""Repository for Document operations.
|
||||||
|
|
@ -320,13 +324,16 @@ class DocumentRepository:
|
||||||
|
|
||||||
Listing reads `document_meta` (uri/title/metadata/timestamps); the
|
Listing reads `document_meta` (uri/title/metadata/timestamps); the
|
||||||
SQL `filter` is evaluated against those columns. When `include_content`
|
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:
|
Args:
|
||||||
limit: Maximum number of documents to return.
|
limit: Maximum number of documents to return.
|
||||||
offset: Number of documents to skip.
|
offset: Number of documents to skip.
|
||||||
filter: Optional SQL WHERE clause over document_meta columns.
|
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:
|
Returns:
|
||||||
List of Document instances matching the criteria.
|
List of Document instances matching the criteria.
|
||||||
|
|
@ -342,26 +349,25 @@ class DocumentRepository:
|
||||||
|
|
||||||
meta_records = await query_to_pydantic(query, DocumentMetaRecord)
|
meta_records = await query_to_pydantic(query, DocumentMetaRecord)
|
||||||
|
|
||||||
if not include_content:
|
content_by_id: dict[str, str] = {}
|
||||||
return [
|
if include_content:
|
||||||
self._merge_to_document(DocumentRecord(id=m.id, content=""), m)
|
for start in range(0, len(meta_records), _CONTENT_BATCH):
|
||||||
for m in meta_records
|
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] = []
|
return [
|
||||||
for meta in meta_records:
|
self._merge_to_document(
|
||||||
safe_id = escape_sql_string(meta.id)
|
DocumentRecord(id=m.id, content=content_by_id.get(m.id, "")), m
|
||||||
doc_results = await query_to_pydantic(
|
|
||||||
self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
|
|
||||||
DocumentRecord,
|
|
||||||
)
|
)
|
||||||
doc_record = (
|
for m in meta_records
|
||||||
doc_results[0]
|
]
|
||||||
if doc_results
|
|
||||||
else DocumentRecord(id=meta.id, content="")
|
|
||||||
)
|
|
||||||
documents.append(self._merge_to_document(doc_record, meta))
|
|
||||||
return documents
|
|
||||||
|
|
||||||
async def count(self, filter: str | None = None) -> int:
|
async def count(self, filter: str | None = None) -> int:
|
||||||
"""Count documents with optional filtering (over document_meta columns)."""
|
"""Count documents with optional filtering (over document_meta columns)."""
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ import pytest
|
||||||
from haiku.rag.store.engine import Store
|
from haiku.rag.store.engine import Store
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
from haiku.rag.store.models.document_item import DocumentItem
|
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 import DocumentRepository
|
||||||
from haiku.rag.store.repositories.document_item import DocumentItemRepository
|
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
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path):
|
async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path):
|
||||||
"""Test listing documents with filter clause."""
|
"""Test listing documents with filter clause."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue