mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
lint fix
Signed-off-by: Nishchay Veer <nishchayveer19@gmail.com>
This commit is contained in:
Nishchay Veer
parent
b882bec39a
commit
f781d932f8
1 changed files with 37 additions and 40 deletions
|
|
@ -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"])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue