Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -174,20 +174,55 @@ ___

### Reranking strategy

Traversal-based search incorporates reranking at two key points during the retrieval process:
Traversal-based search incorporates two independent reranking stages:

- When generating entity network contexts, both entities and entity networks are reranked
- Before finalizing search results, the complete set of statements undergoes reranking
- When generating entity network contexts, entities are always reranked with TF-IDF
- Before finalizing search results, the complete set of statements can be reranked with a configured strategy

Reranking is managed through a single parameter:
The `reranker` and `reranker_fallback_policy` parameters control only the final statement reranking stage. They do not change entity reranking.

##### `reranker`

Parameters options:

- `model`: Uses a LlamaIndex-based `SentenceReranker` to rerank all statements in the result set
- `tfidf` (default): Applies a term frequency-inverse document frequency measure to rank statements
- `None`: Disables the reranking feature completely
- `bedrock`: Uses the configured Amazon Bedrock reranking model to rerank all statements in the result set
- `tfidf` (default): Applies a term frequency-inverse document frequency measure to rank statements
- `None` or `none`: Disables final statement reranking; entity reranking still uses TF-IDF

You can also pass an ordered list of rerankers to try them in sequence. This is useful when you want to prefer a model-based reranker, but fall back to the built-in `tfidf` reranker if the first reranker raises an exception.

```python
query_engine = LexicalGraphQueryEngine.for_traversal_based_search(
graph_store,
vector_store,
reranker=['bedrock', 'tfidf']
)
```

Reranker names in a chain are validated before reranking runs; unknown names raise a `ValueError` on the first query. By default, chained rerankers fall back to the next strategy on any exception. To restrict this to specific failures, provide a `reranker_fallback_policy` function; the reranker falls back only when the function returns `True`. An exception from the final reranker in a chain always propagates to the caller, just as with a single reranker.

Use `none` by itself to disable statement reranking, or as the final item in a chain to return unreranked statements when an earlier reranker fails. The preceding failure must be allowed by `reranker_fallback_policy` before the chain can reach `none`; otherwise, the original exception propagates. Once reached, `none` performs no scoring and cannot raise.

```python
def fallback_on_throttling(reranker, *, error=None):
response = getattr(error, 'response', None)
if not isinstance(response, dict):
return False
code = response.get('Error', {}).get('Code')
return code in {'ThrottlingException', 'TooManyRequestsException'}

query_engine = LexicalGraphQueryEngine.for_traversal_based_search(
graph_store,
vector_store,
reranker=['bedrock', 'tfidf'],
reranker_fallback_policy=fallback_on_throttling
)
```

Fallback is decided per query, so a reranker that fails on one query is attempted again on the next; transient conditions such as throttling recover automatically once they clear server-side.

A reranker that returns an empty score map has completed without raising, so the result is terminal: the empty map is logged as an error, no fallback is attempted, and all statements are dropped.

The tfidf-based option is significantly faster than the model-based approach. To use the model reranker, you must first install the following additional dependencies:

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,15 @@ class ProcessorArgs():
derive_subqueries (bool): Specifies if subqueries should be derived
from the main query.
debug_results (list): A list for storing intermediate debug results.
reranker (str): Defines the reranking strategy to be employed.
reranker (str or list): Defines the final statement reranking strategy
or ordered fallback strategies. ``None`` or ``'none'`` disables
statement reranking only; entity reranking independently uses
TF-IDF.
reranker_fallback_policy (Optional[Callable]): Optional function for
statement reranker chains, called as
``fallback_policy(reranker, *, error=None)``, that decides whether
a reranker exception should fall back to the next strategy,
including ``'none'``. Defaults to falling back on any exception.
max_statements (int): The maximum number of statements to process.
max_search_results (int): The maximum number of search results to
retrieve.
Expand Down Expand Up @@ -61,6 +69,7 @@ def __init__(self, **kwargs):
self.derive_subqueries = kwargs.get('derive_subqueries', False)
self.debug_results = kwargs.get('debug_results', [])
self.reranker = kwargs.get('reranker', 'tfidf')
self.reranker_fallback_policy = kwargs.get('reranker_fallback_policy', None)
self.disaggregate_results = kwargs.get('disaggregate_results', False)
self.max_statements = kwargs.get('max_statements', 200)
self.max_search_results = kwargs.get('max_search_results', 5)
Expand Down Expand Up @@ -143,4 +152,4 @@ def __repr__(self):
Returns:
str: A string representation of the object in dictionary format.
"""
return str(self.to_dict())
return str(self.to_dict())
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
from llama_index.core.schema import QueryBundle, NodeWithScore, TextNode
from llama_index.core.node_parser import TokenTextSplitter

from graphrag_toolkit.lexical_graph.retrieval.processors.reranker_chain import normalize_reranker_chain

logger = logging.getLogger(__name__)

def default_reranking_source_metadata_fn(source:Source):
Expand Down Expand Up @@ -52,13 +54,25 @@ def format_value(s):
return ', '.join([format_value(v) for v in source.metadata.values()])

class RerankStatements(ProcessorBase):
"""
"""Rerank the final statements in a search result collection.

``reranker`` may specify one strategy or an ordered fallback chain. A
fallback policy is consulted after a scorer raises and before the next
strategy is attempted. ``None`` or ``'none'`` disables statement reranking;
entity reranking is an independent TF-IDF stage.
"""
def __init__(self, args:ProcessorArgs, filter_config:FilterConfig, reranking_model=None):
self.reranking_model = reranking_model or GraphRAGConfig.reranking_model
super().__init__(args, filter_config)
self.reranking_source_metadata_fn = self.args.reranking_source_metadata_fn or default_reranking_source_metadata_fn
self.reranker_chain = normalize_reranker_chain(self.args.reranker)
fallback_policy = getattr(self.args, 'reranker_fallback_policy', None)
if fallback_policy is not None and not callable(fallback_policy):
raise TypeError(
f'reranker_fallback_policy must be a callable accepting '
f'(reranker, *, error=None), got {type(fallback_policy).__name__}'
)
self.reranker_fallback_policy = fallback_policy

