From afd91db0e971b86d7a7d8eefd096ba2c48c2d4a4 Mon Sep 17 00:00:00 2001 From: bryan davis Date: Mon, 6 Apr 2026 14:51:29 -0500 Subject: [PATCH] feat: add optional gzip feat: add optional cache size test: force expand context false --- haiku_rag_slim/haiku/rag/agents/qa/agent.py | 9 ++---- haiku_rag_slim/haiku/rag/client.py | 29 ++++++++++++------- haiku_rag_slim/haiku/rag/config/models.py | 2 ++ haiku_rag_slim/haiku/rag/store/compression.py | 27 +++++++++++++---- haiku_rag_slim/haiku/rag/store/engine.py | 5 ++++ .../haiku/rag/store/models/document.py | 23 +++++++++++++++ 6 files changed, 73 insertions(+), 22 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/qa/agent.py b/haiku_rag_slim/haiku/rag/agents/qa/agent.py index 10078d72..d567bce1 100644 --- a/haiku_rag_slim/haiku/rag/agents/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/qa/agent.py @@ -58,6 +58,7 @@ class QuestionAnswerAgent: tool_name="search", on_results=accumulated_results.extend, max_searches=max_searches, + expand_context=False, ) # Agent created per-call: toolset varies with filter, and Agent @@ -81,9 +82,7 @@ class QuestionAnswerAgent: t0 = time.perf_counter() result = await agent.run(question, deps=deps) agent_duration = time.perf_counter() - t0 - logger.info( - "qa.agent_run took %.3fs", agent_duration - ) + logger.info("qa.agent_run took %.3fs", agent_duration) t0 = time.perf_counter() output = result.output @@ -93,7 +92,5 @@ class QuestionAnswerAgent: len(citations), time.perf_counter() - t0, ) - logger.info( - "qa.answer completed total=%.3fs", agent_duration - ) + logger.info("qa.answer completed total=%.3fs", agent_duration) return output.answer, citations diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 3f441e51..d2fd89f0 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -103,6 +103,13 @@ class HaikuRAG: """Whether the client is in read-only mode.""" return self.store.is_read_only + def _compress_docling(self, json_str: str) -> bytes: + """Compress a docling JSON string, respecting config.""" + return compress_json( + json_str, + enabled=self._config.storage.compress_docling, + ) + async def __aenter__(self): """Async context manager entry.""" return self @@ -495,7 +502,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata or {}, - docling_document=compress_json(docling_document.model_dump_json()), + docling_document=self._compress_docling(docling_document.model_dump_json()), docling_version=docling_document.version, ) @@ -535,7 +542,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata or {}, - docling_document=compress_json(docling_document.model_dump_json()), + docling_document=self._compress_docling(docling_document.model_dump_json()), docling_version=docling_document.version, ) @@ -672,7 +679,7 @@ class HaikuRAG: # Update existing document and rechunk existing_doc.content = stored_content existing_doc.metadata = metadata - existing_doc.docling_document = compress_json( + existing_doc.docling_document = self._compress_docling( docling_document.model_dump_json() ) existing_doc.docling_version = docling_document.version @@ -694,7 +701,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata, - docling_document=compress_json(docling_document.model_dump_json()), + docling_document=self._compress_docling(docling_document.model_dump_json()), docling_version=docling_document.version, ) return await self._store_document_with_chunks(document, embedded_chunks) @@ -789,7 +796,7 @@ class HaikuRAG: # Update existing document and rechunk existing_doc.content = stored_content existing_doc.metadata = metadata - existing_doc.docling_document = compress_json( + existing_doc.docling_document = self._compress_docling( docling_document.model_dump_json() ) existing_doc.docling_version = docling_document.version @@ -811,7 +818,7 @@ class HaikuRAG: uri=url, title=title, metadata=metadata, - docling_document=compress_json(docling_document.model_dump_json()), + docling_document=self._compress_docling(docling_document.model_dump_json()), docling_version=docling_document.version, ) return await self._store_document_with_chunks(document, embedded_chunks) @@ -971,7 +978,7 @@ class HaikuRAG: # Store docling data if provided if docling_document is not None: existing_doc.content = docling_document.export_to_markdown() - existing_doc.docling_document = compress_json( + existing_doc.docling_document = self._compress_docling( docling_document.model_dump_json() ) existing_doc.docling_version = docling_document.version @@ -983,7 +990,7 @@ class HaikuRAG: # DoclingDocument provided without chunks - chunk and embed using primitives if docling_document is not None: existing_doc.content = docling_document.export_to_markdown() - existing_doc.docling_document = compress_json( + existing_doc.docling_document = self._compress_docling( docling_document.model_dump_json() ) existing_doc.docling_version = docling_document.version @@ -998,7 +1005,7 @@ class HaikuRAG: assert content is not None existing_doc.content = content converted_docling = await self.convert(existing_doc.content) - existing_doc.docling_document = compress_json( + existing_doc.docling_document = self._compress_docling( converted_docling.model_dump_json() ) existing_doc.docling_version = converted_docling.version @@ -1966,7 +1973,7 @@ class HaikuRAG: embedded_chunks = await embed_chunks(chunks, self._config) # Update document fields - doc.docling_document = compress_json(docling_document.model_dump_json()) + doc.docling_document = self._compress_docling(docling_document.model_dump_json()) doc.docling_version = docling_document.version # Prepare chunks with document_id and order @@ -2045,7 +2052,7 @@ class HaikuRAG: chunks = await self.chunk(docling_document) embedded_chunks = await embed_chunks(chunks, self._config) - doc.docling_document = compress_json(docling_document.model_dump_json()) + doc.docling_document = self._compress_docling(docling_document.model_dump_json()) doc.docling_version = docling_document.version # Prepare chunks with document_id and order diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index de1d6cad..f9c06208 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -47,6 +47,8 @@ class StorageConfig(BaseModel): data_dir: Path = Field(default_factory=get_default_data_dir) auto_vacuum: bool = True vacuum_retention_seconds: int = 86400 + docling_cache_size: int = 100 + compress_docling: bool = True class MonitorConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/store/compression.py b/haiku_rag_slim/haiku/rag/store/compression.py index 881efb89..d0bfc9c7 100644 --- a/haiku_rag_slim/haiku/rag/store/compression.py +++ b/haiku_rag_slim/haiku/rag/store/compression.py @@ -1,11 +1,28 @@ import gzip +# Gzip magic number: first two bytes of any gzip stream +_GZIP_MAGIC = b"\x1f\x8b" -def compress_json(json_str: str) -> bytes: - """Compress a JSON string with gzip.""" - return gzip.compress(json_str.encode("utf-8")) + +def compress_json(json_str: str, *, enabled: bool = True) -> bytes: + """Compress a JSON string, optionally with gzip. + + Args: + json_str: The JSON string to compress. + enabled: If False, returns raw UTF-8 bytes without compression. + """ + data = json_str.encode("utf-8") + if enabled: + return gzip.compress(data) + return data def decompress_json(data: bytes) -> str: - """Decompress gzip-compressed data to a JSON string.""" - return gzip.decompress(data).decode("utf-8") + """Decompress data to a JSON string. + + Automatically detects gzip-compressed data via magic bytes, + so it handles both compressed and uncompressed storage. + """ + if data[:2] == _GZIP_MAGIC: + return gzip.decompress(data).decode("utf-8") + return data.decode("utf-8") diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index aeb6748a..07570ad2 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -107,6 +107,11 @@ class Store: if not db_path.parent.exists(): Path.mkdir(db_path.parent, parents=True) + # Configure docling document cache size from config + from haiku.rag.store.models.document import configure_docling_cache + + configure_docling_cache(config.storage.docling_cache_size) + # Connect to LanceDB self.db = self._connect_to_lancedb(db_path) diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 0956be82..0e702cb1 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -16,6 +16,22 @@ logger = logging.getLogger(__name__) _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100) +def configure_docling_cache(maxsize: int) -> None: + """Resize the DoclingDocument LRU cache. + + Existing entries are preserved up to the new maxsize. + """ + global _docling_document_cache + if _docling_document_cache.maxsize == maxsize: + return + old = _docling_document_cache + _docling_document_cache = LRUCache(maxsize=maxsize) + # Copy existing entries (LRU order preserved by iteration) + for key in old: + _docling_document_cache[key] = old[key] + logger.info("docling.cache_resized maxsize=%d", maxsize) + + def _get_cached_docling_document( document_id: str, compressed_data: bytes ) -> "DoclingDocument": @@ -26,6 +42,13 @@ def _get_cached_docling_document( from docling_core.types.doc.document import DoclingDocument + logger.info( + "docling.cache_miss doc=%s cache_size=%d/%d", + document_id[:8], + len(_docling_document_cache), + _docling_document_cache.maxsize, + ) + t0 = time.perf_counter() json_str = decompress_json(compressed_data) decompress_time = time.perf_counter() - t0