fix: ad more logging
This commit is contained in:
bryan davis
parent
d95d8291af
commit
7f31166829
2 changed files with 64 additions and 10 deletions
|
|
@ -1075,8 +1075,8 @@ class HaikuRAG:
|
|||
# Step 1: Get reranker
|
||||
t0 = time.perf_counter()
|
||||
reranker = get_reranker(config=self._config)
|
||||
logger.debug(
|
||||
"search reranker_init took %.3fs",
|
||||
logger.info(
|
||||
"search.reranker_init took %.3fs",
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
|
|
@ -1086,8 +1086,8 @@ class HaikuRAG:
|
|||
chunk_results = await self.chunk_repository.search(
|
||||
query, limit, search_type, filter
|
||||
)
|
||||
logger.debug(
|
||||
"search chunk_search type=%s limit=%d results=%d took %.3fs",
|
||||
logger.info(
|
||||
"search.chunk_search type=%s limit=%d results=%d took %.3fs",
|
||||
search_type,
|
||||
limit,
|
||||
len(chunk_results),
|
||||
|
|
@ -1098,8 +1098,8 @@ class HaikuRAG:
|
|||
raw_results = await self.chunk_repository.search(
|
||||
query, search_limit, search_type, filter
|
||||
)
|
||||
logger.debug(
|
||||
"search chunk_search type=%s limit=%d results=%d took %.3fs",
|
||||
logger.info(
|
||||
"search.chunk_search type=%s limit=%d results=%d took %.3fs",
|
||||
search_type,
|
||||
search_limit,
|
||||
len(raw_results),
|
||||
|
|
@ -1112,8 +1112,8 @@ class HaikuRAG:
|
|||
chunk_results = await reranker.rerank(
|
||||
query, chunks, top_n=limit
|
||||
)
|
||||
logger.debug(
|
||||
"search rerank candidates=%d top_n=%d took %.3fs",
|
||||
logger.info(
|
||||
"search.rerank candidates=%d top_n=%d took %.3fs",
|
||||
len(chunks),
|
||||
limit,
|
||||
time.perf_counter() - t0,
|
||||
|
|
@ -1125,8 +1125,8 @@ class HaikuRAG:
|
|||
SearchResult.from_chunk(chunk, score)
|
||||
for chunk, score in chunk_results
|
||||
]
|
||||
logger.debug(
|
||||
"search build_results count=%d took %.3fs",
|
||||
logger.info(
|
||||
"search.build_results count=%d took %.3fs",
|
||||
len(results),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -234,10 +234,13 @@ class ChunkRepository:
|
|||
Returns:
|
||||
List of (chunk, score) tuples ordered by relevance.
|
||||
"""
|
||||
import time
|
||||
|
||||
if not query.strip():
|
||||
return []
|
||||
filtered_doc_ids = None
|
||||
if filter:
|
||||
t0 = time.perf_counter()
|
||||
# We perform filtering as a two-step process, first filtering documents, then
|
||||
# filtering chunks based on those document IDs.
|
||||
# This is because LanceDB does not support joins directly in search queries.
|
||||
|
|
@ -252,10 +255,19 @@ class ChunkRepository:
|
|||
return []
|
||||
# Keep as pandas Series for efficient vectorized operations
|
||||
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
|
||||
if search_type == "vector":
|
||||
t0 = time.perf_counter()
|
||||
query_embedding = await self.embedder.embed_query(query)
|
||||
logger.info(
|
||||
"search.embed took %.3fs", time.perf_counter() - t0
|
||||
)
|
||||
vector_query = cast(
|
||||
"LanceVectorQueryBuilder",
|
||||
self.store.chunks_table.search(
|
||||
|
|
@ -270,7 +282,11 @@ class ChunkRepository:
|
|||
results = self.store.chunks_table.search(query, query_type="fts")
|
||||
|
||||
else: # hybrid (default)
|
||||
t0 = time.perf_counter()
|
||||
query_embedding = await self.embedder.embed_query(query)
|
||||
logger.info(
|
||||
"search.embed took %.3fs", time.perf_counter() - t0
|
||||
)
|
||||
# Create RRF reranker
|
||||
reranker = RRFReranker()
|
||||
# Perform native hybrid search with RRF reranking
|
||||
|
|
@ -286,10 +302,21 @@ class ChunkRepository:
|
|||
|
||||
# Apply filtering if needed (common for all search types)
|
||||
if filtered_doc_ids is not None:
|
||||
t0 = time.perf_counter()
|
||||
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[
|
||||
chunks_df["document_id"].isin(filtered_doc_ids)
|
||||
].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)
|
||||
|
||||
# No filtering needed, apply limit and return
|
||||
|
|
@ -453,6 +480,8 @@ class ChunkRepository:
|
|||
Args:
|
||||
query_result: Either a pandas DataFrame or a LanceDB query result
|
||||
"""
|
||||
import time
|
||||
|
||||
import pandas as pd
|
||||
|
||||
def extract_scores(df: pd.DataFrame) -> list[float]:
|
||||
|
|
@ -470,17 +499,25 @@ class ChunkRepository:
|
|||
raise ValueError("Unknown search result format, cannot extract scores")
|
||||
|
||||
# Convert everything to DataFrame for uniform processing
|
||||
t0 = time.perf_counter()
|
||||
if isinstance(query_result, pd.DataFrame):
|
||||
df = query_result
|
||||
else:
|
||||
# Convert LanceDB query result to DataFrame
|
||||
# (this is where the actual DB query executes)
|
||||
df = query_result.to_pandas()
|
||||
logger.info(
|
||||
"search.execute rows=%d took %.3fs",
|
||||
len(df),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
# Extract scores
|
||||
scores = extract_scores(df)
|
||||
|
||||
# Convert DataFrame rows to ChunkRecords using to_dict
|
||||
# (avoids slow .iterrows() overhead)
|
||||
t0 = time.perf_counter()
|
||||
rows = df.to_dict(orient="records")
|
||||
pydantic_results = [
|
||||
self.store.ChunkRecord(
|
||||
|
|
@ -493,11 +530,17 @@ class ChunkRepository:
|
|||
)
|
||||
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
|
||||
document_ids = list(set(chunk.document_id for chunk in pydantic_results))
|
||||
|
||||
# Batch fetch all documents at once
|
||||
t0 = time.perf_counter()
|
||||
documents_map = {}
|
||||
if document_ids:
|
||||
# Use IN clause for efficient batch lookup
|
||||
|
|
@ -509,8 +552,14 @@ class ChunkRepository:
|
|||
.to_pydantic(DocumentRecord)
|
||||
)
|
||||
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
|
||||
t0 = time.perf_counter()
|
||||
chunks_with_scores = []
|
||||
for i, chunk_record in enumerate(pydantic_results):
|
||||
doc = documents_map.get(chunk_record.document_id)
|
||||
|
|
@ -526,5 +575,10 @@ class ChunkRepository:
|
|||
)
|
||||
score = scores[i] if i < len(scores) else 1.0
|
||||
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
|
||||
|
|
|
|||
Loading…
Reference in a new issue