From c17c623936f1838b1f70006d40b568ddef71b982 Mon Sep 17 00:00:00 2001 From: Jeroen Schweitzer Date: Tue, 14 Jul 2026 14:23:34 +0200 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_01QbFZyDvYksazX6nYQYZ67L --- CHANGELOG.md | 1 + src/services/graph_service.py | 62 +++++++++++++++ src/services/hybrid_rag_service.py | 93 +++++++++++++++-------- tests/test_enrichment_topk.py | 118 +++++++++++++++++++++++++++++ tests/test_tenant_scoping.py | 12 +++ 5 files changed, 255 insertions(+), 31 deletions(-) create mode 100644 tests/test_enrichment_topk.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 083af5f..410a03f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed (performance) +- **Top-k enrichment with one batched lookup** - Phase 3 (`_enrich_with_related_dossiers`) ran a sequential Neo4j round-trip for EVERY fused result and the final trim then discarded most of the output. It now enriches only the results that can still reach the response (the Phase 4 rerank slice of 20 when reranking is enabled, otherwise `final_result_count`) and resolves all of them in ONE UNWIND-batched Cypher query (`GraphService.get_related_documents_batch`, tenant-scoped like the single-page variant, per-page ordering/limit preserved). Unenriched tail results carry an empty `related_dossiers` list as before. - **Search persistence off the hot path, one atomic transaction** - HybridRAG Phase 6 (`_persist_search_for_librarian`) no longer gates the `/query/hybrid` response: the `search_id` is generated up front and returned immediately while the Neo4j write runs as a background task (strong task references held so tasks are not GC'd mid-flight). The write itself collapsed from ~21+ sequential auto-commit queries (SearchQuery node + per-document FOUND links + per-web-result WebResult nodes) into ONE UNWIND-based `execute_write` transaction, so a mid-way failure can no longer leave a partial SearchQuery graph behind. The persisted shape (SearchQuery properties incl. `processed: false`, tenant labels, `FOUND` relationship properties, WebResult properties) is unchanged and pinned by `tests/test_search_persistence.py` against exactly what the consolidation service queries. `timing.persistence_ms` now reports 0 (no longer on the request path). - **Batched embeddings + delete-last reindex** - `OllamaClient.embed_batch` now sends ONE batched `/api/embed` request (verified against the live Ollama; the old "batch" looped one `/api/embeddings` call per chunk) with a per-text fallback preserving partial-success semantics. `VectorService.update_from_page` embeds all chunks in that single call and upserts them in one Qdrant batch, and the reindex order is reversed: new points are upserted BEFORE stale points are pruned (deterministic uuid5 chunk ids make the overwrite safe), so a mid-way failure can no longer leave a page with zero vectors — the old order deleted everything first. The summary now reports `status` (`success`/`partial`/`failed`) and `chunks_skipped` instead of unconditional `success=True`; a fully failed embedding pass keeps the old vectors and reports failure. Measured on a real 7-chunk page ingest as `llm_tester` against the local server: ~375ms → ~181ms median (3 runs each). - **Document sync indexing fixed** - `DocumentSyncService._index_vectors` now awaits `ensure_collection` (the coroutine was created but never ran, so fresh tenants had no collection at upsert time), filters out `None` entries from `embed_batch` so one failed chunk embedding no longer aborts the whole document upsert (all-failed still reports failure), and routes the raw `client.delete`/`client.upsert` calls through the async wrapper (`delete_by_filter` / new batch `upsert_points`). Offline unit tests added. diff --git a/src/services/graph_service.py b/src/services/graph_service.py index 29e98c5..b87ff7b 100644 --- a/src/services/graph_service.py +++ b/src/services/graph_service.py @@ -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. diff --git a/src/services/hybrid_rag_service.py b/src/services/hybrid_rag_service.py index 6a575c1..a8ab1f9 100644 --- a/src/services/hybrid_rag_service.py +++ b/src/services/hybrid_rag_service.py @@ -44,6 +44,10 @@ class HybridRAGService: Service for HybridRAG multi-source search with fusion and re-ranking. """ + # Phase 4 reranking only ever considers this many fused results; results + # beyond the slice cannot reach the response when reranking is enabled. + RERANK_SLICE_SIZE = 20 + # Maps internal retrieval leg names to source_status keys in the response SOURCE_STATUS_KEYS = { "vector": "vector", @@ -148,10 +152,19 @@ class HybridRAGService: ) timing["fusion_ms"] = (time.time() - phase2_start) * 1000 - # Phase 3: Enrichment + # Phase 3: Enrichment — only for results that can still reach the + # response: the rerank slice (Phase 4 reorders within it) when + # reranking is on, otherwise just the final result count. Enriching + # the full fused set queried Neo4j per result and threw most of the + # output away at the final trim. phase3_start = time.time() if config.enable_enrichment: - enriched_results = await self._enrich_with_related_dossiers(fused_results, user) + enrich_top_k = config.final_result_count + if config.enable_reranking: + enrich_top_k = max(enrich_top_k, self.RERANK_SLICE_SIZE) + enriched_results = await self._enrich_with_related_dossiers( + fused_results, user, top_k=enrich_top_k + ) else: enriched_results = fused_results timing["enrichment_ms"] = (time.time() - phase3_start) * 1000 @@ -159,7 +172,9 @@ class HybridRAGService: # Phase 4: LLM Re-ranking phase4_start = time.time() if config.enable_reranking and len(enriched_results) > 1: - reranked_results = await self._rerank_with_llm(enriched_results[:20], query) + reranked_results = await self._rerank_with_llm( + enriched_results[:self.RERANK_SLICE_SIZE], query + ) else: reranked_results = enriched_results timing["reranking_ms"] = (time.time() - phase4_start) * 1000 @@ -736,49 +751,65 @@ JSON:""" async def _enrich_with_related_dossiers( self, results: List[Dict[str, Any]], - user: str + user: str, + top_k: Optional[int] = None ) -> List[Dict[str, Any]]: """ Phase 3: Enrich results with related documents via shared entities. + Only the top_k results are enriched (the rest get an empty + related_dossiers list — they cannot survive the final trim anyway), + and all lookups go through ONE UNWIND-batched Neo4j query instead of + a sequential round-trip per result. + Args: results: Fused results user: User identifier + top_k: How many leading results to enrich (None = all) Returns: Results with related_dossiers added """ + enrich_slice = results if top_k is None else results[:top_k] + + # Dedupe while preserving order; volatile/web results have no page_id + page_ids: List[int] = [] + for result in enrich_slice: + page_id = result.get("result", {}).get("page_id") + if page_id and page_id not in page_ids: + page_ids.append(page_id) + + try: + related_map = await self.graph.get_related_documents_batch( + page_ids=page_ids, + user=user, + limit_per_page=5 + ) + except Exception as e: + logger.warning(f"Failed batched related-docs lookup for {len(page_ids)} pages: {e}") + related_map = {} + for result in results: - result_data = result.get("result", {}) - page_id = result_data.get("page_id") + result["related_dossiers"] = [] - if page_id: - try: - related_docs = await self.graph.get_related_documents( - page_id=page_id, - user=user, - limit=5 - ) + for result in enrich_slice: + page_id = result.get("result", {}).get("page_id") + if not page_id: + continue - # Convert to RelatedDossier format - related_dossiers = [] - for doc in related_docs: - for tag in doc.get("tags", [])[:3]: # Max 3 tags per doc - related_dossiers.append({ - "page_id": doc["page_id"], - "title": doc["title"], - "path": doc["path"], - "tag": tag, - "shared_entities": doc["shared_entities"] - }) + # Convert to RelatedDossier format + related_dossiers = [] + for doc in related_map.get(page_id, []): + for tag in (doc.get("tags") or [])[:3]: # Max 3 tags per doc + related_dossiers.append({ + "page_id": doc["page_id"], + "title": doc["title"], + "path": doc["path"], + "tag": tag, + "shared_entities": doc["shared_entities"] + }) - result["related_dossiers"] = related_dossiers[:5] # Limit to 5 total - - except Exception as e: - logger.warning(f"Failed to get related docs for page {page_id}: {e}") - result["related_dossiers"] = [] - else: - result["related_dossiers"] = [] + result["related_dossiers"] = related_dossiers[:5] # Limit to 5 total return results diff --git a/tests/test_enrichment_topk.py b/tests/test_enrichment_topk.py new file mode 100644 index 0000000..97c53de --- /dev/null +++ b/tests/test_enrichment_topk.py @@ -0,0 +1,118 @@ +""" +Offline tests for Phase 3 enrichment running only on the top-k slice with +ONE batched related-documents lookup (perf: the old code ran a sequential +Neo4j query per fused result and the final trim discarded most of it). +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from src.services.hybrid_rag_service import HybridRAGService + +TENANT = "llm_tester" + + +@pytest.fixture +def graph(): + g = MagicMock() + g.get_related_documents_batch = AsyncMock(return_value={}) + return g + + +@pytest.fixture +def service(graph): + settings = MagicMock() + settings.ollama_llm_model = "test-model" + return HybridRAGService( + vector_service=MagicMock(), + graph_service=graph, + searxng_client=MagicMock(), + ollama_client=MagicMock(), + content_extractor=MagicMock(), + settings=settings, + ) + + +def _fused(n): + return [ + { + "result": {"page_id": i + 1, "title": f"page {i + 1}"}, + "rrf_score": 1.0 / (i + 1), + "sources": ["vector"], + } + for i in range(n) + ] + + +@pytest.mark.unit +class TestTopKEnrichment: + async def test_only_top_k_enriched_with_single_batched_lookup( + self, service, graph + ): + results = _fused(30) + + enriched = await service._enrich_with_related_dossiers( + results, user=TENANT, top_k=5 + ) + + # ONE batched lookup, only for the top-k page ids + graph.get_related_documents_batch.assert_awaited_once() + kwargs = graph.get_related_documents_batch.await_args.kwargs + assert kwargs["page_ids"] == [1, 2, 3, 4, 5] + assert kwargs["user"] == TENANT + # The single-page method must not be used anymore + graph.get_related_documents.assert_not_called() + + # Every result still carries the key (tail is empty) + assert all("related_dossiers" in r for r in enriched) + assert all(r["related_dossiers"] == [] for r in enriched[5:]) + + async def test_related_docs_mapped_onto_results(self, service, graph): + graph.get_related_documents_batch = AsyncMock(return_value={ + 1: [{ + "page_id": 9, "title": "rel", "path": "p/rel", + "tags": ["docker", "infra", "misc", "extra-tag-ignored"], + "shared_entities": 3, + }], + }) + results = _fused(3) + + enriched = await service._enrich_with_related_dossiers( + results, user=TENANT, top_k=2 + ) + + dossiers = enriched[0]["related_dossiers"] + assert len(dossiers) == 3 # max 3 tags per related doc + assert dossiers[0] == { + "page_id": 9, "title": "rel", "path": "p/rel", + "tag": "docker", "shared_entities": 3, + } + assert enriched[1]["related_dossiers"] == [] + assert enriched[2]["related_dossiers"] == [] + + async def test_results_without_page_id_are_skipped(self, service, graph): + results = [ + {"result": {"url": "http://x", "title": "web"}, "sources": ["web"]}, + {"result": {"page_id": 7, "title": "wiki"}, "sources": ["vector"]}, + ] + + enriched = await service._enrich_with_related_dossiers( + results, user=TENANT, top_k=10 + ) + + kwargs = graph.get_related_documents_batch.await_args.kwargs + assert kwargs["page_ids"] == [7] + assert enriched[0]["related_dossiers"] == [] + + async def test_batch_lookup_failure_degrades_gracefully(self, service, graph): + graph.get_related_documents_batch = AsyncMock( + side_effect=RuntimeError("neo4j down") + ) + results = _fused(3) + + enriched = await service._enrich_with_related_dossiers( + results, user=TENANT, top_k=3 + ) + + assert all(r["related_dossiers"] == [] for r in enriched) diff --git a/tests/test_tenant_scoping.py b/tests/test_tenant_scoping.py index 8cd83db..1195821 100644 --- a/tests/test_tenant_scoping.py +++ b/tests/test_tenant_scoping.py @@ -337,6 +337,18 @@ class TestEnrichmentScoping: assert f"(d1:{TENANT_DOC_LABEL}:Document" in cypher assert f"(d2:{TENANT_DOC_LABEL}:Document)" in cypher + async def test_batched_related_documents_scoped_to_tenant( + self, graph_service, mock_neo4j + ): + await graph_service.get_related_documents_batch( + page_ids=[1, 2], user=TENANT + ) + + cypher = _all_cypher(mock_neo4j) + assert "UNWIND $page_ids" in cypher + assert f"(d1:{TENANT_DOC_LABEL}:Document" in cypher + assert f"(d2:{TENANT_DOC_LABEL}:Document)" in cypher + # ============================================================================= # Persistence phase (SearchQuery / FOUND / WebResult)