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

add more logging

Browse source
This commit is contained in:
bryan davis 2026-04-06 14:26:54 -05:00
parent 712c073e2b
commit 804f76f301
2 changed files with 31 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1109,9 +1109,7 @@ class HaikuRAG:
# Step 3: Reranking # Step 3: Reranking
t0 = time.perf_counter() t0 = time.perf_counter()
chunks = [chunk for chunk, _ in raw_results] chunks = [chunk for chunk, _ in raw_results]
chunk_results = await reranker.rerank( chunk_results = await reranker.rerank(query, chunks, top_n=limit)
query, chunks, top_n=limit
)
logger.info( logger.info(
"search.rerank candidates=%d top_n=%d took %.3fs", "search.rerank candidates=%d top_n=%d took %.3fs",
len(chunks), len(chunks),
@ -1122,8 +1120,7 @@ class HaikuRAG:
# Step 4: Build SearchResult objects # Step 4: Build SearchResult objects
t0 = time.perf_counter() t0 = time.perf_counter()
results = [ results = [
SearchResult.from_chunk(chunk, score) SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results
for chunk, score in chunk_results
] ]
logger.info( logger.info(
"search.build_results count=%d took %.3fs", "search.build_results count=%d took %.3fs",
@ -1209,9 +1206,7 @@ class HaikuRAG:
if has_refs: if has_refs:
# Only fetch docling data (skip content blob) # Only fetch docling data (skip content blob)
t0 = time.perf_counter() t0 = time.perf_counter()
doc = await self.document_repository.get_docling_data( doc = await self.document_repository.get_docling_data(doc_id)
doc_id
)
logger.info( logger.info(
"expand.fetch_docling_data doc=%s took %.3fs", "expand.fetch_docling_data doc=%s took %.3fs",
doc_id[:8], doc_id[:8],
@ -1221,7 +1216,7 @@ class HaikuRAG:
t0 = time.perf_counter() t0 = time.perf_counter()
docling_doc = doc.get_docling_document() docling_doc = doc.get_docling_document()
logger.info( logger.info(
"expand.decompress_docling doc=%s took %.3fs", "expand.get_docling_document doc=%s took %.3fs",
doc_id[:8], doc_id[:8],
time.perf_counter() - t0, time.perf_counter() - t0,
) )
@ -1237,7 +1232,7 @@ class HaikuRAG:
max_chars, max_chars,
) )
logger.info( logger.info(
"expand.docling doc=%s results=%d->%d took %.3fs", "expand._expand_with_docling doc=%s results=%d->%d took %.3fs",
doc_id[:8], doc_id[:8],
len(doc_results), len(doc_results),
len(expanded), len(expanded),
@ -1252,7 +1247,7 @@ class HaikuRAG:
doc_id, doc_results, radius doc_id, doc_results, radius
) )
logger.info( logger.info(
"expand.chunks doc=%s results=%d->%d took %.3fs", "expand._expand_with_chunks doc=%s results=%d->%d took %.3fs",
doc_id[:8], doc_id[:8],
len(doc_results), len(doc_results),
len(expanded), len(expanded),
@ -1531,9 +1526,7 @@ class HaikuRAG:
if not chunk_ids: if not chunk_ids:
return results return results
id_to_order = await self.chunk_repository.get_orders_by_ids( id_to_order = await self.chunk_repository.get_orders_by_ids(chunk_ids)
chunk_ids
)
if not id_to_order: if not id_to_order:
return results return results
@ -1543,11 +1536,9 @@ class HaikuRAG:
range_max = max(orders) + radius range_max = max(orders) + radius
# Fetch only chunks in the needed range # Fetch only chunks in the needed range
chunks_in_doc = ( chunks_in_doc = await self.chunk_repository.get_by_document_id_order_range(
await self.chunk_repository.get_by_document_id_order_range(
doc_id, range_min, range_max doc_id, range_min, range_max
) )
)
if not chunks_in_doc: if not chunks_in_doc:
return results return results

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

@ -1,3 +1,5 @@
import logging
import time
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@ -9,6 +11,7 @@ from haiku.rag.store.compression import decompress_json
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
logger = logging.getLogger(__name__)
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100) _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
@ -18,12 +21,31 @@ def _get_cached_docling_document(
) -> "DoclingDocument": ) -> "DoclingDocument":
"""Get or parse DoclingDocument with LRU caching by document ID.""" """Get or parse DoclingDocument with LRU caching by document ID."""
if document_id in _docling_document_cache: if document_id in _docling_document_cache:
logger.info("docling.cache_hit doc=%s", document_id[:8])
return _docling_document_cache[document_id] return _docling_document_cache[document_id]
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
t0 = time.perf_counter()
json_str = decompress_json(compressed_data) json_str = decompress_json(compressed_data)
decompress_time = time.perf_counter() - t0
logger.info(
"docling.decompress doc=%s bytes=%d json_chars=%d took %.3fs",
document_id[:8],
len(compressed_data),
len(json_str),
decompress_time,
)
t0 = time.perf_counter()
doc = DoclingDocument.model_validate_json(json_str) doc = DoclingDocument.model_validate_json(json_str)
validate_time = time.perf_counter() - t0
logger.info(
"docling.model_validate doc=%s took %.3fs",
document_id[:8],
validate_time,
)
_docling_document_cache[document_id] = doc _docling_document_cache[document_id] = doc
return doc return doc