def _score_values_with_tfidf(self, values:List[str], query:QueryBundle, entity_contexts:EntityContexts):
"""
Expand Down Expand Up @@ -205,13 +219,80 @@ def rerank_text(text_query, text_sources, num_results, model_package_arn):
for result in results
}

def _should_fallback(self, reranker:str, *, error=None) -> bool:
if self.reranker_fallback_policy is None:
return True

try:
return bool(self.reranker_fallback_policy(reranker, error=error))
except Exception:
logger.error(
f'reranker_fallback_policy raised an exception for reranker '
f'"{reranker}"; re-raising',
exc_info=True
)
raise

def _score_values_with_chain(self, values:List[str], query:QueryBundle, entity_contexts:EntityContexts):
"""Score values with the first successful statement reranker.

Raised exceptions may advance to the next strategy when the fallback
policy permits. A returned empty score map is a successful, terminal
response: it is logged at error level and returned without consulting
the fallback policy. Once reached, ``'none'`` performs no scoring and
returns ``None``.
"""
reranker_chain = self.reranker_chain
if not reranker_chain:
return None

reranker_registry = {
'model': self._score_values,
'tfidf': self._score_values_with_tfidf,
'bedrock': self._score_values_with_bedrock
}

for i, reranker in enumerate(reranker_chain):
has_next = i < len(reranker_chain) - 1

if reranker == 'none':
return None

if reranker not in reranker_registry:
# Only the unvalidated legacy single-string form can reach here.
logger.warning(f'Unknown reranker "{reranker}"; skipping reranking')
return None

try:
scored_values = reranker_registry[reranker](values, query, entity_contexts)
except Exception as e:
if not has_next:
raise
if not self._should_fallback(reranker, error=e):
raise
next_reranker = reranker_chain[i + 1]
logger.warning(
f'Reranking with {reranker} failed ({type(e).__name__}); '
f'trying {next_reranker}',
exc_info=True
)
continue

if not scored_values:
logger.error(f'Reranking with {reranker} returned an empty score map; all statements will be dropped')
else:
logger.debug(f'Reranking succeeded with {reranker}')

return scored_values

def _process_results(self, search_results:SearchResultCollection, query:QueryBundle) -> SearchResultCollection:
"""
Processes search results by reranking statements within each topic based on their relevance scores.
Rerank statements within each topic based on their relevance scores.

This method processes a set of search results, and if a reranking approach is specified,
it calculates scores for the statements within topics using the supplied query and entities.
It then reranks statements based on the computed scores and updates the search results accordingly.
The configured chain applies only to this final statement stage. Entity
reranking occurs independently with TF-IDF. A score map is terminal even
when empty; an empty map therefore drops every statement without trying
another fallback.

