From 8c2a101a8a14efdbdcc17bdc4379930725955ad8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 6 Feb 2026 14:07:20 +0100 Subject: [PATCH] Fix out-of-memory in list_documents by adding column projection --- CHANGELOG.md | 4 ++ docs/python.md | 3 ++ haiku_rag_slim/haiku/rag/client.py | 7 ++- .../haiku/rag/store/repositories/document.py | 29 ++++++++++- tests/store/test_time_travel.py | 2 +- tests/test_database_autocreate.py | 4 +- tests/test_document.py | 48 +++++++++++++++++++ 7 files changed, 91 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cb04aa9..c8109d8c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [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 ### Added diff --git a/docs/python.md b/docs/python.md index 2555d13f..c1a148b4 100644 --- a/docs/python.md +++ b/docs/python.md @@ -134,6 +134,9 @@ doc = await client.get_document_by_uri("file:///path/to/document.pdf") List all documents: ```python 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: diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 42d432b9..74482041 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -867,6 +867,7 @@ class HaikuRAG: limit: int | None = None, offset: int | None = None, filter: str | None = None, + include_content: bool = False, ) -> list[Document]: """List all documents with optional pagination and filtering. @@ -874,12 +875,14 @@ 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. Returns: List of Document instances matching the criteria. """ 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: @@ -1480,7 +1483,7 @@ class HaikuRAG: settings_repo = SettingsRepository(self.store) settings_repo.save_current_settings() - documents = await self.list_documents() + documents = await self.list_documents(include_content=True) if mode == RebuildMode.EMBED_ONLY: async for doc_id in self._rebuild_embed_only(documents): diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index dfc01594..f40a6ed3 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -142,11 +142,14 @@ class DocumentRepository: self.store.documents_table.delete(f"id = '{safe_id}'") return True + _LISTING_COLUMNS = ["id", "title", "uri", "metadata", "created_at", "updated_at"] + async def list_all( self, limit: int | None = None, offset: int | None = None, filter: str | None = None, + include_content: bool = False, ) -> list[Document]: """List all documents with optional pagination and filtering. @@ -154,12 +157,16 @@ class DocumentRepository: 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 for listing. Returns: List of Document instances matching the criteria. """ query = self.store.documents_table.search() + if not include_content: + query = query.select(self._LISTING_COLUMNS) if filter is not None: query = query.where(filter) if offset is not None: @@ -167,8 +174,26 @@ class DocumentRepository: if limit is not None: query = query.limit(limit) - results = list(query.to_pydantic(DocumentRecord)) - return [self._record_to_document(doc) for doc in results] + if include_content: + 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: """Count documents with optional filtering. diff --git a/tests/store/test_time_travel.py b/tests/store/test_time_travel.py index bc7474f6..48d75d51 100644 --- a/tests/store/test_time_travel.py +++ b/tests/store/test_time_travel.py @@ -64,7 +64,7 @@ class TestStoreTimeTravel: repo = DocumentRepository(store) # Should only see first document - docs = await repo.list_all() + docs = await repo.list_all(include_content=True) assert len(docs) == 1 assert docs[0].content == "First document" store.close() diff --git a/tests/test_database_autocreate.py b/tests/test_database_autocreate.py index 5b269a16..5f28884e 100644 --- a/tests/test_database_autocreate.py +++ b/tests/test_database_autocreate.py @@ -46,4 +46,6 @@ async def test_operations_work_after_database_created(): async with HaikuRAG(db_path=db_path, config=config) as client: docs = await client.list_documents() 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" diff --git a/tests/test_document.py b/tests/test_document.py index e062ca41..3f43e12b 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -6,6 +6,54 @@ from haiku.rag.store.models.document import Document 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 async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path): """Test listing documents with filter clause."""