[FEATURE] Add a pluggable reranker interface so users can supply custom reranking strategies

[acarbonetto]

Package

lexical-graph

Problem statement

Reranking strategies are selected by string name and dispatched through a hardcoded if/elif chain in the processors. In lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_statements.py:250-258:

  reranker = self.args.reranker.lower()
  if reranker == 'model':
      scored_values = self._score_values(...)
  elif reranker == 'tfidf':
      scored_values = self._score_values_with_tfidf(...)
  elif reranker == 'bedrock':
      scored_values = self._score_values_with_bedrock(...)
  else:
      return search_results

rerank_topics.py repeats the same pattern for topic_reranker ('none' | 'tfidf' | 'bedrock').

Consequences:

• Every new strategy requires a change to toolkit source. Adding Bedrock support ([FEATURE] Add reranker fallback chains to statement reranking #406) meant adding a _score_values_with_bedrock method plus a new branch plus validation of the new name — the same cost for Cohere, a SageMaker endpoint, a cross-encoder, or any customer-specific business logic.
• The dispatch method is a poor extension surface. Provider wiring lives inline in the processor (rerank_statements.py:157-206 constructs the boto3 client, builds the ARN, and shapes the request/response in the middle of the processor), so each addition grows a class whose job is reranking statements, not talking to providers.
• Users with proprietary or domain-specific ranking logic have no supported option other than forking, monkey-patching the processor, or reranking outside the retrieval pipeline — which loses the entity-context enrichment and max_statements truncation the processors apply.
• RerankerMixin already exists but isn't the extension point for this path. retrieval/post_processors/reranker_mixin.py defines batch_size and rerank_pairs(pairs, batch_size), and SentenceReranker/BGEReranker implement it — but the only consumer is the deprecated retrievers/deprecated/rerank_beam_search.py. The statement/topic processors never consult it, so a user implementing RerankerMixin today does not get picked up by the current pipeline.
• The fallback chain work in [FEATURE] Add reranker fallback chains to statement reranking #406 multiplies the cost. Chains (reranker=['bedrock', 'tfidf']) plus reranker_fallback_policy make the set of names something users will want to compose, and each element of a chain is still limited to a name the toolkit ships.

Proposed solution

Introduce a first-class reranker interface and accept instances of it wherever a reranker name is accepted today.

• Define a scorer-shaped protocol/ABC that matches what the processors actually need, e.g.:

  class StatementReranker(ABC):
      @abstractmethod
      def score_values(
          self,
          values: List[str],
          query: QueryBundle,
          entity_contexts: EntityContexts,
      ) -> Dict[str, float]: ...

• Returning a {value: score} map keeps it drop-in compatible with the existing scored_values contract (rerank_statements.py:270-295).
• Reconcile with the existing RerankerMixin: either extend/adapt it so rerank_pairs-style implementations are usable from the processors, or clearly scope the two (node post-processor vs. statement scorer) and document which to implement. Decide whether the deprecated beam-search consumer keeps its own path.
• Allow ProcessorArgs.reranker / topic_reranker to accept an instance (or list of instances, for [FEATURE] Add reranker fallback chains to statement reranking #406-style chains) in addition to the current string names — built-in names resolve to built-in implementations through a small registry so the dispatch chain disappears.
• Extract the built-in strategies (tfidf, model, bedrock) into implementations of the new interface so the built-ins and user-supplied rerankers travel the same code path — this is also what proves the interface is sufficient.
• Ensure custom rerankers participate in the [FEATURE] Add reranker fallback chains to statement reranking #406 fallback chain and reranker_fallback_policy on the same terms as built-ins.

Considerations / open questions

• Should the interface cover statement and topic reranking with one type, or two? Both currently need query + values → scores, but topics carry pre-existing scores.
• Does entity reranking (currently always TF-IDF, per [FEATURE] Add reranker fallback chains to statement reranking #406's "Decisions to review" §1) come in scope, or stay out for now?
• Naming/registry: is a plain Dict[str, Callable] registry enough, or should custom names be registrable so config-driven setups can reference them by string?
• Backward compatibility: all existing string values must keep working unchanged.

Acceptance criteria

• A documented public interface for custom rerankers, exported from a stable module path.
• reranker / topic_reranker accept user-supplied instances alongside the existing string names.
• Built-in tfidf / model / bedrock reimplemented against the interface; no if/elif strategy dispatch left in rerank_statements.py / rerank_topics.py.
• Custom rerankers work inside fallback chains and with reranker_fallback_policy.
• Existing string-based configuration continues to work with no changes (regression tests).
• Unit tests covering a custom reranker end-to-end through the retrieval pipeline.
• Docs: a "writing a custom reranker" example in the public documentation.

Related

• PR [FEATURE] Add reranker fallback chains to statement reranking #406 — reranker fallback chains (review comment that prompted this)
• Issue [FEATURE] Add toolkit-native reranker chains with configurable fallback behavior #403

Alternatives considered

No response