awslabs / awslabs/graphrag-toolkit

[FEATURE] Add toolkit-native reranker chains with configurable fallback behavior

Open
#403 2 comments 0 reactions 1 assignee Claimed by @acarbonetto View on GitHub
enhancement
Dominant language
Python
Stars
442
Forks
106
Avg merge
2d 17h
Merged PRs (30d)
52

Description

### Package

lexical-graph

### Problem statement

I’d like to make the reranking behavior a little easier to use in production.

Right now traversal search lets you pick one reranker, like tfidf, bedrock etc.

That works, but in a real deployment it would be preferable to try a better reranker first, like Bedrock, and fall back to the built-in TF-IDF reranker if Bedrock is throttled, unavailable, misconfigured, or returns no usable scores.

### Proposed solution

Smallest possible implementation would be to keep all of the new behavior inside RerankStatements.

New shape can be something like:
```
ProcessorArgs(
reranker="bedrock,tfidf",
reranker_fallback_policy=my_policy, #optional
)
```

RerankStatements would add a private parser/helper that turns the configured value into an ordered list

Instead of the existing dispatch block, it would loop over the reranker chain
https://github.com/awslabs/graphrag-toolkit/blob/64422d797c776fdb4f68c1025622d66ed8b226f9/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_statements.py#L250-L258

something like:
```
for reranker in chain:
try:
scored_values = scorer(...)
if scored_values:
return scored_values

if not self._should_fallback(reranker, scored_values=scored_values):
return scored_values

except Exception as e:
if not has_next_reranker or not self._should_fallback(reranker, error=e):
raise
logger.warning("Reranker failed; trying fallback", exc_info=True)
```

where user code can decide whether to fall back on throttling, timeouts, empty scores, etc.

The current defaults will be maintained including support for a single reranker along with a basic default policy that mimics the behavior today.

### Alternatives considered

A more involved implementation would be to introduce a small toolkit-owned reranker abstraction inside retrieval, then adapt statement and topic reranking to use it.

Rough shape:
```
retrieval/reranking/
base.py # shared dataclasses/protocols
chain.py # ordered fallback runner
factory.py # turns "bedrock,tfidf" into reranker objects
providers/
tfidf.py
bedrock.py
sentence_transformer.py
noop.py
```

The useful contract could be like:
`score(query, texts) -> scores`

The existing rerankers can then become providers:
```
TfidfReranker
BedrockReranker
SentenceTransformerReranker
NoopReranker
```

A reranker abstraction is a little more upfront work, but will give one place to handle provider selection, fallback policy, lazy optional imports, timing, and future reranker implementations.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.