danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add optional gzip

Browse source 
feat: add optional cache size
test: force expand context false
This commit is contained in:
bryan davis 2026-04-06 14:51:29 -05:00
parent 804f76f301
commit afd91db0e9
6 changed files with 73 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -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

29
haiku_rag_slim/haiku/rag/client.py
View file

@ -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

2
haiku_rag_slim/haiku/rag/config/models.py
View file

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

27
haiku_rag_slim/haiku/rag/store/compression.py
View file

@ -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")

5
haiku_rag_slim/haiku/rag/store/engine.py
View file

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

23
haiku_rag_slim/haiku/rag/store/models/document.py
View file

@ -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