fix(graphrag): query-relevant community reports, rerank floor parity

_community_context accepted the query but never used it: every search
injected the same global top-k community reports sorted by static rank,
contradicting the module's "relevant community reports" contract
(unfinished wiring; the scaffold commit documents the data model but
never connected it). Select reports by reranking the top-30
static-ranked candidates against the query with the existing Cohere
client (one bounded rerank call per query); scores stay pinned at 0.0
so reports remain context-only.

Also applies the classical retriever's RERANK_SCORE_FLOOR in
_apply_token_budget (exempting the score-0.0 community reports) per the
module's fairness contract, and skips oversized results instead of
breaking, matching the classical fix.

Note: a future re-run of the graphrag eval arm will not be directly
comparable to the recorded F28 numbers, which measured the
static-report behavior.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

app/graph_rag.py

@@ -40,6 +40,7 @@ from typing import Any
 from .chroma_rag import (
     DEFAULT_CONTEXT_TOKEN_BUDGET,
     DEFAULT_RERANK_MODEL,
+    RERANK_SCORE_FLOOR,
     SearchResult,
     get_token_encoding,
     rerank_results,
@@ -63,6 +64,10 @@ GRAPHRAG_COMMUNITY_SOURCE = "graphrag_community"
 DEFAULT_ENTITY_TOP_K = 20
 DEFAULT_COMMUNITY_TOP_K = 3
 DEFAULT_RERANK_TOP_K = 5
+# Candidate pool for query-relevant community-report selection: only the top
+# reports by the index's static rank are reranked against the query, so the
+# per-query cost stays bounded (one rerank call over at most this many docs).
+COMMUNITY_RERANK_CANDIDATES = 30
 
 
 class GraphRAGIndexNotBuilt(RuntimeError):
@@ -211,8 +216,8 @@ class GraphRAGRetriever:
         """Assemble GraphRAG context as reranked SearchResult chunks.
 
         Steps: embed query -> nearest entities (LanceDB) -> their text units
-        (mapped to real source/url via document_id) + top community reports ->
-        Cohere rerank -> token budget. Returns [] on any failure so the agent
+        (mapped to real source/url via document_id) + query-relevant community
+        reports -> Cohere rerank -> token budget. Returns [] on any failure so the agent
         degrades to KB browsing, matching the classical backend's contract.
         """
         try:
@@ -287,19 +292,27 @@ class GraphRAGRetriever:
         )
 
     def _community_context(self, query: str) -> list[SearchResult]:
-        """Top community reports as context-only chunks (no source -> no recall
-        credit, by design)."""
+        """Query-relevant community reports as context-only chunks (synthetic
+        source, no url -> no recall/citation credit, by design).
+
+        Candidates are the ``COMMUNITY_RERANK_CANDIDATES`` highest-ranked
+        reports (the index's static rank); the ``community_top_k`` most relevant
+        to the query are selected with the same Cohere rerank used for text
+        units. Scores are pinned back to 0.0 so the reports stay context-only
+        chunks (and stay exempt from the rerank-score floor in
+        ``_apply_token_budget``).
+        """
         if self._reports is None or self._community_top_k <= 0:
             return []
         reports = self._reports
         if "rank" in reports.columns:
             reports = reports.sort_values("rank", ascending=False)
-        results: list[SearchResult] = []
-        for _, row in reports.head(self._community_top_k).iterrows():
+        candidates: list[SearchResult] = []
+        for _, row in reports.head(COMMUNITY_RERANK_CANDIDATES).iterrows():
             content = str(row.get("full_content") or row.get("summary") or "")
             if not content:
                 continue
