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

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.

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.

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
 

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)

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)