From 8b776dd649b926ea186c08e17d2ee7d46ec4bb34 Mon Sep 17 00:00:00 2001 From: Mohit Arora Date: Wed, 15 Jul 2026 20:36:23 -0400 Subject: [PATCH] Add reranker fallback chains to statement reranking --- .../traversal-based-search-configuration.mdx | 47 ++- .../retrieval/processors/processor_args.py | 13 +- .../retrieval/processors/rerank_statements.py | 151 +++++--- .../retrieval/processors/reranker_chain.py | 69 ++++ .../processors/test_rerank_statements.py | 359 +++++++++++++++++- .../processors/test_reranker_chain.py | 49 +++ 6 files changed, 630 insertions(+), 58 deletions(-) create mode 100644 lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/reranker_chain.py create mode 100644 lexical-graph/tests/unit/retrieval/processors/test_reranker_chain.py diff --git a/docs-site/src/content/docs/lexical-graph/traversal-based-search-configuration.mdx b/docs-site/src/content/docs/lexical-graph/traversal-based-search-configuration.mdx index 54713ab30..494d51569 100644 --- a/docs-site/src/content/docs/lexical-graph/traversal-based-search-configuration.mdx +++ b/docs-site/src/content/docs/lexical-graph/traversal-based-search-configuration.mdx @@ -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: diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/processor_args.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/processor_args.py index 410c5cb9e..d95eca09b 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/processor_args.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/processor_args.py @@ -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. @@ -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) @@ -143,4 +152,4 @@ def __repr__(self): Returns: str: A string representation of the object in dictionary format. """ - return str(self.to_dict()) \ No newline at end of file + return str(self.to_dict()) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_statements.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_statements.py index 562600767..85065e02c 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_statements.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_statements.py @@ -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): @@ -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): """ @@ -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 @@ -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 = [] @@ -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() @@ -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: @@ -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) - - diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/reranker_chain.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/reranker_chain.py new file mode 100644 index 000000000..69b20bff1 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/reranker_chain.py @@ -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 diff --git a/lexical-graph/tests/unit/retrieval/processors/test_rerank_statements.py b/lexical-graph/tests/unit/retrieval/processors/test_rerank_statements.py index 9c8706e4c..e03a9a8ee 100644 --- a/lexical-graph/tests/unit/retrieval/processors/test_rerank_statements.py +++ b/lexical-graph/tests/unit/retrieval/processors/test_rerank_statements.py @@ -3,7 +3,8 @@ """Tests for retrieval/processors/rerank_statements.""" -from unittest.mock import MagicMock, Mock, patch +import logging +from unittest.mock import Mock, patch import pytest from llama_index.core.schema import QueryBundle @@ -24,6 +25,7 @@ RerankStatements, default_reranking_source_metadata_fn, ) +from graphrag_toolkit.lexical_graph.retrieval.processors.reranker_chain import KNOWN_RERANKERS def _collection_with_statements(): @@ -146,6 +148,21 @@ def test_none_reranker_returns_results_unchanged(self): result = processor._process_results(collection, QueryBundle('q')) assert result is collection + def test_none_reranker_short_circuits_before_formatting_values(self): + source_metadata_fn = Mock() + processor = RerankStatements( + ProcessorArgs( + reranker='none', + debug_results=[], + reranking_source_metadata_fn=source_metadata_fn, + ), + FilterConfig(), + ) + collection = _collection_with_statements() + result = processor._process_results(collection, QueryBundle('q')) + assert result is collection + source_metadata_fn.assert_not_called() + def test_no_reranker_value_returns_results_unchanged(self): processor = RerankStatements( ProcessorArgs(reranker=None, debug_results=[]), FilterConfig(), @@ -224,3 +241,343 @@ def test_bedrock_reranker_dispatches_and_applies_scores(self): statements = result.results[0].topics[0].statements assert {s.score for s in statements} == {0.3, 0.7} assert statements[0].score >= statements[1].score + + def test_reranker_chain_uses_first_successful_reranker(self): + processor = RerankStatements( + ProcessorArgs(reranker=['bedrock', 'tfidf'], debug_results=[], max_statements=10), + FilterConfig(), + ) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = lambda values, *_a, **_kw: { + values[0]: 0.3, values[1]: 0.7, + } + result = processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_not_called() + statements = result.results[0].topics[0].statements + assert {s.score for s in statements} == {0.3, 0.7} + + def test_reranker_chain_falls_back_after_exception(self, caplog): + processor = RerankStatements( + ProcessorArgs(reranker=['bedrock', 'tfidf'], debug_results=[], max_statements=10), + FilterConfig(), + ) + failure = RuntimeError('throttled') + caplog.set_level(logging.WARNING, logger=mod.__name__) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = failure + tfidf_scorer.side_effect = lambda values, *_a, **_kw: { + values[0]: 0.4, values[1]: 0.9, + } + result = processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_called_once() + statements = result.results[0].topics[0].statements + assert statements[0].score == 0.9 + assert statements[1].score == 0.4 + warning = next( + record for record in caplog.records + if record.levelno == logging.WARNING and 'Reranking with bedrock failed' in record.getMessage() + ) + assert warning.getMessage() == 'Reranking with bedrock failed (RuntimeError); trying tfidf' + assert warning.exc_info is not None + assert warning.exc_info[1] is failure + + def test_reranker_chain_empty_score_map_is_terminal(self, caplog): + fallback_policy = Mock(return_value=True) + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'tfidf'], + reranker_fallback_policy=fallback_policy, + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + caplog.set_level(logging.ERROR, logger=mod.__name__) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.return_value = {} + result = processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_not_called() + fallback_policy.assert_not_called() + assert any( + record.levelno == logging.ERROR + and record.getMessage() + == 'Reranking with bedrock returned an empty score map; all statements will be dropped' + for record in caplog.records + ) + assert result.results == [] + + def test_object_fallback_policy_is_rejected(self): + class PolicyObject: + def should_fallback(self, reranker, *, error=None): + return True + + args = ProcessorArgs( + reranker=['bedrock', 'tfidf'], + reranker_fallback_policy=PolicyObject(), + debug_results=[], + ) + with pytest.raises(TypeError): + RerankStatements(args, FilterConfig()) + + def test_reranker_chain_can_fall_back_to_none(self, caplog): + processor = RerankStatements( + ProcessorArgs(reranker=['bedrock', 'none'], debug_results=[], max_statements=10), + FilterConfig(), + ) + collection = _collection_with_statements() + caplog.set_level(logging.WARNING, logger=mod.__name__) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = RuntimeError('throttled') + result = processor._process_results(collection, QueryBundle('q')) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_not_called() + assert any( + record.levelno == logging.WARNING + and record.getMessage() == 'Reranking with bedrock failed (RuntimeError); trying none' + for record in caplog.records + ) + assert result is collection + + def test_policy_denied_failure_does_not_reach_none(self): + fallback_policy = Mock(return_value=False) + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'none'], + reranker_fallback_policy=fallback_policy, + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + failure = RuntimeError('not eligible for fallback') + with patch.object(processor, '_score_values_with_bedrock', side_effect=failure): + with pytest.raises(RuntimeError) as raised: + processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + assert raised.value is failure + fallback_policy.assert_called_once_with('bedrock', error=failure) + + def test_terminal_reranker_exception_propagates(self): + policy_calls = [] + + def fallback_only_from_bedrock(reranker, *, error=None): + policy_calls.append((reranker, error)) + return reranker == 'bedrock' + + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'tfidf'], + reranker_fallback_policy=fallback_only_from_bedrock, + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + bedrock_failure = RuntimeError('retryable failure') + terminal_failure = RuntimeError('terminal failure') + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = bedrock_failure + tfidf_scorer.side_effect = terminal_failure + with pytest.raises(RuntimeError) as raised: + processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + assert raised.value is terminal_failure + assert policy_calls == [('bedrock', bedrock_failure)] + + def test_terminal_exception_propagates_by_default(self): + processor = RerankStatements( + ProcessorArgs(reranker=['bedrock', 'tfidf'], debug_results=[], max_statements=10), + FilterConfig(), + ) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = RuntimeError('retryable failure') + tfidf_scorer.side_effect = RuntimeError('terminal failure') + with pytest.raises(RuntimeError, match='terminal failure'): + processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_called_once() + + def test_custom_fallback_policy_can_stop_fallback(self): + calls = [] + + def fallback_policy(reranker, *, error=None): + calls.append(reranker) + return False + + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'tfidf'], + reranker_fallback_policy=fallback_policy, + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = RuntimeError('do not fallback') + with pytest.raises(RuntimeError): + processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + assert calls == ['bedrock'] + tfidf_scorer.assert_not_called() + + def test_three_reranker_chain_stops_at_middle_success(self): + processor = RerankStatements( + ProcessorArgs(reranker=['bedrock', 'tfidf', 'model'], debug_results=[], max_statements=10), + FilterConfig(), + ) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer, \ + patch.object(processor, '_score_values') as model_scorer: + bedrock_scorer.side_effect = RuntimeError('throttled') + tfidf_scorer.side_effect = lambda values, *_a, **_kw: {values[0]: 0.5, values[1]: 0.8} + result = processor._process_results(_collection_with_statements(), QueryBundle('q')) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_called_once() + model_scorer.assert_not_called() + statements = result.results[0].topics[0].statements + assert statements[0].score == 0.8 + + def test_three_reranker_chain_reaches_third_after_two_failures(self): + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'tfidf', 'model'], + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer, \ + patch.object(processor, '_score_values') as model_scorer: + bedrock_scorer.side_effect = RuntimeError('first failure') + tfidf_scorer.side_effect = RuntimeError('second failure') + model_scorer.side_effect = lambda values, *_a, **_kw: { + values[0]: 0.2, values[1]: 0.95, + } + result = processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + bedrock_scorer.assert_called_once() + tfidf_scorer.assert_called_once() + model_scorer.assert_called_once() + statements = result.results[0].topics[0].statements + assert [statement.score for statement in statements] == [0.95, 0.2] + + def test_all_rerankers_fail_with_last_exception(self): + policy_calls = [] + + def fallback_policy(reranker, *, error=None): + policy_calls.append((reranker, error)) + return True + + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'tfidf', 'model'], + reranker_fallback_policy=fallback_policy, + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + bedrock_failure = RuntimeError('first failure') + tfidf_failure = RuntimeError('second failure') + model_failure = RuntimeError('last failure') + with patch.object(processor, '_score_values_with_bedrock', side_effect=bedrock_failure), \ + patch.object(processor, '_score_values_with_tfidf', side_effect=tfidf_failure), \ + patch.object(processor, '_score_values', side_effect=model_failure): + with pytest.raises(RuntimeError) as raised: + processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + assert raised.value is model_failure + assert policy_calls == [ + ('bedrock', bedrock_failure), + ('tfidf', tfidf_failure), + ] + + def test_successful_reranking_logs_only_at_debug(self, caplog): + processor = RerankStatements( + ProcessorArgs(reranker=['bedrock', 'tfidf'], debug_results=[], max_statements=10), + FilterConfig(), + ) + caplog.set_level(logging.DEBUG, logger=mod.__name__) + with patch.object(processor, '_score_values_with_bedrock', return_value={'value': 1.0}), \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + result = processor._score_values_with_chain( + ['value'], QueryBundle('q'), _entity_contexts(), + ) + assert result == {'value': 1.0} + tfidf_scorer.assert_not_called() + success_records = [ + record for record in caplog.records + if record.getMessage() == 'Reranking succeeded with bedrock' + ] + assert len(success_records) == 1 + assert success_records[0].levelno == logging.DEBUG + assert not any( + record.levelno == logging.INFO and 'Reranking succeeded' in record.getMessage() + for record in caplog.records + ) + + def test_policy_exception_propagates(self): + def broken_policy(reranker, *, error=None): + raise KeyError('broken policy') + + processor = RerankStatements( + ProcessorArgs( + reranker=['bedrock', 'tfidf'], + reranker_fallback_policy=broken_policy, + debug_results=[], + max_statements=10, + ), + FilterConfig(), + ) + with patch.object(processor, '_score_values_with_bedrock') as bedrock_scorer, \ + patch.object(processor, '_score_values_with_tfidf') as tfidf_scorer: + bedrock_scorer.side_effect = RuntimeError('throttled') + with pytest.raises(KeyError): + processor._process_results( + _collection_with_statements(), QueryBundle('q'), + ) + tfidf_scorer.assert_not_called() + + +class TestKnownRerankerCoverage: + @pytest.mark.parametrize('reranker', [r for r in KNOWN_RERANKERS if r != 'none']) + def test_every_known_reranker_dispatches_to_a_scorer(self, reranker): + # Drift guard: a name added to KNOWN_RERANKERS without a matching entry + # in _score_values_with_chain's scorer registry passes chain validation + # but silently skips reranking at query time. + processor = RerankStatements( + ProcessorArgs(reranker=[reranker], debug_results=[], max_statements=10), + FilterConfig(), + ) + with patch.object(processor, '_score_values', return_value={'v': 1.0}), \ + patch.object(processor, '_score_values_with_tfidf', return_value={'v': 1.0}), \ + patch.object(processor, '_score_values_with_bedrock', return_value={'v': 1.0}): + scored = processor._score_values_with_chain( + ['v'], QueryBundle('q'), EntityContexts(contexts=[], keywords=[]), + ) + assert scored == {'v': 1.0}, f'no scorer wired up for reranker "{reranker}"' diff --git a/lexical-graph/tests/unit/retrieval/processors/test_reranker_chain.py b/lexical-graph/tests/unit/retrieval/processors/test_reranker_chain.py new file mode 100644 index 000000000..78c3b7a11 --- /dev/null +++ b/lexical-graph/tests/unit/retrieval/processors/test_reranker_chain.py @@ -0,0 +1,49 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for retrieval/processors/reranker_chain.""" + +import pytest + +from graphrag_toolkit.lexical_graph.retrieval.processors.reranker_chain import normalize_reranker_chain + + +class TestNormalizeRerankerChain: + @pytest.mark.parametrize( + 'reranker, expected', + [ + (None, []), + (False, []), + (0, []), + ({}, []), + ([], []), + (' TFIDF ', ['tfidf']), + (' ', []), + (['Bedrock', ' tfidf'], ['bedrock', 'tfidf']), + (['', 'tfidf'], ['tfidf']), + ([' '], []), + (['none'], ['none']), + ], + ) + def test_normalizes_supported_values(self, reranker, expected): + assert normalize_reranker_chain(reranker) == expected + + @pytest.mark.parametrize( + 'reranker, error_type, match', + [ + (('bedrock', 'tfidf'), TypeError, 'string or list'), + (42, TypeError, 'string or list'), + ({'reranker': 'tfidf'}, TypeError, 'string or list'), + ([42, 'tfidf'], ValueError, 'must be strings'), + ([None, 'tfidf'], ValueError, 'must be strings'), + (['bedrock', 'none', 'tfidf'], ValueError, 'none can only be used'), + (['bedrock', 'typo'], ValueError, 'Unknown reranker'), + ], + ) + def test_rejects_invalid_chain_values(self, reranker, error_type, match): + with pytest.raises(error_type, match=match): + normalize_reranker_chain(reranker) + + def test_unknown_single_string_is_preserved_for_legacy_behaviour(self): + # Legacy string form is not validated against known rerankers. + assert normalize_reranker_chain('something-weird') == ['something-weird']