Add diversity-aware selection after Document-RAG reranking (#1014)

* Add Document-RAG diversity selection helper

* Add optional MMR diversity selection after reranking

* Fix Document-RAG diversity test method signatures

trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py

@@ -20,6 +20,8 @@ from trustgraph.provenance import (
     GRAPH_RETRIEVAL,
 )
 
+from .rerank import RerankCandidate, mmr_select
+
 # Module logger
 logger = logging.getLogger(__name__)
 
@@ -150,6 +152,8 @@ class DocumentRag:
             fetch_chunk,
             reranker_client=None,
             verbose=False,
+            rerank_diversity_mode="none",
+            rerank_diversity_lambda=0.7,
     ):
 
         self.verbose = verbose
@@ -162,6 +166,8 @@ class DocumentRag:
         # Optional cross-encoder reranker. When None, the retrieval path is
         # byte-identical to the pre-reranker behaviour.
         self.reranker_client = reranker_client
+        self.rerank_diversity_mode = rerank_diversity_mode
+        self.rerank_diversity_lambda = rerank_diversity_lambda
 
         if self.verbose:
             logger.debug("DocumentRag initialized")
@@ -277,30 +283,74 @@ class DocumentRag:
         # skipped entirely and behaviour is byte-identical to before.
         reranked = False
         if self.reranker_client is not None and docs:
+            use_diversity = self.rerank_diversity_mode == "mmr"
+
+            # Without diversity selection, preserve the existing #1011
+            # behavior: ask the reranker for exactly doc_limit results.
+            #
+            # With diversity selection enabled, ask the reranker to score the
+            # full fetched candidate pool first, then let MMR choose the final
+            # doc_limit context set.
+            rerank_limit = len(docs) if use_diversity else doc_limit
+
             results = await self.reranker_client.rerank(
                 queries=[{"id": "0", "text": query}],
                 documents=[
                     {"id": str(i), "text": d} for i, d in enumerate(docs)
                 ],
-                # Narrow the over-fetched candidate pool down to the final
-                # doc_limit requested for synthesis.
-                limit=doc_limit,
+                limit=rerank_limit,
             )
 
-            # results are sorted desc by score and truncated to limit by the
-            # reranker service, so order gives the surviving top-N directly.
-            order = [int(r.document_id) for r in results]
-            docs = [docs[i] for i in order]
-            chunk_ids = [chunk_ids[i] for i in order]
+            source_docs = docs
+            source_chunk_ids = chunk_ids
+
+            if use_diversity:
+                candidates = [
+                    RerankCandidate(
+                        index=int(r.document_id),
+                        chunk_id=source_chunk_ids[int(r.document_id)],
+                        text=source_docs[int(r.document_id)],
+                        reranker_score=r.score,
+                    )
+                    for r in results
+                ]
+
+                selected_candidates = mmr_select(
+                    candidates,
+                    limit=doc_limit,
+                    lambda_mult=self.rerank_diversity_lambda,
+                )
+
+                docs = [candidate.text for candidate in selected_candidates]
+                chunk_ids = [
+                    candidate.chunk_id for candidate in selected_candidates
+                ]
+
+                selected_chunks_with_scores = [
+                    {
+                        "chunk_id": candidate.chunk_id,
+                        "score": candidate.reranker_score,
+                    }
+                    for candidate in selected_candidates
+                ]
+
+            else:
+                # results are sorted desc by score and truncated to limit by the
+                # reranker service, so order gives the surviving top-N directly.
+                order = [int(r.document_id) for r in results]
+                docs = [source_docs[i] for i in order]
+                chunk_ids = [source_chunk_ids[i] for i in order]
+
+                selected_chunks_with_scores = [
+                    {"chunk_id": chunk_ids[i], "score": r.score}
+                    for i, r in enumerate(results)
+                ]
+
             reranked = True
 
             # Emit chunk-selection (focus) explainability: surviving chunks
             # with their cross-encoder scores, derived from exploration.
             if explain_callback:
