diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py index 0d3e54b5bdc..dc441736efa 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py @@ -4,7 +4,6 @@ Integration tests for WebSearch interception with chat completions API. Tests the end-to-end flow of websearch_interception callback with litellm.acompletion() for transparent server-side web search execution. """ - import os from unittest.mock import AsyncMock, MagicMock, patch @@ -46,7 +45,7 @@ def websearch_logger(): ) async def test_websearch_chat_completion_with_openai(): """Test websearch interception with OpenAI chat completions API. - + This test verifies that: 1. Model calls litellm_web_search tool 2. Server executes web search automatically @@ -59,15 +58,12 @@ async def test_websearch_chat_completion_with_openai(): enabled_providers=[LlmProviders.OPENAI] ) litellm.callbacks = [websearch_logger] - + try: response = await litellm.acompletion( model="gpt-4o-mini", # Use cheaper model for testing messages=[ - { - "role": "user", - "content": "What's the weather in San Francisco today?", - } + {"role": "user", "content": "What's the weather in San Francisco today?"} ], tools=[ { @@ -89,12 +85,12 @@ async def test_websearch_chat_completion_with_openai(): } ], ) - + # Verify response structure assert isinstance(response, ModelResponse) assert response.choices[0].message.content is not None assert len(response.choices[0].message.content) > 0 - + # If agentic loop worked, we should NOT have tool_calls in final response # (they should have been executed and replaced with final answer) if hasattr(response.choices[0].message, "tool_calls"): @@ -103,9 +99,10 @@ async def test_websearch_chat_completion_with_openai(): pytest.skip( "Agentic loop did not execute - search tool may not be configured" ) + # Verify we got a meaningful response assert response.choices[0].finish_reason in ["stop", "end_turn"] - + finally: # Restore original callbacks litellm.callbacks = original_callbacks @@ -120,11 +117,11 @@ async def test_websearch_chat_completion_hook_detection(): Function, Message, ) - + websearch_logger = WebSearchInterceptionLogger( enabled_providers=[LlmProviders.OPENAI] ) - + # Mock response with litellm_web_search tool call mock_response = ModelResponse( id="test-123", @@ -145,14 +142,14 @@ async def test_websearch_chat_completion_hook_detection(): ), ) ], - ), + ) ) ], model="gpt-4o", object="chat.completion", created=1234567890, ) - + # Test should_run_chat_completion_agentic_loop should_run, tools_dict = ( await websearch_logger.async_should_run_chat_completion_agentic_loop( @@ -170,7 +167,7 @@ async def test_websearch_chat_completion_hook_detection(): kwargs={}, ) ) - + # Verify hook detected the tool call assert should_run is True assert "tool_calls" in tools_dict @@ -183,11 +180,11 @@ async def test_websearch_chat_completion_hook_detection(): async def test_websearch_not_triggered_without_tool(): """Test that websearch hook is NOT triggered when no web search tool in request.""" from litellm.types.utils import Choices, Message - + websearch_logger = WebSearchInterceptionLogger( enabled_providers=[LlmProviders.OPENAI] ) - + mock_response = ModelResponse( id="test-123", choices=[ @@ -198,14 +195,14 @@ async def test_websearch_not_triggered_without_tool(): role="assistant", content="Here's the answer", tool_calls=None, - ), + ) ) ], model="gpt-4o", object="chat.completion", created=1234567890, ) - + # Test without web search tool should_run, tools_dict = ( await websearch_logger.async_should_run_chat_completion_agentic_loop( @@ -223,7 +220,7 @@ async def test_websearch_not_triggered_without_tool(): kwargs={}, ) ) - + # Verify hook did NOT trigger assert should_run is False assert tools_dict == {} @@ -243,7 +240,7 @@ async def test_websearch_not_triggered_for_disabled_provider(): websearch_logger = WebSearchInterceptionLogger( enabled_providers=[LlmProviders.BEDROCK] ) - + mock_response = ModelResponse( id="test-123", choices=[ @@ -263,14 +260,14 @@ async def test_websearch_not_triggered_for_disabled_provider(): ), ) ], - ), + ) ) ], model="gpt-4o", object="chat.completion", created=1234567890, ) - + # Test with OpenAI provider (not enabled) should_run, tools_dict = ( await websearch_logger.async_should_run_chat_completion_agentic_loop( @@ -288,7 +285,7 @@ async def test_websearch_not_triggered_for_disabled_provider(): kwargs={}, ) ) - + # Verify hook did NOT trigger assert should_run is False assert tools_dict == {} @@ -297,7 +294,7 @@ async def test_websearch_not_triggered_for_disabled_provider(): @pytest.mark.asyncio async def test_websearch_json_serialization_fix(): """Test that tool call arguments are properly JSON serialized. - + Regression test for the bug where arguments were converted to Python string representation instead of proper JSON, causing providers like MiniMax to reject requests with 'invalid function arguments json string'. @@ -314,25 +311,25 @@ async def test_websearch_json_serialization_fix(): "input": {"query": "weather in SF"}, # Dict input } ] - + search_results = ["Weather: 65°F, partly cloudy"] - + # Transform to OpenAI format assistant_message, tool_messages = WebSearchTransformation.transform_response( tool_calls=tool_calls, search_results=search_results, response_format="openai", ) - + # Verify arguments are properly JSON serialized import json - + arguments_str = assistant_message["tool_calls"][0]["function"]["arguments"] - + # Should be valid JSON parsed_args = json.loads(arguments_str) assert parsed_args == {"query": "weather in SF"} - + # Should NOT be Python string representation like "{'query': 'weather in SF'}" assert arguments_str == '{"query": "weather in SF"}' assert arguments_str != "{'query': 'weather in SF'}" @@ -346,7 +343,7 @@ async def test_websearch_json_serialization_fix(): ) async def test_websearch_streaming_conversion(): """Test that streaming requests are converted to non-streaming for web search. - + When stream=True is passed with web search tools, the handler should: 1. Convert stream=True to stream=False for initial request 2. Execute web search @@ -356,11 +353,13 @@ async def test_websearch_streaming_conversion(): enabled_providers=[LlmProviders.OPENAI], search_tool_name="perplexity-search" ) litellm.callbacks = [websearch_logger] - + try: response = await litellm.acompletion( model="gpt-4o-mini", - messages=[{"role": "user", "content": "What's the latest AI news?"}], + messages=[ + {"role": "user", "content": "What's the latest AI news?"} + ], tools=[ { "type": "function", @@ -376,24 +375,23 @@ async def test_websearch_streaming_conversion(): ], stream=True, ) - + # Response should be a streaming iterator chunks = [] async for chunk in response: chunks.append(chunk) - + # Verify we got streaming chunks assert len(chunks) > 0 - + # Verify chunks have expected structure for chunk in chunks: assert hasattr(chunk, "choices") assert len(chunk.choices) > 0 - + finally: litellm.callbacks = [] - @pytest.mark.asyncio async def test_bedrock_converse_agentic_hook_triggers(): """Test that BedrockConverseLLM._call_agentic_chat_completion_hooks @@ -607,7 +605,6 @@ async def test_bedrock_converse_agentic_hook_skips_when_no_tool_call(): finally: litellm.callbacks = original_callbacks - if __name__ == "__main__": # Run with: pytest test_websearch_chat_completion.py -v -s pytest.main([__file__, "-v", "-s"])