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 54713ab3..67181389 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 @@ -20,6 +20,7 @@ title: Traversal-Based Search Configuration - [When to use different retrievers](#when-to-use-different-retrievers) - [Reranking strategy](#reranking-strategy) - [reranker](#reranker) + - [bedrock_reranker_client_config](#bedrock_reranker_client_config) - [Choosing a reranker strategy](#choosing-a-reranker-strategy) - [Troubleshooting reranking results](#troubleshooting-reranking-results) - [Graph and vector search parameters](#graph-and-vector-search-parameters) @@ -186,6 +187,7 @@ Reranking is managed through a single parameter: Parameters options: - `model`: Uses a LlamaIndex-based `SentenceReranker` to rerank all statements in the result set + - `bedrock`: Uses AWS Bedrock API for reranking. Requires AWS credentials and can be configured with `bedrock_reranker_client_config` - `tfidf` (default): Applies a term frequency-inverse document frequency measure to rank statements - `None`: Disables the reranking feature completely @@ -205,6 +207,25 @@ query_engine = LexicalGraphQueryEngine.for_traversal_based_search( ) ``` +##### `bedrock_reranker_client_config` + +An optional mapping of arguments for [`botocore.config.Config`](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html). It applies only to the `bedrock-agent-runtime` clients created for statement- and topic-level Bedrock reranking. Each processor creates its own client from the mapping; the shared boto3 session and other AWS clients are not modified. The default is `None`, which passes no explicit `Config`. + +```python +args = ProcessorArgs( + reranker='bedrock', + topic_reranker='bedrock', + bedrock_reranker_client_config={ + 'connect_timeout': 2, + 'read_timeout': 2, + 'retries': { + 'total_max_attempts': 1, + 'mode': 'standard', + }, + }, +) +``` + #### Choosing a reranker strategy The tfidf reranker option provides a fast, cost-effective, and generally effective solution for most use cases. However, if you find that the results don't meet your requirements, consider switching to the model reranker. Be aware that while model may provide different results, it operates significantly slower than tfidf and doesn't guarantee improved outcomes. 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 410c5cb9..87c6fd9c 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 @@ -1,7 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Dict, Any +from typing import Any, Dict, Mapping class ProcessorArgs(): """ProcessorArgs is a configuration class for managing various options and @@ -22,6 +22,8 @@ class ProcessorArgs(): from the main query. debug_results (list): A list for storing intermediate debug results. reranker (str): Defines the reranking strategy to be employed. + bedrock_reranker_client_config (Mapping[str, Any] | None): Optional + botocore client configuration used only for Bedrock reranking. max_statements (int): The maximum number of statements to process. max_search_results (int): The maximum number of search results to retrieve. @@ -53,6 +55,8 @@ class ProcessorArgs(): ecs_max_entities_per_context (int): Restriction on how many entities are considered per context. """ + bedrock_reranker_client_config: Mapping[str, Any] | None + def __init__(self, **kwargs): self.expand_entities = kwargs.get('expand_entities', True) @@ -61,6 +65,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.bedrock_reranker_client_config = kwargs.get('bedrock_reranker_client_config', 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 +148,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 56260076..fdb81f8f 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 @@ -3,10 +3,10 @@ import logging import time -import boto3 import json from typing import List, Dict from dateutil.parser import parse +from botocore.config import Config from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph import GraphRAGConfig @@ -154,9 +154,13 @@ def _score_values_with_bedrock(self, values:List[str], query:QueryBundle, entity else f'{query.query_str} (keywords: {extras})' ) - region = boto3.Session().region_name + session = GraphRAGConfig.session + region = session.region_name + client_kwargs = {'region_name': region} + if self.args.bedrock_reranker_client_config is not None: + client_kwargs['config'] = Config(**self.args.bedrock_reranker_client_config) - bedrock_agent_runtime = boto3.client('bedrock-agent-runtime', region_name=region) + bedrock_agent_runtime = session.client('bedrock-agent-runtime', **client_kwargs) modelId = GraphRAGConfig.bedrock_reranking_model model_package_arn = f"arn:aws:bedrock:{region}::foundation-model/{modelId}" @@ -315,4 +319,3 @@ def rerank_search_result(index:int, search_result:SearchResult): return self._apply_to_search_results(search_results, rerank_search_result) - diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_topics.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_topics.py index a6db9b27..a3f60746 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_topics.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/retrieval/processors/rerank_topics.py @@ -20,7 +20,7 @@ import time from typing import List, Dict -import boto3 +from botocore.config import Config from graphrag_toolkit.lexical_graph import GraphRAGConfig from graphrag_toolkit.lexical_graph.retrieval.processors.processor_base import ProcessorBase @@ -50,8 +50,12 @@ def _score_with_tfidf(self, texts: List[str], query: QueryBundle) -> List[float] return [scored.get(t, 0.0) for t in texts] def _score_with_bedrock(self, texts: List[str], query: QueryBundle) -> List[float]: - region = boto3.Session().region_name or getattr(GraphRAGConfig, 'aws_region', None) or 'us-east-1' - client = boto3.client('bedrock-agent-runtime', region_name=region) + session = GraphRAGConfig.session + region = session.region_name or GraphRAGConfig.aws_region or 'us-east-1' + client_kwargs = {'region_name': region} + if self.args.bedrock_reranker_client_config is not None: + client_kwargs['config'] = Config(**self.args.bedrock_reranker_client_config) + client = session.client('bedrock-agent-runtime', **client_kwargs) model_arn = f'arn:aws:bedrock:{region}::foundation-model/{GraphRAGConfig.bedrock_reranking_model}' sources = [{'type': 'INLINE', 'inlineDocumentSource': {'type': 'TEXT', 'textDocument': {'text': t}}} for t in texts] 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 9c8706e4..b074e1d2 100644 --- a/lexical-graph/tests/unit/retrieval/processors/test_rerank_statements.py +++ b/lexical-graph/tests/unit/retrieval/processors/test_rerank_statements.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, Mock, patch import pytest +from botocore.config import Config from llama_index.core.schema import QueryBundle from graphrag_toolkit.lexical_graph.metadata import FilterConfig @@ -119,22 +120,98 @@ def test_builds_score_map_from_reranker(self): class TestScoreValuesWithBedrock: - def test_maps_results_back_to_values(self): + def test_configures_session_client_and_preserves_request(self): + client_config = { + 'connect_timeout': 2, + 'read_timeout': 3, + 'retries': {'total_max_attempts': 1, 'mode': 'standard'}, + } processor = RerankStatements( - ProcessorArgs(reranker='bedrock', debug_results=[], max_statements=5), + ProcessorArgs( + reranker='bedrock', + debug_results=[], + max_statements=5, + bedrock_reranker_client_config=client_config, + ), FilterConfig(), ) - with patch.object(mod, 'boto3') as boto3_mod, \ - patch.object(mod, 'GraphRAGConfig') as config: - boto3_mod.Session.return_value.region_name = 'us-east-1' + session = MagicMock() + session.region_name = 'us-east-1' + client = session.client.return_value + client.rerank.return_value = { + 'results': [{'index': 0, 'relevanceScore': 0.7}], + } + + with patch.object(mod, 'GraphRAGConfig') as config: + config.session = session config.bedrock_reranking_model = 'model-x' - boto3_mod.client.return_value.rerank.return_value = { - 'results': [{'index': 0, 'relevanceScore': 0.7}], - } out = processor._score_values_with_bedrock( ['a', 'b'], QueryBundle('q'), _entity_contexts(), ) + assert out == {'a': 0.7} + session.client.assert_called_once() + service_name, = session.client.call_args.args + client_kwargs = session.client.call_args.kwargs + assert service_name == 'bedrock-agent-runtime' + assert client_kwargs['region_name'] == 'us-east-1' + assert isinstance(client_kwargs['config'], Config) + assert client_kwargs['config'].connect_timeout == 2 + assert client_kwargs['config'].read_timeout == 3 + assert client_kwargs['config'].retries == { + 'total_max_attempts': 1, + 'mode': 'standard', + } + client.rerank.assert_called_once_with( + queries=[{'type': 'TEXT', 'textQuery': {'text': 'q'}}], + sources=[ + { + 'type': 'INLINE', + 'inlineDocumentSource': { + 'type': 'TEXT', + 'textDocument': {'text': 'a'}, + }, + }, + { + 'type': 'INLINE', + 'inlineDocumentSource': { + 'type': 'TEXT', + 'textDocument': {'text': 'b'}, + }, + }, + ], + rerankingConfiguration={ + 'type': 'BEDROCK_RERANKING_MODEL', + 'bedrockRerankingConfiguration': { + 'numberOfResults': 2, + 'modelConfiguration': { + 'modelArn': ( + 'arn:aws:bedrock:us-east-1::foundation-model/model-x' + ), + }, + }, + }, + ) + + def test_omitted_config_does_not_pass_botocore_config(self): + processor = RerankStatements( + ProcessorArgs(reranker='bedrock', debug_results=[], max_statements=5), + FilterConfig(), + ) + session = MagicMock() + session.region_name = 'us-east-1' + session.client.return_value.rerank.return_value = {'results': []} + + with patch.object(mod, 'GraphRAGConfig') as config: + config.session = session + config.bedrock_reranking_model = 'model-x' + processor._score_values_with_bedrock( + ['a'], QueryBundle('q'), _entity_contexts(), + ) + + session.client.assert_called_once_with( + 'bedrock-agent-runtime', region_name='us-east-1', + ) class TestProcessResults: diff --git a/lexical-graph/tests/unit/retrieval/processors/test_rerank_topics.py b/lexical-graph/tests/unit/retrieval/processors/test_rerank_topics.py index 03fce854..2636ae45 100644 --- a/lexical-graph/tests/unit/retrieval/processors/test_rerank_topics.py +++ b/lexical-graph/tests/unit/retrieval/processors/test_rerank_topics.py @@ -3,10 +3,14 @@ """Tests for the RerankTopics processor (topic-level reranking).""" +from unittest.mock import MagicMock, patch + import pytest +from botocore.config import Config from graphrag_toolkit.lexical_graph.metadata import FilterConfig from graphrag_toolkit.lexical_graph.retrieval.processors import ProcessorArgs, RerankTopics +from graphrag_toolkit.lexical_graph.retrieval.processors import rerank_topics as mod from graphrag_toolkit.lexical_graph.retrieval.model import ( SearchResultCollection, SearchResult, Topic, Statement, Source, Versioning, EntityContexts, ) @@ -47,6 +51,13 @@ def test_noop_when_topic_reranker_none(query): assert _topic_ids(out) == ['t1', 't2', 't3'] # unchanged +def test_processor_args_preserve_reranking_defaults(): + args = ProcessorArgs() + assert args.reranker == 'tfidf' + assert args.topic_reranker == 'none' + assert args.bedrock_reranker_client_config is None + + def test_prunes_to_max_topics(query): proc = RerankTopics(ProcessorArgs(topic_reranker='tfidf', max_topics=2), FilterConfig()) out = proc._process_results(_collection(), query) @@ -61,3 +72,97 @@ def test_propagates_topic_score_to_unscored_statements(query): out = proc._process_results(_collection(), query) scores = [s.score for r in out.results for t in r.topics for s in t.statements] assert any(sc and sc > 0.0 for sc in scores) + + +def test_bedrock_uses_configured_session_client_and_preserves_request(query): + proc = RerankTopics( + ProcessorArgs( + topic_reranker='bedrock', + bedrock_reranker_client_config={ + 'connect_timeout': 2, + 'read_timeout': 3, + 'retries': {'total_max_attempts': 1, 'mode': 'standard'}, + }, + ), + FilterConfig(), + ) + session = MagicMock() + session.region_name = 'us-west-2' + client = session.client.return_value + client.rerank.return_value = { + 'results': [ + {'index': 1, 'relevanceScore': 0.8}, + {'index': 0, 'relevanceScore': 0.2}, + ], + } + + with patch.object(mod, 'GraphRAGConfig') as config: + config.session = session + config.bedrock_reranking_model = 'model-y' + scores = proc._score_with_bedrock(['first', 'second'], query) + + assert scores == [0.2, 0.8] + session.client.assert_called_once() + service_name, = session.client.call_args.args + client_kwargs = session.client.call_args.kwargs + assert service_name == 'bedrock-agent-runtime' + assert client_kwargs['region_name'] == 'us-west-2' + assert isinstance(client_kwargs['config'], Config) + assert client_kwargs['config'].connect_timeout == 2 + assert client_kwargs['config'].read_timeout == 3 + assert client_kwargs['config'].retries == { + 'total_max_attempts': 1, + 'mode': 'standard', + } + client.rerank.assert_called_once_with( + queries=[{ + 'type': 'TEXT', + 'textQuery': {'text': query.query_str}, + }], + sources=[ + { + 'type': 'INLINE', + 'inlineDocumentSource': { + 'type': 'TEXT', + 'textDocument': {'text': 'first'}, + }, + }, + { + 'type': 'INLINE', + 'inlineDocumentSource': { + 'type': 'TEXT', + 'textDocument': {'text': 'second'}, + }, + }, + ], + rerankingConfiguration={ + 'type': 'BEDROCK_RERANKING_MODEL', + 'bedrockRerankingConfiguration': { + 'numberOfResults': 2, + 'modelConfiguration': { + 'modelArn': ( + 'arn:aws:bedrock:us-west-2::foundation-model/model-y' + ), + }, + }, + }, + ) + + +def test_bedrock_omitted_config_does_not_pass_botocore_config(query): + proc = RerankTopics( + ProcessorArgs(topic_reranker='bedrock'), + FilterConfig(), + ) + session = MagicMock() + session.region_name = 'us-west-2' + session.client.return_value.rerank.return_value = {'results': []} + + with patch.object(mod, 'GraphRAGConfig') as config: + config.session = session + config.bedrock_reranking_model = 'model-y' + proc._score_with_bedrock(['first'], query) + + session.client.assert_called_once_with( + 'bedrock-agent-runtime', region_name='us-west-2', + )