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

fix: ad more logging

Browse source
This commit is contained in:
bryan davis 2026-04-06 14:02:59 -05:00
parent d95d8291af
commit 7f31166829
2 changed files with 64 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1075,8 +1075,8 @@ class HaikuRAG:
# Step 1: Get reranker # Step 1: Get reranker
t0 = time.perf_counter() t0 = time.perf_counter()
reranker = get_reranker(config=self._config) reranker = get_reranker(config=self._config)
logger.debug( logger.info(
"search reranker_init took %.3fs", "search.reranker_init took %.3fs",
time.perf_counter() - t0, time.perf_counter() - t0,
) )
@ -1086,8 +1086,8 @@ class HaikuRAG:
chunk_results = await self.chunk_repository.search( chunk_results = await self.chunk_repository.search(
query, limit, search_type, filter query, limit, search_type, filter
) )
logger.debug( logger.info(
"search chunk_search type=%s limit=%d results=%d took %.3fs", "search.chunk_search type=%s limit=%d results=%d took %.3fs",
search_type, search_type,
limit, limit,
len(chunk_results), len(chunk_results),
@ -1098,8 +1098,8 @@ class HaikuRAG:
raw_results = await self.chunk_repository.search( raw_results = await self.chunk_repository.search(
query, search_limit, search_type, filter query, search_limit, search_type, filter
) )
logger.debug( logger.info(
"search chunk_search type=%s limit=%d results=%d took %.3fs", "search.chunk_search type=%s limit=%d results=%d took %.3fs",
search_type, search_type,
search_limit, search_limit,
len(raw_results), len(raw_results),
@ -1112,8 +1112,8 @@ class HaikuRAG:
chunk_results = await reranker.rerank( chunk_results = await reranker.rerank(
query, chunks, top_n=limit query, chunks, top_n=limit
) )
logger.debug( logger.info(
"search rerank candidates=%d top_n=%d took %.3fs", "search.rerank candidates=%d top_n=%d took %.3fs",
len(chunks), len(chunks),
limit, limit,
time.perf_counter() - t0, time.perf_counter() - t0,
@ -1125,8 +1125,8 @@ class HaikuRAG:
SearchResult.from_chunk(chunk, score) SearchResult.from_chunk(chunk, score)
for chunk, score in chunk_results for chunk, score in chunk_results
] ]
logger.debug( logger.info(
"search build_results count=%d took %.3fs", "search.build_results count=%d took %.3fs",
len(results), len(results),
time.perf_counter() - t0, time.perf_counter() - t0,
) )

54
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -234,10 +234,13 @@ class ChunkRepository:
Returns: Returns:
List of (chunk, score) tuples ordered by relevance. List of (chunk, score) tuples ordered by relevance.
""" """
import time
if not query.strip(): if not query.strip():
return [] return []
filtered_doc_ids = None filtered_doc_ids = None
if filter: if filter:
t0 = time.perf_counter()
# We perform filtering as a two-step process, first filtering documents, then # We perform filtering as a two-step process, first filtering documents, then
# filtering chunks based on those document IDs. # filtering chunks based on those document IDs.
# This is because LanceDB does not support joins directly in search queries. # This is because LanceDB does not support joins directly in search queries.
@ -252,10 +255,19 @@ class ChunkRepository:
return [] return []
# Keep as pandas Series for efficient vectorized operations # Keep as pandas Series for efficient vectorized operations
filtered_doc_ids = docs_df["id"] filtered_doc_ids = docs_df["id"]
logger.info(
"search.filter docs=%d took %.3fs",
len(filtered_doc_ids),
time.perf_counter() - t0,
)
# Prepare search query based on search type # Prepare search query based on search type
if search_type == "vector": if search_type == "vector":
t0 = time.perf_counter()
query_embedding = await self.embedder.embed_query(query) query_embedding = await self.embedder.embed_query(query)
logger.info(
"search.embed took %.3fs", time.perf_counter() - t0
)
vector_query = cast( vector_query = cast(
"LanceVectorQueryBuilder", "LanceVectorQueryBuilder",
self.store.chunks_table.search( self.store.chunks_table.search(
@ -270,7 +282,11 @@ class ChunkRepository:
results = self.store.chunks_table.search(query, query_type="fts") results = self.store.chunks_table.search(query, query_type="fts")
else: # hybrid (default) else: # hybrid (default)
t0 = time.perf_counter()
query_embedding = await self.embedder.embed_query(query) query_embedding = await self.embedder.embed_query(query)
logger.info(
"search.embed took %.3fs", time.perf_counter() - t0
)
# Create RRF reranker # Create RRF reranker
reranker = RRFReranker() reranker = RRFReranker()
# Perform native hybrid search with RRF reranking # Perform native hybrid search with RRF reranking
@ -286,10 +302,21 @@ class ChunkRepository:
# Apply filtering if needed (common for all search types) # Apply filtering if needed (common for all search types)
if filtered_doc_ids is not None: if filtered_doc_ids is not None:
t0 = time.perf_counter()
chunks_df = results.to_pandas() chunks_df = results.to_pandas()
logger.info(
"search.execute took %.3fs", time.perf_counter() - t0
)
t0 = time.perf_counter()
filtered_chunks_df = chunks_df.loc[ filtered_chunks_df = chunks_df.loc[
chunks_df["document_id"].isin(filtered_doc_ids) chunks_df["document_id"].isin(filtered_doc_ids)
].head(limit) ].head(limit)
logger.info(
"search.doc_filter rows=%d->%d took %.3fs",
len(chunks_df),
len(filtered_chunks_df),
time.perf_counter() - t0,
)
return await self._process_search_results(filtered_chunks_df) return await self._process_search_results(filtered_chunks_df)
# No filtering needed, apply limit and return # No filtering needed, apply limit and return
@ -453,6 +480,8 @@ class ChunkRepository:
Args: Args:
query_result: Either a pandas DataFrame or a LanceDB query result query_result: Either a pandas DataFrame or a LanceDB query result
""" """
import time
import pandas as pd import pandas as pd
def extract_scores(df: pd.DataFrame) -> list[float]: def extract_scores(df: pd.DataFrame) -> list[float]:
@ -470,17 +499,25 @@ class ChunkRepository:
raise ValueError("Unknown search result format, cannot extract scores") raise ValueError("Unknown search result format, cannot extract scores")
# Convert everything to DataFrame for uniform processing # Convert everything to DataFrame for uniform processing
t0 = time.perf_counter()
if isinstance(query_result, pd.DataFrame): if isinstance(query_result, pd.DataFrame):
df = query_result df = query_result
else: else:
# Convert LanceDB query result to DataFrame # Convert LanceDB query result to DataFrame
# (this is where the actual DB query executes)
df = query_result.to_pandas() df = query_result.to_pandas()
logger.info(
"search.execute rows=%d took %.3fs",
len(df),
time.perf_counter() - t0,
)
# Extract scores # Extract scores
scores = extract_scores(df) scores = extract_scores(df)
# Convert DataFrame rows to ChunkRecords using to_dict # Convert DataFrame rows to ChunkRecords using to_dict
# (avoids slow .iterrows() overhead) # (avoids slow .iterrows() overhead)
t0 = time.perf_counter()
rows = df.to_dict(orient="records") rows = df.to_dict(orient="records")
pydantic_results = [ pydantic_results = [
self.store.ChunkRecord( self.store.ChunkRecord(
@ -493,11 +530,17 @@ class ChunkRepository:
) )
for row in rows for row in rows
] ]
logger.info(
"search.to_records count=%d took %.3fs",
len(pydantic_results),
time.perf_counter() - t0,
)
# Collect all unique document IDs for batch lookup # Collect all unique document IDs for batch lookup
document_ids = list(set(chunk.document_id for chunk in pydantic_results)) document_ids = list(set(chunk.document_id for chunk in pydantic_results))
# Batch fetch all documents at once # Batch fetch all documents at once
t0 = time.perf_counter()
documents_map = {} documents_map = {}
if document_ids: if document_ids:
# Use IN clause for efficient batch lookup # Use IN clause for efficient batch lookup
@ -509,8 +552,14 @@ class ChunkRepository:
.to_pydantic(DocumentRecord) .to_pydantic(DocumentRecord)
) )
documents_map = {doc.id: doc for doc in doc_results} documents_map = {doc.id: doc for doc in doc_results}
logger.info(
"search.doc_lookup docs=%d took %.3fs",
len(documents_map),
time.perf_counter() - t0,
)
# Build final results with document info # Build final results with document info
t0 = time.perf_counter()
chunks_with_scores = [] chunks_with_scores = []
for i, chunk_record in enumerate(pydantic_results): for i, chunk_record in enumerate(pydantic_results):
doc = documents_map.get(chunk_record.document_id) doc = documents_map.get(chunk_record.document_id)
@ -526,5 +575,10 @@ class ChunkRepository:
) )
score = scores[i] if i < len(scores) else 1.0 score = scores[i] if i < len(scores) else 1.0
chunks_with_scores.append((chunk, score)) chunks_with_scores.append((chunk, score))
logger.info(
"search.build_chunks count=%d took %.3fs",
len(chunks_with_scores),
time.perf_counter() - t0,
)
return chunks_with_scores return chunks_with_scores