From 41a2a6251157a33ac923cccdedd712d424f36c25 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 3 Jun 2025 11:55:28 -0700 Subject: [PATCH] fix: bedrock kb test --- .../test_bedrock_knowledgebase_hook.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 1947a5f785b..96efbea7adc 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -52,7 +52,7 @@ def setup_vector_store_registry(): @pytest.mark.asyncio async def test_basic_bedrock_knowledgebase_retrieval(setup_vector_store_registry): - bedrock_knowledgebase_hook = BedrockVectorStore() + bedrock_knowledgebase_hook = BedrockVectorStore(aws_region_name="us-west-2") response = await bedrock_knowledgebase_hook.make_bedrock_kb_retrieve_request( knowledge_base_id="T37J8R4WTM", query="what is litellm?", @@ -116,7 +116,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call(setup_vecto # Init client litellm._turn_on_debug() async_client = AsyncHTTPHandler() - litellm.callbacks = [BedrockVectorStore()] + litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")] response = await litellm.acompletion( model="anthropic/claude-3-5-haiku-latest", messages=[{"role": "user", "content": "what is litellm?"}], @@ -133,7 +133,7 @@ async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registr """ Tests that knowledge base content is correctly passed to the OpenAI API call """ - litellm.callbacks = [BedrockVectorStore()] + litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")] litellm.set_verbose = True from openai import AsyncOpenAI @@ -181,7 +181,7 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai(setup_vecto This is the OpenAI format """ - litellm.callbacks = [BedrockVectorStore()] + litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")] litellm.set_verbose = True from openai import AsyncOpenAI @@ -229,7 +229,7 @@ async def test_logging_with_knowledge_base_hook(setup_vector_store_registry): Test that the knowledge base request was logged in standard logging payload """ test_custom_logger = TestCustomLogger() - litellm.callbacks = [BedrockVectorStore(), test_custom_logger] + litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2"), test_custom_logger] litellm.set_verbose = True await litellm.acompletion( model="gpt-4", @@ -282,7 +282,7 @@ async def test_logging_with_knowledge_base_hook_no_vector_store_registry(setup_v Test that the knowledge base request was logged in standard logging payload """ test_custom_logger = TestCustomLogger() - litellm.callbacks = [BedrockVectorStore(), test_custom_logger] + litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2"), test_custom_logger] litellm.vector_store_registry = None await litellm.acompletion( model="gpt-4",