-                selected_chunks_with_scores = [
-                    {"chunk_id": chunk_ids[i], "score": r.score}
-                    for i, r in enumerate(results)
-                ]
                 foc_triples = set_graph(
                     docrag_chunk_selection_triples(
                         foc_uri, exp_uri,

trustgraph-flow/trustgraph/retrieval/document_rag/rag.py

@@ -33,17 +33,23 @@ class Processor(FlowProcessor):
         # reranking; the rerank step narrows it back down to doc_limit for the
         # LLM. 0 means the core derives it (OVERFETCH_FACTOR x doc_limit).
         fetch_limit = params.get("fetch_limit", 0)
+        rerank_diversity_mode = params.get("rerank_diversity_mode", "none")
+        rerank_diversity_lambda = params.get("rerank_diversity_lambda", 0.7)
 
         super(Processor, self).__init__(
             **params | {
                 "id": id,
                 "doc_limit": doc_limit,
                 "fetch_limit": fetch_limit,
+                "rerank_diversity_mode": rerank_diversity_mode,
+                "rerank_diversity_lambda": rerank_diversity_lambda,
             }
         )
 
         self.doc_limit = doc_limit
         self.fetch_limit = fetch_limit
+        self.rerank_diversity_mode = rerank_diversity_mode
+        self.rerank_diversity_lambda = rerank_diversity_lambda
 
         self.register_specification(
             ConsumerSpec(
@@ -122,6 +128,8 @@ class Processor(FlowProcessor):
                 fetch_chunk = fetch_chunk,
                 reranker_client = flow("reranker-request"),
                 verbose=True,
+                rerank_diversity_mode=self.rerank_diversity_mode,
+                rerank_diversity_lambda=self.rerank_diversity_lambda,
             )
 
             if v.doc_limit:
@@ -277,6 +285,20 @@ class Processor(FlowProcessor):
                  '(default: derive from doc-limit)'
         )
 
+        parser.add_argument(
+            '--rerank-diversity-mode',
+            choices=['none', 'mmr'],
+            default='none',
+            help='Optional diversity-aware selection after reranking (default: none)'
+        )
+
+        parser.add_argument(
+            '--rerank-diversity-lambda',
+            type=float,
+            default=0.7,
+            help='MMR relevance/diversity tradeoff, higher values prefer relevance'
+        )
+
 def run():
 
     Processor.launch(default_ident, __doc__)

trustgraph-flow/trustgraph/retrieval/document_rag/rerank.py

@@ -0,0 +1,142 @@
+import re
+from dataclasses import dataclass, replace
+from typing import List, Sequence, Set
+
+
+@dataclass(frozen=True)
+class RerankCandidate:
+    """
+    Candidate chunk after cross-encoder reranking.
+
+    reranker_score is the raw score returned by the reranker backend. It may
+    not be normalized, so MMR should use normalized_score instead.
+    """
+    index: int
+    chunk_id: str
+    text: str
+    reranker_score: float
+    normalized_score: float = 0.0
+
+
+_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
+
+
+def _clamp01(value: float) -> float:
+    return max(0.0, min(1.0, value))
+
+
+def _token_set(text: str) -> Set[str]:
+    return set(token.lower() for token in _TOKEN_RE.findall(text or ""))
+
+
+def _jaccard(a: str, b: str) -> float:
+    a_tokens = _token_set(a)
+    b_tokens = _token_set(b)
+
+    if not a_tokens or not b_tokens:
+        return 0.0
+
+    return len(a_tokens & b_tokens) / len(a_tokens | b_tokens)
+
+
+def normalize_candidate_scores(
+    candidates: Sequence[RerankCandidate],
+) -> List[RerankCandidate]:
+    """
+    Min-max normalize reranker scores within the current candidate set.
+
+    Reranker backends may return different score scales: probabilities,
+    logits, or prompt-defined scores. MMR needs a stable [0, 1] relevance
+    signal, so normalize per candidate set instead of assuming a global range.
+    """
+    if not candidates:
+        return []
+
+    scores = [float(candidate.reranker_score) for candidate in candidates]
+    min_score = min(scores)
+    max_score = max(scores)
+
+    if max_score == min_score:
+        return [
+            replace(candidate, normalized_score=0.5)
+            for candidate in candidates
+        ]
+
+    score_range = max_score - min_score
+
+    return [
+        replace(
+            candidate,
+            normalized_score=(float(candidate.reranker_score) - min_score) / score_range,
+        )
+        for candidate in candidates
+    ]
+
+
+def _pair_diversity_penalty(
+    candidate: RerankCandidate,
+    selected: RerankCandidate,
+    token_overlap_weight: float,
+) -> float:
+    """
+    Pairwise diversity penalty between two candidate chunks.
+
+    The first revision only uses token overlap because the current Document-RAG
+    reranker document_id is the candidate index, not a source document id.
+    """
+    penalty = token_overlap_weight * _jaccard(candidate.text, selected.text)
+    return _clamp01(penalty)
+
+
+def mmr_select(
+    candidates: Sequence[RerankCandidate],
+    limit: int,
+    lambda_mult: float = 0.7,
+    token_overlap_weight: float = 1.0,
+) -> List[RerankCandidate]:
+    """
+    Select a diverse final context set using MMR.
+
+    Relevance comes from normalized cross-encoder reranker scores.
+    Diversity comes from token overlap against already selected chunks.
+    """
+    if limit <= 0:
+        return []
+
+    lambda_mult = _clamp01(lambda_mult)
+    token_overlap_weight = max(0.0, token_overlap_weight)
+
+    remaining = normalize_candidate_scores(candidates)
+    selected: List[RerankCandidate] = []
+
+    while remaining and len(selected) < limit:
+        best_idx = 0
+        best_score = None
+
+        for idx, candidate in enumerate(remaining):
+            relevance = candidate.normalized_score
+
+            if selected:
+                diversity_penalty = max(
+                    _pair_diversity_penalty(
+                        candidate,
+                        chosen,
+                        token_overlap_weight=token_overlap_weight,
+                    )
+                    for chosen in selected
+                )
+            else:
+                diversity_penalty = 0.0
+
+            mmr_score = (
+                lambda_mult * relevance
+                - (1.0 - lambda_mult) * diversity_penalty
+            )
+
+            if best_score is None or mmr_score > best_score:
+                best_score = mmr_score
+                best_idx = idx
+
+        selected.append(remaining.pop(best_idx))
+
+    return selected