feat: add optional gzip

feat: add optional cache size
test: force expand context false

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(
+        logger.info("qa.agent_run took %.3fs", agent_duration)
-            "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(
+        logger.info("qa.answer completed total=%.3fs", agent_duration)
-            "qa.answer completed total=%.3fs", agent_duration
-        )
         return output.answer, citations

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

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):

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."""
+def compress_json(json_str: str, *, enabled: bool = True) -> bytes:
-    return gzip.compress(json_str.encode("utf-8"))
+    """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."""
+    """Decompress data to a JSON string.
-    return gzip.decompress(data).decode("utf-8")
+
+    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")

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)
 

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