awslabs / awslabs/graphrag-toolkit
[FEATURE] Add a pluggable reranker interface so users can supply custom reranking strategies
- 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
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