import os import litellm import litellm.vector_stores.main import json from typing import Optional from unittest.mock import AsyncMock, patch, Mock import pytest import litellm from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( VectorStorePreCallHook, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import ( StandardLoggingPayload, ) from litellm.types.vector_stores import ( VectorStoreSearchResponse, VectorStoreResultContent, VectorStoreSearchResult, ) class MockCustomLogger(CustomLogger): def __init__(self): self.standard_logging_payload: Optional[StandardLoggingPayload] = None self.completion_logging_payload: Optional[StandardLoggingPayload] = None super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): payload = kwargs.get("standard_logging_object") # Store the payload - completion calls have call_type='acompletion' if payload and payload.get("call_type") == "acompletion": self.completion_logging_payload = payload self.standard_logging_payload = payload pass @pytest.fixture(autouse=True) def add_aws_region_to_env(monkeypatch): monkeypatch.setenv("AWS_REGION", "us-west-2") @pytest.fixture def setup_vector_store_registry(): from litellm.vector_stores.vector_store_registry import ( VectorStoreRegistry, LiteLLM_ManagedVectorStore, ) # Init vector store registry litellm.vector_store_registry = VectorStoreRegistry( vector_stores=[ LiteLLM_ManagedVectorStore( vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock" ) ] ) @pytest.mark.asyncio async def test_vector_store_hook_routes_search_through_proxy_router( setup_vector_store_registry, ): proxy_router = Mock() proxy_router.avector_store_search = AsyncMock( return_value=VectorStoreSearchResponse( object="vector_store.search_results.page", search_query="what is litellm?", data=[ VectorStoreSearchResult( score=1.0, content=[VectorStoreResultContent(text="routed context", type="text")], ) ], ) ) logging_obj = Mock() logging_obj.model_call_details = { "litellm_params": {"metadata": {"user_api_key_team_id": "team-a"}} } with patch("litellm.proxy.proxy_server.llm_router", proxy_router): _, messages, _ = await VectorStorePreCallHook().async_get_chat_completion_prompt( model="chat-model", messages=[{"role": "user", "content": "what is litellm?"}], non_default_params={"vector_store_ids": ["T37J8R4WTM"]}, prompt_id=None, prompt_variables=None, dynamic_callback_params={}, litellm_logging_obj=logging_obj, ) proxy_router.avector_store_search.assert_awaited_once_with( vector_store_id="T37J8R4WTM", query="what is litellm?", custom_llm_provider="bedrock", metadata={"user_api_key_team_id": "team-a"}, ) assert messages[0]["content"] == "Context:\n\nrouted context\n\n" @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_completion( setup_vector_store_registry, ): litellm.turn_on_debug() client = AsyncHTTPHandler() print("value of litellm.vector_store_registry:", litellm.vector_store_registry) with patch.object(client, "post") as mock_post: # Mock the response for the LLM call mock_response = Mock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} # Provide proper JSON response content mock_response.text = json.dumps( { "id": "msg_01ABC123", "type": "message", "role": "assistant", "content": [ { "type": "text", "text": "LiteLLM is a library that simplifies LLM API access.", } ], "model": "claude-3.5-sonnet", "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 100, "output_tokens": 50}, } ) mock_response.json = lambda: json.loads(mock_response.text) mock_post.return_value = mock_response try: response = await litellm.acompletion( model="anthropic/claude-3.5-sonnet", messages=[{"role": "user", "content": "what is litellm?"}], vector_store_ids=["T37J8R4WTM"], client=client, ) except Exception as e: print(f"Error: {e}") # Verify the LLM request was made mock_post.assert_called_once() # Verify the request body print("call args:", mock_post.call_args) request_body = mock_post.call_args.kwargs["json"] print("Request body:", json.dumps(request_body, indent=4, default=str)) # Assert content from the knowedge base was applied to the request # 1. we should have 2 content blocks, the first is the context from the knowledge base, the second is the user message content = request_body["messages"][0]["content"] assert len(content) == 2 assert content[0]["type"] == "text" assert content[1]["type"] == "text" # 2. the first content block should have the bedrock knowledge base prefix string # this helps confirm that the context from the knowledge base was applied to the request assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in content[0]["text"] @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call( setup_vector_store_registry, ): """ Test that the Bedrock Knowledge Base Hook works when making a real llm api call and returns citations. """ # Init client litellm.turn_on_debug() async_client = AsyncHTTPHandler() response = await litellm.acompletion( model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", messages=[{"role": "user", "content": "what is litellm?"}], vector_store_ids=["T37J8R4WTM"], client=async_client, ) print("OPENAI RESPONSE:", json.dumps(dict(response), indent=4, default=str)) assert response is not None # Check that search_results are present in provider_specific_fields assert hasattr(response.choices[0].message, "provider_specific_fields") provider_fields = response.choices[0].message.provider_specific_fields assert provider_fields is not None assert "search_results" in provider_fields search_results = provider_fields["search_results"] assert search_results is not None assert len(search_results) > 0 # Check search result structure (OpenAI-compatible format) first_search_result = search_results[0] assert "object" in first_search_result assert first_search_result["object"] == "vector_store.search_results.page" assert "data" in first_search_result assert len(first_search_result["data"]) > 0 # Check individual result structure first_result = first_search_result["data"][0] assert "score" in first_result assert "content" in first_result print(f"Search results returned: {len(search_results)}") print(f"First search result has {len(first_search_result['data'])} items") @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_streaming( setup_vector_store_registry, ): """ Test that the Bedrock Knowledge Base Hook works with streaming and returns search_results in chunks. """ # Init client # litellm.turn_on_debug() async_client = AsyncHTTPHandler() response = await litellm.acompletion( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], vector_store_ids=["T37J8R4WTM"], stream=True, client=async_client, ) # Collect chunks chunks = [] search_results_found = False async for chunk in response: chunks.append(chunk) print(f"Chunk: {chunk}") # Check if this chunk has search_results in provider_specific_fields if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: if hasattr(choice, "delta") and choice.delta: provider_fields = getattr( choice.delta, "provider_specific_fields", None ) if provider_fields and "search_results" in provider_fields: search_results = provider_fields["search_results"] print( f"Found search_results in streaming chunk: {len(search_results)} results" ) # Verify structure assert search_results is not None assert len(search_results) > 0 first_search_result = search_results[0] assert "object" in first_search_result assert ( first_search_result["object"] == "vector_store.search_results.page" ) assert "data" in first_search_result assert len(first_search_result["data"]) > 0 search_results_found = True print(f"Total chunks received: {len(chunks)}") assert len(chunks) > 0 assert search_results_found, "search_results should be present in streaming chunks" @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools( setup_vector_store_registry, ): """ Test that the Bedrock Knowledge Base Hook works when making a real llm api call """ # Init client litellm.turn_on_debug() response = await litellm.acompletion( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], max_tokens=10, tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], ) assert response is not None @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_and_filters( setup_vector_store_registry, ): """ Test that filters from file_search tools are properly passed through to vector store search. This test verifies the entire flow: tool parsing -> filter extraction -> vector store API call. In this case we filter for a non-existent user_id, which should return no results. """ litellm.turn_on_debug() response = await litellm.acompletion( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], max_tokens=10, tools=[ { "type": "file_search", "vector_store_ids": ["T37J8R4WTM"], "filters": { "key": "user_id", "value": "fake-user-id", "operator": "eq", }, } ], ) # Verify response is not None assert response is not None # Verify search results were added to the response (this proves the search was called) assert hasattr(response.choices[0].message, "provider_specific_fields") provider_fields = response.choices[0].message.provider_specific_fields assert provider_fields is not None assert ( "search_results" in provider_fields ), "search_results not in provider_specific_fields" search_results = provider_fields["search_results"] assert ( search_results is not None and len(search_results) > 0 ), "No search results found" # The search was performed - this confirms filters were passed through # The logs above show: litellm.asearch(... filters={'key': 'user_id', 'value': 'fake-user-id', 'operator': 'eq'}) # And the Bedrock API request contains: {'filter': {'equals': {'key': 'user_id', 'value': 'fake-user-id'}}} print("✅ Filters were successfully passed through to vector store search") print(f" Search was performed and {len(search_results)} result(s) returned") @pytest.mark.asyncio async def test_bedrock_kb_request_body_has_transformed_filters( setup_vector_store_registry, ): """ Validate that the Bedrock Knowledge Base request body contains the transformed filters. """ captured_request_body: dict = {} async def fake_async_vector_store_search_handler( vector_store_id, query, vector_store_search_optional_params, vector_store_provider_config, custom_llm_provider, litellm_params, logging_obj, embedding_executor=None, extra_headers=None, extra_body=None, timeout=None, client=None, _is_async=False, ): litellm_params_dict = ( litellm_params.model_dump(exclude_none=False) if hasattr(litellm_params, "model_dump") else dict(litellm_params) ) api_base = vector_store_provider_config.get_complete_url( api_base=litellm_params_dict.get("api_base"), litellm_params=litellm_params_dict, ) url, request_body = ( vector_store_provider_config.transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=litellm_params_dict, extra_body=None, ) ) captured_request_body["url"] = url captured_request_body["body"] = request_body return VectorStoreSearchResponse( object="vector_store.search_results.page", search_query=query if isinstance(query, str) else " ".join(query), data=[ VectorStoreSearchResult( score=0.9, content=[ VectorStoreResultContent( text="LiteLLM is a library", type="text" ) ], ) ], ) with patch.object( litellm.vector_stores.main.base_llm_http_handler, "async_vector_store_search_handler", new=AsyncMock(side_effect=fake_async_vector_store_search_handler), ): response = await litellm.acompletion( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], max_tokens=10, tools=[ { "type": "file_search", "vector_store_ids": ["T37J8R4WTM"], "filters": { "key": "user_id", "value": "fake-user-id", "operator": "eq", }, } ], ) assert response is not None print( "captured_request_body:", json.dumps(captured_request_body, indent=4, default=str), ) assert "body" in captured_request_body, "Bedrock KB request body was not captured" vector_search = captured_request_body["body"]["retrievalConfiguration"][ "vectorSearchConfiguration" ] aws_filter = vector_search["filter"] assert "equals" in aws_filter, f"Expected 'equals' in AWS format, got: {aws_filter}" assert aws_filter["equals"]["key"] == "user_id" assert aws_filter["equals"]["value"] == "fake-user-id" print("✅ Filters transformed correctly: OpenAI format -> AWS Bedrock format") @pytest.mark.asyncio async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registry): """ Tests that knowledge base content is correctly passed to the OpenAI API call """ litellm.set_verbose = True from openai import AsyncOpenAI client = AsyncOpenAI(api_key="fake-api-key") # Variable to capture the request captured_request = {} with patch.object( client.chat.completions.with_raw_response, "create" ) as mock_client: # Create async mock that returns proper structure async def mock_create(**kwargs): mock_response = Mock() mock_response.choices = [ Mock( message=Mock(content="Mock response from OpenAI", role="assistant") ) ] mock_response.usage = Mock( prompt_tokens=100, completion_tokens=50, total_tokens=150 ) mock_response.id = "chatcmpl-123" mock_response.object = "chat.completion" mock_response.created = 1234567890 mock_response.model = "gpt-5.5" # Store the request for verification captured_request.update(kwargs) # Return wrapper with parse method wrapper = Mock() wrapper.parse.return_value = mock_response return wrapper mock_client.side_effect = mock_create try: await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], vector_store_ids=["T37J8R4WTM"], client=client, ) except Exception as e: print(f"Error: {e}") # Verify the API was called mock_client.assert_called_once() request_body = captured_request # Verify the request contains messages with knowledge base context assert "messages" in request_body messages = request_body["messages"] # We expect at least 2 messages: # 1. User message with the knowledge base context # 2. User message with the question assert len(messages) >= 2 print("request messages:", json.dumps(messages, indent=4, default=str)) # assert message[0] is the user message with the knowledge base context assert messages[0]["role"] == "user" assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"] @pytest.mark.asyncio async def test_openai_with_vector_store_ids_in_tool_call_mock_openai( setup_vector_store_registry, ): """ Tests that vector store ids can be passed as tools This is the OpenAI format """ litellm.set_verbose = True from openai import AsyncOpenAI client = AsyncOpenAI(api_key="fake-api-key") # Variable to capture the request captured_request = {} with patch.object( client.chat.completions.with_raw_response, "create" ) as mock_client: # Create async mock that returns proper structure async def mock_create(**kwargs): mock_response = Mock() mock_response.choices = [ Mock( message=Mock(content="Mock response from OpenAI", role="assistant") ) ] mock_response.usage = Mock( prompt_tokens=100, completion_tokens=50, total_tokens=150 ) mock_response.id = "chatcmpl-123" mock_response.object = "chat.completion" mock_response.created = 1234567890 mock_response.model = "gpt-5.5" # Store the request for verification captured_request.update(kwargs) # Return wrapper with parse method wrapper = Mock() wrapper.parse.return_value = mock_response return wrapper mock_client.side_effect = mock_create try: await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], client=client, ) except Exception as e: print(f"Error: {e}") # Verify the API was called mock_client.assert_called_once() request_body = captured_request print("request body:", json.dumps(request_body, indent=4, default=str)) # Verify the request contains messages with knowledge base context assert "messages" in request_body messages = request_body["messages"] # We expect at least 2 messages: # 1. User message with the knowledge base context # 2. User message with the question assert len(messages) >= 2 print("request messages:", json.dumps(messages, indent=4, default=str)) # assert message[0] is the user message with the knowledge base context assert messages[0]["role"] == "user" assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"] # assert that the tool call was not sent to the upstream llm API if it's a litellm vector store assert "tools" not in request_body @pytest.mark.asyncio async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_registry): """Ensure unrecognized vector store tools are forwarded to the provider""" from openai import AsyncOpenAI client = AsyncOpenAI(api_key="fake-api-key") # Variable to capture the request captured_request = {} with patch.object( client.chat.completions.with_raw_response, "create" ) as mock_client: # Create async mock that returns proper structure async def mock_create(**kwargs): mock_response = Mock() mock_response.choices = [ Mock( message=Mock(content="Mock response from OpenAI", role="assistant") ) ] mock_response.usage = Mock( prompt_tokens=100, completion_tokens=50, total_tokens=150 ) mock_response.id = "chatcmpl-123" mock_response.object = "chat.completion" mock_response.created = 1234567890 mock_response.model = "gpt-5.5" # Store the request for verification captured_request.update(kwargs) # Return wrapper with parse method wrapper = Mock() wrapper.parse.return_value = mock_response return wrapper mock_client.side_effect = mock_create try: await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], tools=[ {"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}, {"type": "file_search", "vector_store_ids": ["unknownVS"]}, ], client=client, ) except Exception as e: print(f"Error: {e}") mock_client.assert_called_once() request_body = captured_request assert "messages" in request_body messages = request_body["messages"] assert len(messages) >= 2 assert messages[0]["role"] == "user" assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"] assert "tools" in request_body tools = request_body["tools"] assert len(tools) == 1 assert tools[0]["vector_store_ids"] == ["unknownVS"] # @pytest.mark.asyncio # 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 = MockCustomLogger() # litellm.set_verbose = True # await litellm.acompletion( # model="gpt-5.5", # messages=[{"role": "user", "content": "what is litellm?"}], # vector_store_ids = [ # "T37J8R4WTM" # ], # ) # # sleep for 1 second to allow the logging callback to run # await asyncio.sleep(1) # # assert that the knowledge base request was logged in the standard logging payload # standard_logging_payload: Optional[StandardLoggingPayload] = test_custom_logger.standard_logging_payload # assert standard_logging_payload is not None # metadata = standard_logging_payload["metadata"] # standard_logging_vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = metadata["vector_store_request_metadata"] # print("standard_logging_vector_store_request_metadata:", json.dumps(standard_logging_vector_store_request_metadata, indent=4, default=str)) # # 1 vector store request was made, expect 1 vector store request metadata object # assert len(standard_logging_vector_store_request_metadata) == 1 # # expect the vector store request metadata object to have the correct values # vector_store_request_metadata = standard_logging_vector_store_request_metadata[0] # assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM" # assert vector_store_request_metadata.get("query") == "what is litellm?" # assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock" # vector_store_search_response: VectorStoreSearchResponse = vector_store_request_metadata.get("vector_store_search_response") # assert vector_store_search_response is not None # assert vector_store_search_response.get("search_query") == "what is litellm?" # assert len(vector_store_search_response.get("data", [])) >=0 # for item in vector_store_search_response.get("data", []): # assert item.get("score") is not None # assert item.get("content") is not None # assert len(item.get("content", [])) >= 0 # for content_item in item.get("content", []): # text_content = content_item.get("text") # assert text_content is not None # assert len(text_content) > 0 @pytest.mark.asyncio async def test_provider_specific_fields_in_proxy_http_response( setup_vector_store_registry, ): """ Test that provider_specific_fields (like search_results) are included in the proxy HTTP JSON response, not just in Python SDK objects. This test catches serialization bugs where exclude=True would strip provider_specific_fields from the HTTP response. """ from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app, initialize from unittest.mock import patch as mock_patch # Initialize proxy await initialize( model="gpt-5-mini", alias=None, api_base=None, debug=False, temperature=None, max_tokens=None, request_timeout=600, max_budget=None, drop_params=True, add_function_to_prompt=False, headers=None, save=False, use_queue=False, config=None, ) # Create test client client = TestClient(app) # Create mock response with provider_specific_fields mock_response = litellm.ModelResponse( id="test-123", model="gpt-5-mini", created=1234567890, object="chat.completion", ) # Create message with provider_specific_fields mock_message = litellm.Message( content="LiteLLM is a tool that simplifies working with multiple LLMs.", role="assistant", provider_specific_fields={ "search_results": [ { "object": "vector_store.search_results.page", "search_query": "what is litellm?", "data": [ { "score": 0.95, "content": [{"text": "Test content", "type": "text"}], "file_id": "test-file", "filename": "test.txt", } ], } ] }, ) mock_choice = litellm.Choices(finish_reason="stop", index=0, message=mock_message) mock_response.choices = [mock_choice] mock_response.usage = litellm.Usage( prompt_tokens=10, completion_tokens=20, total_tokens=30 ) # Patch the completion call at the proxy level with mock_patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)): # Make HTTP request to proxy response = client.post( "/v1/chat/completions", json={ "model": "gpt-5-mini", "messages": [{"role": "user", "content": "What is litellm?"}], }, ) # Check HTTP response assert response.status_code == 200 result = response.json() print("HTTP Response JSON:", json.dumps(result, indent=2)) # THE KEY ASSERTIONS - These would FAIL with exclude=True! assert "choices" in result assert len(result["choices"]) > 0 choice = result["choices"][0] assert "message" in choice message = choice["message"] # Verify provider_specific_fields is in the JSON response assert ( "provider_specific_fields" in message ), "provider_specific_fields missing from HTTP JSON response! This means exclude=True is preventing serialization." assert "search_results" in message["provider_specific_fields"] search_results = message["provider_specific_fields"]["search_results"] assert len(search_results) > 0 # Verify search result structure first_result = search_results[0] assert first_result["object"] == "vector_store.search_results.page" assert "data" in first_result assert len(first_result["data"]) > 0 print("✅ provider_specific_fields successfully serialized in HTTP response")