[FEATURE] : Add rerank bedrock client configuration support

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.

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())
+        return str(self.to_dict())

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)
 
-

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]

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:

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',
+    )