-            results.append(
+            candidates.append(
                 SearchResult(
                     chunk_id=f"community:{row.get('community', '')}",
                     doc_id="",
@@ -315,7 +328,18 @@ class GraphRAGRetriever:
                     retrieval_method="graphrag_community",
                 )
             )
-        return results
+        if len(candidates) <= self._community_top_k:
+            return candidates
+        selected = rerank_results(
+            self._cohere,
+            query,
+            candidates,
+            model=self._rerank_model,
+            top_n=self._community_top_k,
+        )
+        for result in selected:
+            result.score = 0.0
+        return selected
 
     def _apply_token_budget(
         self, results: list[SearchResult], token_budget: int | None
@@ -324,11 +348,22 @@ class GraphRAGRetriever:
         filtered: list[SearchResult] = []
         total = 0
         for result in results:
+            # Same low-relevance floor as the classical retriever, so both eval
+            # arms drop weak reranked hits (the module's fairness contract).
+            # Community reports are context-only chunks pinned at score 0.0 and
+            # are exempt, or the floor would drop them all.
+            if (
+                result.source != GRAPHRAG_COMMUNITY_SOURCE
+                and result.score < RERANK_SCORE_FLOOR
+            ):
+                continue
             result_tokens = len(
                 self._encoding.encode(result.content, disallowed_special=())
             )
             if total + result_tokens > budget:
-                break
+                # Skip an oversized result instead of cutting off the whole
+                # list; keep filling the budget with later results that fit.
+                continue
             total += result_tokens
             filtered.append(result)
         return filtered

tests/test_graph_rag.py

@@ -14,7 +14,9 @@ pytest.importorskip("graphrag", reason="install the optional 'graphrag' extra to
 import pandas as pd
 
 from app.chat_types import ChatRequest
+from app.chroma_rag import SearchResult, get_token_encoding
 from app.graph_rag import (
+    COMMUNITY_RERANK_CANDIDATES,
     GRAPHRAG_COMMUNITY_SOURCE,
     GraphRAGIndexNotBuilt,
     GraphRAGRetriever,
@@ -22,6 +24,40 @@ from app.graph_rag import (
 )
 
 
+class _FakeRerankItem:
+    def __init__(self, index: int, score: float) -> None:
+        self.index = index
+        self.relevance_score = score
+
+
+class _FakeRerankResponse:
+    def __init__(self, items: list[_FakeRerankItem]) -> None:
+        self.results = items
+
+
+class _FakeCohere:
+    """Captures rerank calls; ranks documents containing the query text first."""
+
+    def __init__(self) -> None:
+        self.calls: list[dict] = []
+
+    def rerank(self, *, model, query, documents, top_n):  # type: ignore[no-untyped-def]
+        self.calls.append(
+            {"model": model, "query": query, "documents": list(documents)}
+        )
+        order = sorted(
+            range(len(documents)),
+            key=lambda i: (query.lower() in documents[i].lower(), -i),
+            reverse=True,
+        )
+        return _FakeRerankResponse(
+            [
+                _FakeRerankItem(doc_index, 1.0 - rank * 0.1)
+                for rank, doc_index in enumerate(order[:top_n])
+            ]
+        )
+
+
 class _FakeSearch:
     def __init__(self, df: pd.DataFrame) -> None:
         self._df = df
@@ -120,6 +156,8 @@ class GraphRagMappingTestCase(unittest.TestCase):
         self.assertEqual({res.source for res in only_a}, {"src_a"})
 
     def test_community_reports_are_context_only(self) -> None:
+        # With candidates <= community_top_k every report is returned without a
+        # rerank call (no _cohere set on the bare retriever proves that).
         r = _bare_retriever()
         r._community_top_k = 2
         r._reports = pd.DataFrame(
@@ -137,9 +175,132 @@ class GraphRagMappingTestCase(unittest.TestCase):
             # no url -> never resolves as a cited source card.
             self.assertEqual(res.source, GRAPHRAG_COMMUNITY_SOURCE)
             self.assertEqual(res.url, "")
+            self.assertEqual(res.score, 0.0)
         # Sorted by rank desc: highest-rank community first.
         self.assertEqual(results[0].title, "C1")
 
+    def test_community_reports_selected_by_query_relevance(self) -> None:
+        # The query, not the static community rank, decides which reports come
+        # back: the top-ranked candidates are reranked against the query.
+        r = _bare_retriever()
+        r._community_top_k = 1
+        r._rerank_model = "fake-rerank"
+        fake = _FakeCohere()
+        r._cohere = fake
+        r._reports = pd.DataFrame(
+            {
+                "community": [1, 2, 3, 4],
+                "title": ["C1", "C2", "C3", "C4"],
+                "full_content": [
+                    "agents and tools",
+                    "prompt engineering tips",
+                    "vector databases overview",
+                    "fine-tuning walkthrough",
+                ],
+                "rank": [9.0, 8.0, 7.0, 6.0],
+            }
+        )
+
+        results = r._community_context("vector databases")
+
+        # Not C1 (highest static rank): the query-relevant report wins.
+        self.assertEqual([res.title for res in results], ["C3"])
+        # Every candidate reached the reranker in one call.
+        self.assertEqual(len(fake.calls), 1)
+        self.assertEqual(len(fake.calls[0]["documents"]), 4)
+        # Still context-only: synthetic source, no url, score pinned to 0.0.
+        for res in results:
+            self.assertEqual(res.source, GRAPHRAG_COMMUNITY_SOURCE)
+            self.assertEqual(res.url, "")
+            self.assertEqual(res.score, 0.0)
+            self.assertEqual(res.retrieval_method, "graphrag_community")
+
+    def test_community_rerank_candidate_pool_is_bounded(self) -> None:
+        # Cost stays bounded: only the top COMMUNITY_RERANK_CANDIDATES reports
+        # by static rank are sent to the reranker.
+        r = _bare_retriever()
+        r._community_top_k = 1
+        r._rerank_model = "fake-rerank"
+        fake = _FakeCohere()
+        r._cohere = fake
+        total = COMMUNITY_RERANK_CANDIDATES + 5
+        r._reports = pd.DataFrame(
+            {
+                "community": list(range(total)),
+                "title": [f"C{i}" for i in range(total)],
+                "full_content": [f"report {i}" for i in range(total)],
+                "rank": [float(total - i) for i in range(total)],
+            }
+        )
+
+        r._community_context("q")
+
+        documents = fake.calls[0]["documents"]
+        self.assertEqual(len(documents), COMMUNITY_RERANK_CANDIDATES)
+        self.assertIn("report 0", documents)  # highest static rank kept
+        self.assertNotIn(f"report {total - 1}", documents)  # lowest dropped
+
+
+class GraphRagTokenBudgetTestCase(unittest.TestCase):
+    def _retriever(self) -> GraphRAGRetriever:
+        r = _bare_retriever()
+        r._encoding = get_token_encoding(None)
+        r._token_budget = 100_000
+        return r
+
+    def _result(
+        self,
+        chunk_id: str,
+        score: float,
+        *,
+        content: str = "x",
+        source: str = "src_a",
+        method: str = "graphrag",
+    ) -> SearchResult:
+        return SearchResult(
+            chunk_id=chunk_id,
+            doc_id=chunk_id,
+            title=chunk_id,
+            url="",
+            source=source,
+            retrieve_doc=False,
+            tokens=10,
+            score=score,
+            content=content,
+            chunk_content=content,
+            heading_path="",
+            retrieval_method=method,
+        )
+
+    def test_budget_skips_oversized_result_and_fills_with_smaller(self) -> None:
+        # Same regression as the classical retriever: an oversized rank-1
+        # result under a small per-request budget must not empty the list.
+        r = self._retriever()
+        oversized = self._result("big", 0.9, content="word " * 500)
+        small = self._result("small", 0.8, content="word " * 15)
+
+        kept = r._apply_token_budget([oversized, small], token_budget=50)
+
+        self.assertEqual([res.chunk_id for res in kept], ["small"])
+
+    def test_low_score_floor_exempts_community_reports(self) -> None:
+        # Fairness with the classical arm: weak reranked text units are dropped
+        # by the same score floor, but community reports (context-only chunks
+        # pinned at score 0.0) are exempt.
+        r = self._retriever()
+        strong = self._result("strong", 0.5)
+        weak = self._result("weak", 0.05)
+        community = self._result(
+            "community:1",
+            0.0,
+            source=GRAPHRAG_COMMUNITY_SOURCE,
+            method="graphrag_community",
+        )
+
+        kept = r._apply_token_budget([strong, weak, community], token_budget=None)
+
+        self.assertEqual([res.chunk_id for res in kept], ["strong", "community:1"])
+
 
 if __name__ == "__main__":
     unittest.main()