perf: enrich only top-k results with one batched related-docs query
Phase 3 enrichment ran a sequential Neo4j query per fused result and the final trim then discarded most of the output. Enrichment now covers only results that can still reach the response - the Phase 4 rerank slice (RERANK_SLICE_SIZE = 20, results beyond it are dropped when reranking) or final_result_count, whichever applies - and resolves every page in a single UNWIND $page_ids Cypher query via the new tenant-scoped GraphService.get_related_documents_batch (per-page ordering by shared_entities and per-page limit preserved via ORDER BY + collect()). Unenriched tail results still carry related_dossiers: [] so the response shape is unchanged. Query validated with EXPLAIN against the live Neo4j. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QbFZyDvYksazX6nYQYZ67L
This commit is contained in:
@@ -1185,6 +1185,68 @@ Feel free to expand it with more details!
|
||||
logger.error(f"Failed to get related documents for page {page_id}: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
async def get_related_documents_batch(
|
||||
self,
|
||||
page_ids: List[int],
|
||||
user: str,
|
||||
limit_per_page: int = 5
|
||||
) -> Dict[int, List[Dict[str, Any]]]:
|
||||
"""
|
||||
Batch variant of get_related_documents: ONE UNWIND query for all pages
|
||||
instead of one round-trip per page.
|
||||
|
||||
Args:
|
||||
page_ids: Page IDs to find related documents for
|
||||
user: User identifier
|
||||
limit_per_page: Maximum related documents per page
|
||||
|
||||
Returns:
|
||||
Mapping of page_id -> related-document rows (same shape as
|
||||
get_related_documents). Pages with no related documents are
|
||||
absent from the mapping.
|
||||
"""
|
||||
from src.core.multi_tenancy import get_neo4j_user_base_label
|
||||
|
||||
if not page_ids:
|
||||
return {}
|
||||
|
||||
user_base_label = get_neo4j_user_base_label(user)
|
||||
user_doc_label = get_neo4j_user_label(user)
|
||||
|
||||
# ORDER BY runs before collect() so each page's list is sorted by
|
||||
# shared_entities descending; [..$limit] trims per page.
|
||||
query = f"""
|
||||
UNWIND $page_ids AS pid
|
||||
MATCH (d1:{user_doc_label}:Document {{page_id: pid}})
|
||||
MATCH (d1)-[:MENTIONS]->(e:{user_base_label})<-[:MENTIONS]-(d2:{user_doc_label}:Document)
|
||||
WHERE d1 <> d2 AND NOT e:Document
|
||||
WITH pid, d2, d2.tags as tags, count(DISTINCT e) as shared_entities
|
||||
WHERE tags IS NOT NULL AND size(tags) > 0
|
||||
WITH pid, d2, tags, shared_entities
|
||||
ORDER BY shared_entities DESC
|
||||
WITH pid, collect({{
|
||||
page_id: d2.page_id,
|
||||
title: d2.title,
|
||||
path: d2.path,
|
||||
tags: tags,
|
||||
shared_entities: shared_entities
|
||||
}})[..$limit] AS related
|
||||
RETURN pid AS page_id, related
|
||||
"""
|
||||
|
||||
try:
|
||||
rows = await self.neo4j.execute_query(
|
||||
query,
|
||||
{"page_ids": page_ids, "limit": limit_per_page}
|
||||
)
|
||||
return {row["page_id"]: row["related"] for row in rows}
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to get related documents for {len(page_ids)} pages: {e}",
|
||||
exc_info=True
|
||||
)
|
||||
return {}
|
||||
|
||||
async def get_all_entities(self, user: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all entities from the knowledge graph for a user.
|
||||
|
||||
Reference in New Issue
Block a user