awslabs / awslabs/graphrag-toolkit

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

Open
#435 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
442
Forks
106
Avg merge
2d 3h
Merged PRs (30d)
41

Description

### 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 (#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 #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 #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 #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 #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 #406 — reranker fallback chains (review comment that prompted this)
- Issue #403

### Alternatives considered

_No response_

Contributor guide

Open the contributing guide

Research direction

Start with rerank_statements.py and rerank_topics.py, especially their dispatch and scored_values paths, then compare retrieval/post_processors/reranker_mixin.py with the deprecated beam-search consumer. Review ProcessorArgs and the #406 fallback behavior before deciding the interface boundary. Done means built-ins and supplied instances share the documented path, string configuration still works, fallback chains are covered by tests, and public custom-reranker documentation exists.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.