Fix out-of-memory in list_documents by adding column projection

This commit is contained in:
Yiorgis Gozadinos 2026-02-06 14:07:20 +01:00
parent 17c2136b4e
commit 8c2a101a8a
No known key found for this signature in database
7 changed files with 91 additions and 6 deletions

View file

@ -1,6 +1,10 @@
# Changelog # Changelog
## [Unreleased] ## [Unreleased]
### Fixed
- **Document listing memory usage**: `list_documents` no longer loads full document content and docling blobs by default, preventing out-of-memory errors on large databases. Use `include_content=True` when content is needed.
## [0.29.0] - 2026-02-06 ## [0.29.0] - 2026-02-06
### Added ### Added

View file

@ -134,6 +134,9 @@ doc = await client.get_document_by_uri("file:///path/to/document.pdf")
List all documents: 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)
docs = await client.list_documents(include_content=True)
``` ```
Filter documents by properties: Filter documents by properties:

View file

@ -867,6 +867,7 @@ class HaikuRAG:
limit: int | None = None, limit: int | None = None,
offset: int | None = None, offset: int | None = None,
filter: str | None = None, filter: str | None = None,
include_content: bool = False,
) -> list[Document]: ) -> list[Document]:
"""List all documents with optional pagination and filtering. """List all documents with optional pagination and filtering.
@ -874,12 +875,14 @@ 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.
Defaults to False to avoid loading large blobs.
Returns: Returns:
List of Document instances matching the criteria. List of Document instances matching the criteria.
""" """
return await self.document_repository.list_all( return await self.document_repository.list_all(
limit=limit, offset=offset, filter=filter limit=limit, offset=offset, filter=filter, include_content=include_content
) )
async def count_documents(self, filter: str | None = None) -> int: async def count_documents(self, filter: str | None = None) -> int:
@ -1480,7 +1483,7 @@ class HaikuRAG:
settings_repo = SettingsRepository(self.store) settings_repo = SettingsRepository(self.store)
settings_repo.save_current_settings() settings_repo.save_current_settings()
documents = await self.list_documents() documents = await self.list_documents(include_content=True)
if mode == RebuildMode.EMBED_ONLY: if mode == RebuildMode.EMBED_ONLY:
async for doc_id in self._rebuild_embed_only(documents): async for doc_id in self._rebuild_embed_only(documents):

View file

@ -142,11 +142,14 @@ class DocumentRepository:
self.store.documents_table.delete(f"id = '{safe_id}'") self.store.documents_table.delete(f"id = '{safe_id}'")
return True return True
_LISTING_COLUMNS = ["id", "title", "uri", "metadata", "created_at", "updated_at"]
async def list_all( async def list_all(
self, self,
limit: int | None = None, limit: int | None = None,
offset: int | None = None, offset: int | None = None,
filter: str | None = None, filter: str | None = None,
include_content: bool = False,
) -> list[Document]: ) -> list[Document]:
"""List all documents with optional pagination and filtering. """List all documents with optional pagination and filtering.
@ -154,12 +157,16 @@ class DocumentRepository:
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.
Defaults to False to avoid loading large blobs for listing.
Returns: Returns:
List of Document instances matching the criteria. List of Document instances matching the criteria.
""" """
query = self.store.documents_table.search() query = self.store.documents_table.search()
if not include_content:
query = query.select(self._LISTING_COLUMNS)
if filter is not None: if filter is not None:
query = query.where(filter) query = query.where(filter)
if offset is not None: if offset is not None:
@ -167,8 +174,26 @@ class DocumentRepository:
if limit is not None: if limit is not None:
query = query.limit(limit) query = query.limit(limit)
results = list(query.to_pydantic(DocumentRecord)) if include_content:
return [self._record_to_document(doc) for doc in results] results = list(query.to_pydantic(DocumentRecord))
return [self._record_to_document(doc) for doc in results]
return [
Document(
id=row["id"],
content="",
title=row.get("title"),
uri=row.get("uri"),
metadata=json.loads(row.get("metadata", "{}")),
created_at=datetime.fromisoformat(row["created_at"])
if row.get("created_at")
else datetime.now(),
updated_at=datetime.fromisoformat(row["updated_at"])
if row.get("updated_at")
else datetime.now(),
)
for row in query.to_list()
]
async def count(self, filter: str | None = None) -> int: async def count(self, filter: str | None = None) -> int:
"""Count documents with optional filtering. """Count documents with optional filtering.

View file

@ -64,7 +64,7 @@ class TestStoreTimeTravel:
repo = DocumentRepository(store) repo = DocumentRepository(store)
# Should only see first document # Should only see first document
docs = await repo.list_all() docs = await repo.list_all(include_content=True)
assert len(docs) == 1 assert len(docs) == 1
assert docs[0].content == "First document" assert docs[0].content == "First document"
store.close() store.close()

View file

@ -46,4 +46,6 @@ async def test_operations_work_after_database_created():
async with HaikuRAG(db_path=db_path, config=config) as client: async with HaikuRAG(db_path=db_path, config=config) as client:
docs = await client.list_documents() docs = await client.list_documents()
assert len(docs) == 1 assert len(docs) == 1
assert docs[0].content == "Test content" doc = await client.get_document_by_id(docs[0].id)
assert doc is not None
assert doc.content == "Test content"

View file

@ -6,6 +6,54 @@ from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_document_list_excludes_content_by_default(
qa_corpus: Dataset, temp_db_path
):
"""list_all excludes content and docling_document by default."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
title="Test Document",
metadata={"key": "value"},
)
created = await doc_repo.create(doc)
docs = await doc_repo.list_all()
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].title == "Test Document"
assert docs[0].uri == "https://example.com/doc.txt"
assert docs[0].metadata == {"key": "value"}
assert docs[0].content == ""
assert docs[0].docling_document is None
store.close()
@pytest.mark.asyncio
async def test_document_list_includes_content_when_requested(
qa_corpus: Dataset, temp_db_path
):
"""list_all returns content when include_content=True."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
content = qa_corpus[0]["document_extracted"]
doc = Document(content=content, uri="https://example.com/doc.txt")
created = await doc_repo.create(doc)
docs = await doc_repo.list_all(include_content=True)
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].content == content
store.close()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path): async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path):
"""Test listing documents with filter clause.""" """Test listing documents with filter clause."""