fix: more performance tweaks
feat: add perf logging
This commit is contained in:
bryan davis
parent
9da345067c
commit
d95d8291af
3 changed files with 198 additions and 35 deletions
|
|
@ -894,13 +894,21 @@ class HaikuRAG:
|
|||
return doc
|
||||
|
||||
safe_input = _escape_sql_string(id_or_title)
|
||||
docs = await self.list_documents(filter=f"title = '{safe_input}'")
|
||||
if docs and docs[0].id:
|
||||
return await self.get_document_by_id(docs[0].id)
|
||||
docs = await self.list_documents(
|
||||
filter=f"title = '{safe_input}'",
|
||||
include_content=True,
|
||||
limit=1,
|
||||
)
|
||||
if docs:
|
||||
return docs[0]
|
||||
|
||||
docs = await self.list_documents(filter=f"uri = '{safe_input}'")
|
||||
if docs and docs[0].id:
|
||||
return await self.get_document_by_id(docs[0].id)
|
||||
docs = await self.list_documents(
|
||||
filter=f"uri = '{safe_input}'",
|
||||
include_content=True,
|
||||
limit=1,
|
||||
)
|
||||
if docs:
|
||||
return docs[0]
|
||||
|
||||
return None
|
||||
|
||||
|
|
@ -1055,24 +1063,95 @@ class HaikuRAG:
|
|||
Returns:
|
||||
List of SearchResult objects ordered by relevance.
|
||||
"""
|
||||
import time
|
||||
|
||||
import logfire
|
||||
|
||||
search_start = time.perf_counter()
|
||||
|
||||
if limit is None:
|
||||
limit = self._config.search.limit
|
||||
|
||||
# Step 1: Get reranker
|
||||
t0 = time.perf_counter()
|
||||
reranker = get_reranker(config=self._config)
|
||||
logger.debug(
|
||||
"search reranker_init took %.3fs",
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
# Step 2: Chunk search
|
||||
t0 = time.perf_counter()
|
||||
if reranker is None:
|
||||
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",
|
||||
search_type,
|
||||
limit,
|
||||
len(chunk_results),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
else:
|
||||
search_limit = limit * 10
|
||||
raw_results = await self.chunk_repository.search(
|
||||
query, search_limit, search_type, filter
|
||||
)
|
||||
chunks = [chunk for chunk, _ in raw_results]
|
||||
chunk_results = await reranker.rerank(query, chunks, top_n=limit)
|
||||
logger.debug(
|
||||
"search chunk_search type=%s limit=%d results=%d took %.3fs",
|
||||
search_type,
|
||||
search_limit,
|
||||
len(raw_results),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
return [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results]
|
||||
# Step 3: Reranking
|
||||
t0 = time.perf_counter()
|
||||
chunks = [chunk for chunk, _ in raw_results]
|
||||
chunk_results = await reranker.rerank(
|
||||
query, chunks, top_n=limit
|
||||
)
|
||||
logger.debug(
|
||||
"search rerank candidates=%d top_n=%d took %.3fs",
|
||||
len(chunks),
|
||||
limit,
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
# Step 4: Build SearchResult objects
|
||||
t0 = time.perf_counter()
|
||||
results = [
|
||||
SearchResult.from_chunk(chunk, score)
|
||||
for chunk, score in chunk_results
|
||||
]
|
||||
logger.debug(
|
||||
"search build_results count=%d took %.3fs",
|
||||
len(results),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
duration_s = time.perf_counter() - search_start
|
||||
logger.info(
|
||||
"search completed query=%r type=%s results=%d duration=%.3fs",
|
||||
query[:80],
|
||||
search_type,
|
||||
len(results),
|
||||
duration_s,
|
||||
)
|
||||
logfire.metric_histogram(
|
||||
"search.duration",
|
||||
unit="s",
|
||||
).record(
|
||||
duration_s,
|
||||
attributes={
|
||||
"search_type": search_type,
|
||||
"result_count": len(results),
|
||||
"reranked": reranker is not None,
|
||||
},
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
async def expand_context(
|
||||
self,
|
||||
|
|
@ -1114,19 +1193,18 @@ class HaikuRAG:
|
|||
expanded_results.extend(doc_results)
|
||||
continue
|
||||
|
||||
# Fetch the document to get DoclingDocument
|
||||
doc = await self.get_document_by_id(doc_id)
|
||||
if doc is None:
|
||||
expanded_results.extend(doc_results)
|
||||
continue
|
||||
has_refs = any(r.doc_item_refs for r in doc_results)
|
||||
docling_doc = None
|
||||
|
||||
if has_refs:
|
||||
# Only fetch docling data (skip content blob)
|
||||
doc = await self.document_repository.get_docling_data(
|
||||
doc_id
|
||||
)
|
||||
if doc is not None:
|
||||
docling_doc = doc.get_docling_document()
|
||||
|
||||
# Check if we can use DoclingDocument-based expansion
|
||||
has_docling = docling_doc is not None
|
||||
has_refs = any(r.doc_item_refs for r in doc_results)
|
||||
|
||||
if has_docling and has_refs:
|
||||
if docling_doc is not None and has_refs:
|
||||
# Use DoclingDocument-based expansion
|
||||
expanded = await self._expand_with_docling(
|
||||
doc_results,
|
||||
|
|
@ -1320,12 +1398,14 @@ class HaikuRAG:
|
|||
Structural content (tables, code, lists) expands to complete structures.
|
||||
Text content uses radius-based expansion.
|
||||
"""
|
||||
all_items = list(docling_doc.iterate_items())
|
||||
ref_to_index = {
|
||||
getattr(item, "self_ref", None): i
|
||||
for i, (item, _) in enumerate(all_items)
|
||||
if getattr(item, "self_ref", None)
|
||||
}
|
||||
# Single-pass: build items list and ref index together
|
||||
all_items = []
|
||||
ref_to_index = {}
|
||||
for i, item_tuple in enumerate(docling_doc.iterate_items()):
|
||||
all_items.append(item_tuple)
|
||||
ref = getattr(item_tuple[0], "self_ref", None)
|
||||
if ref:
|
||||
ref_to_index[ref] = i
|
||||
|
||||
# Compute expanded ranges
|
||||
ranges: list[tuple[int, int, SearchResult]] = []
|
||||
|
|
@ -1403,25 +1483,46 @@ class HaikuRAG:
|
|||
radius: int,
|
||||
) -> list[SearchResult]:
|
||||
"""Expand results using chunk-based adjacency."""
|
||||
all_chunks = await self.chunk_repository.get_by_document_id(doc_id)
|
||||
if not all_chunks:
|
||||
# Get orders for result chunks (lightweight query)
|
||||
chunk_ids = [r.chunk_id for r in results if r.chunk_id]
|
||||
if not chunk_ids:
|
||||
return results
|
||||
|
||||
content_to_chunk = {c.content: c for c in all_chunks}
|
||||
chunk_by_order = {c.order: c for c in all_chunks}
|
||||
min_order, max_order = min(chunk_by_order.keys()), max(chunk_by_order.keys())
|
||||
id_to_order = await self.chunk_repository.get_orders_by_ids(
|
||||
chunk_ids
|
||||
)
|
||||
if not id_to_order:
|
||||
return results
|
||||
|
||||
# Compute the order range we actually need
|
||||
orders = list(id_to_order.values())
|
||||
range_min = min(orders) - radius
|
||||
range_max = max(orders) + radius
|
||||
|
||||
# Fetch only chunks in the needed range
|
||||
chunks_in_doc = (
|
||||
await self.chunk_repository.get_by_document_id_order_range(
|
||||
doc_id, range_min, range_max
|
||||
)
|
||||
)
|
||||
if not chunks_in_doc:
|
||||
return results
|
||||
|
||||
chunk_by_order = {c.order: c for c in chunks_in_doc}
|
||||
actual_min = min(chunk_by_order.keys())
|
||||
actual_max = max(chunk_by_order.keys())
|
||||
|
||||
# Build ranges
|
||||
ranges: list[tuple[int, int, SearchResult]] = []
|
||||
passthrough: list[SearchResult] = []
|
||||
|
||||
for result in results:
|
||||
chunk = content_to_chunk.get(result.content)
|
||||
if chunk is None:
|
||||
order = id_to_order.get(result.chunk_id)
|
||||
if order is None:
|
||||
passthrough.append(result)
|
||||
continue
|
||||
start = max(min_order, chunk.order - radius)
|
||||
end = min(max_order, chunk.order + radius)
|
||||
start = max(actual_min, order - radius)
|
||||
end = min(actual_max, order + radius)
|
||||
ranges.append((start, end, result))
|
||||
|
||||
# Merge and build results
|
||||
|
|
|
|||
|
|
@ -407,6 +407,44 @@ class ChunkRepository:
|
|||
for rec in results
|
||||
]
|
||||
|
||||
async def get_orders_by_ids(self, chunk_ids: list[str]) -> dict[str, int]:
|
||||
"""Get order values for specific chunk IDs."""
|
||||
if not chunk_ids:
|
||||
return {}
|
||||
id_list = "', '".join(chunk_ids)
|
||||
rows = list(
|
||||
self.store.chunks_table.search()
|
||||
.select(["id", "order"])
|
||||
.where(f"id IN ('{id_list}')")
|
||||
.to_list()
|
||||
)
|
||||
return {str(row["id"]): int(row["order"]) for row in rows}
|
||||
|
||||
async def get_by_document_id_order_range(
|
||||
self, document_id: str, min_order: int, max_order: int
|
||||
) -> list[Chunk]:
|
||||
"""Get chunks for a document within an order range."""
|
||||
where = (
|
||||
f"document_id = '{document_id}'"
|
||||
f" AND `order` >= {min_order}"
|
||||
f" AND `order` <= {max_order}"
|
||||
)
|
||||
results = list(
|
||||
self.store.chunks_table.search()
|
||||
.where(where)
|
||||
.to_pydantic(self.store.ChunkRecord)
|
||||
)
|
||||
return [
|
||||
Chunk(
|
||||
id=rec.id,
|
||||
document_id=rec.document_id,
|
||||
content=rec.content,
|
||||
metadata=json.loads(rec.metadata),
|
||||
order=rec.order,
|
||||
)
|
||||
for rec in results
|
||||
]
|
||||
|
||||
async def _process_search_results(
|
||||
self, query_result: "pd.DataFrame | LanceQueryBuilder"
|
||||
) -> list[tuple[Chunk, float]]:
|
||||
|
|
|
|||
|
|
@ -90,6 +90,30 @@ class DocumentRepository:
|
|||
|
||||
return self._record_to_document(results[0])
|
||||
|
||||
_DOCLING_COLUMNS = ["id", "docling_document", "docling_version"]
|
||||
|
||||
async def get_docling_data(self, entity_id: str) -> Document | None:
|
||||
"""Get document with only docling data loaded (skips content blob)."""
|
||||
safe_id = _escape_sql_string(entity_id)
|
||||
results = list(
|
||||
self.store.documents_table.search()
|
||||
.select(self._DOCLING_COLUMNS)
|
||||
.where(f"id = '{safe_id}'")
|
||||
.limit(1)
|
||||
.to_list()
|
||||
)
|
||||
|
||||
if not results:
|
||||
return None
|
||||
|
||||
row = results[0]
|
||||
return Document(
|
||||
id=row["id"],
|
||||
content="",
|
||||
docling_document=row.get("docling_document"),
|
||||
docling_version=row.get("docling_version"),
|
||||
)
|
||||
|
||||
async def update(self, entity: Document) -> Document:
|
||||
"""Update an existing document."""
|
||||
self.store._assert_writable()
|
||||
|
|
|
|||
Loading…
Reference in a new issue