scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix: bedrock kb test

Browse source
This commit is contained in:
Ishaan Jaff 2025-06-03 11:55:28 -07:00
parent 50d5ee9138
commit 41a2a62511
1 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
View file

@ -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",