Args:
search_results (SearchResultCollection): A collection of search results that contain topics
Expand All @@ -224,9 +305,14 @@ def _process_results(self, search_results:SearchResultCollection, query:QueryBun
topic are reranked based on the relevance scores computed by the specified reranking method.

Raises:
Any exception raised during the reranking process will propagate to the caller.
Exceptions raised by the final (or only) reranker propagate to the
caller. Earlier rerankers fall back when the configured policy
permits it (by default, on any exception). Reaching a final
``'none'`` also requires the policy to permit the preceding failure;
``'none'`` itself performs no scoring and cannot raise.
"""
if not self.args.reranker or self.args.reranker.lower() == 'none':
reranker_chain = self.reranker_chain
if not reranker_chain or reranker_chain[0] == 'none':
return search_results

values_to_score = []
Expand All @@ -245,16 +331,13 @@ def _process_results(self, search_results:SearchResultCollection, query:QueryBun

start = time.time()

scored_values = None

reranker = self.args.reranker.lower()
if reranker == 'model':
scored_values = self._score_values(values_to_score, query, search_results.entity_contexts)
elif reranker == 'tfidf':
scored_values = self._score_values_with_tfidf(values_to_score, query, search_results.entity_contexts)
elif reranker == 'bedrock':
scored_values = self._score_values_with_bedrock(values_to_score, query, search_results.entity_contexts)
else:
scored_values = self._score_values_with_chain(
values_to_score,
query,
search_results.entity_contexts
)

if scored_values is None:
return search_results

end = time.time()
Expand All @@ -268,21 +351,7 @@ def _process_results(self, search_results:SearchResultCollection, query:QueryBun
logger.debug('Scored values:\n' + '\n--------------\n'.join([str(scored_value) for scored_value in scored_values.items()]))

def rerank_statements(topic:Topic, source_str:str):
"""
Represents a processor that reranks statements within topics based on their scores.

This class is used to process and reorder statements in a `SearchResultCollection`
object using scores associated with the statements. Statements are filtered and
sorted in descending order by their scores.

Attributes:
No attributes are defined directly for this class.

Methods:
_process_results: Processes the search results to rerank the statements based
on their scores.

"""
"""Apply scores to one topic, dropping statements without scores."""
topic_str = topic.topic
surviving_statements = []
for statement in topic.statements:
Expand All @@ -295,24 +364,8 @@ def rerank_statements(topic:Topic, source_str:str):
return topic

def rerank_search_result(index:int, search_result:SearchResult):
"""
Re-ranks search results based on specific criteria and applies modifications to their
topics using a metadata source function and a specified reranking method. This
processor works by overriding the `_process_results` method and provides the
necessary mechanisms to handle individual search results in a collection.

Methods:
_process_results: Processes a collection of search results and applies the reranking
to individual entries based on the provided query.

Args:
index (int): Position index of the search result to rerank within the collection.
search_result (SearchResult): Individual search result object containing the
data to be reranked.
"""
"""Apply statement reranking to one search result."""
source_str = self.reranking_source_metadata_fn(search_result.source)
return self._apply_to_topics(search_result, rerank_statements, source_str=source_str)

return self._apply_to_search_results(search_results, rerank_search_result)


Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

"""Statement reranker chain configuration and fallback policies.

A reranker chain is an ordered list of reranking strategies. Processors such as
``RerankStatements`` try each strategy in turn, consulting a fallback policy to
decide whether a failure should be recovered by moving to the next strategy.
Entity reranking is a separate TF-IDF stage and is not controlled by this chain.

A fallback policy is a callable ``policy(reranker, *, error=None) -> bool``,
consulted when a strategy raises an exception and a next strategy remains.
When no policy is configured, the chain falls back on any exception. An
exception from the final (or only) strategy always propagates. Before any next
entry is reached, including ``'none'``, the preceding failure must be permitted
by the fallback policy. Once reached, ``'none'`` performs no scoring and cannot
raise. Fallback is per-query and stateless: a strategy that failed on one query
is attempted again on the next, so transient conditions such as throttling
recover automatically once they clear server-side.
"""

KNOWN_RERANKERS = (
'model',
'tfidf',
'bedrock',
'none',
)


def normalize_reranker_chain(reranker):
"""Normalize a statement reranker configuration into strategy names.

Any falsy value returns an empty chain before type validation. Strings are
stripped and lowercased. Lists are stripped, lowercased, and emptied of
blank entries; their names and placement of ``'none'`` are validated. The
legacy single-string form is not validated against ``KNOWN_RERANKERS``,
preserving the historical silently-skip behaviour for unknown names.

Raises:
TypeError: If a truthy value is neither a string nor a list.
ValueError: If a list contains a non-string entry, an unknown strategy,
or ``'none'`` anywhere except the final position.
"""
if not reranker:
return []

if isinstance(reranker, str):
return [reranker.strip().lower()] if reranker.strip() else []

if not isinstance(reranker, list):
raise TypeError(
f'reranker must be a string or list of strings, got {type(reranker).__name__}'
)

for r in reranker:
if not isinstance(r, str):
raise ValueError(
f'reranker list entries must be strings, got {type(r).__name__}: {r!r}'
)
chain = [r.strip().lower() for r in reranker if r.strip()]
unknown = [r for r in chain if r not in KNOWN_RERANKERS]
if unknown:
raise ValueError(
f'Unknown reranker(s) in chain: {unknown}. '
f'Expected values from: {list(KNOWN_RERANKERS)}'
)
if 'none' in chain[:-1]:
raise ValueError('none can only be used alone or as the final fallback in a reranker chain')
return chain
Loading