mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
462 lines
17 KiB
Python
462 lines
17 KiB
Python
"""
|
|
Base test class for LiteLLM Realtime API E2E tests.
|
|
|
|
Provides common test infrastructure for testing realtime WebSocket connections
|
|
across different providers (OpenAI, xAI, etc.)
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from typing import Optional, Tuple, Union
|
|
|
|
import pytest
|
|
import websockets
|
|
|
|
|
|
import litellm
|
|
|
|
|
|
class RealTimeWebSocketClient:
|
|
"""
|
|
Mock WebSocket client for testing realtime connections.
|
|
Captures messages sent from the backend and provides a simple interface
|
|
for testing connection success.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.messages_sent = []
|
|
self.messages_received = []
|
|
self.received_initial_event = False
|
|
self.connection_successful = False
|
|
self.close_code = None
|
|
self.close_reason = None
|
|
# Required by realtime_streaming.py - import exceptions module
|
|
from websockets import exceptions as websockets_exceptions
|
|
|
|
self.exceptions = websockets_exceptions
|
|
|
|
async def accept(self):
|
|
"""Accept the WebSocket connection"""
|
|
pass
|
|
|
|
async def send_text(self, message):
|
|
"""Receive message from backend and store it"""
|
|
self.messages_sent.append(message)
|
|
try:
|
|
if isinstance(message, bytes):
|
|
message_str = message.decode("utf-8")
|
|
else:
|
|
message_str = message
|
|
|
|
msg_data = json.loads(message_str)
|
|
msg_type = msg_data.get("type", "unknown")
|
|
|
|
# Pretty print API response
|
|
print(f"\n{'='*80}")
|
|
print(
|
|
f"API RESPONSE #{len(self.messages_received) + 1} - Event: {msg_type}"
|
|
)
|
|
print(f"{'='*80}")
|
|
print(json.dumps(msg_data, indent=2, sort_keys=False))
|
|
print(f"{'='*80}\n")
|
|
|
|
self.messages_received.append(msg_data)
|
|
|
|
# Check for initial connection event
|
|
if not self.received_initial_event and self._is_initial_event(msg_type):
|
|
self.received_initial_event = True
|
|
self.connection_successful = True
|
|
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
|
# Non-JSON messages are acceptable
|
|
print(f"\n[Non-JSON message: {e}]")
|
|
print(f"Raw content: {str(message)[:200]}\n")
|
|
pass
|
|
|
|
def _is_initial_event(self, msg_type: str) -> bool:
|
|
"""Check if message type is an initial connection event"""
|
|
# OpenAI and xAI send "session.created"; some providers send "conversation.created"
|
|
return msg_type in ["session.created", "conversation.created"]
|
|
|
|
async def receive_text(self):
|
|
"""
|
|
Wait briefly for messages, then close connection.
|
|
This allows the backend forwarding task to send messages.
|
|
"""
|
|
print(f"\nWaiting for connection to establish...")
|
|
max_wait = 5.0
|
|
check_interval = 0.1
|
|
waited = 0.0
|
|
|
|
while waited < max_wait:
|
|
if self.connection_successful:
|
|
print(f"Connection successful after {waited:.1f}s\n")
|
|
break
|
|
await asyncio.sleep(check_interval)
|
|
waited += check_interval
|
|
|
|
if not self.connection_successful:
|
|
print(f"Warning: No initial event received after {max_wait}s\n")
|
|
|
|
# If we have a pending message to send, send it now
|
|
if hasattr(self, "_pending_client_message") and self._pending_client_message:
|
|
print(f"Sending client message to backend...\n")
|
|
# This simulates receiving a message from the client that needs to be forwarded to backend
|
|
# We return it as if it came from the client
|
|
msg = self._pending_client_message
|
|
self._pending_client_message = None
|
|
return msg
|
|
|
|
# Close connection to end the test
|
|
print(f"\n{'='*80}")
|
|
print(f"TEST COMPLETE - Closing connection")
|
|
print(f"Total messages received from API: {len(self.messages_received)}")
|
|
print(f"{'='*80}\n")
|
|
raise websockets.exceptions.ConnectionClosed(None, None)
|
|
|
|
def queue_client_message(self, message: str):
|
|
"""Queue a message to be sent from 'client' to backend"""
|
|
self._pending_client_message = message
|
|
|
|
async def close(self, code=1000, reason=""):
|
|
"""Close the WebSocket"""
|
|
self.close_code = code
|
|
self.close_reason = reason
|
|
|
|
@property
|
|
def headers(self):
|
|
return {}
|
|
|
|
|
|
class BaseRealtimeTest(ABC):
|
|
"""
|
|
Abstract base test class for realtime API tests.
|
|
|
|
Child classes must implement:
|
|
- get_model(): Return the model name to test
|
|
- get_api_key_env_var(): Return the environment variable name for the API key
|
|
- get_initial_event_type(): Return the expected initial event type (e.g., "session.created")
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_model(self) -> str:
|
|
"""Return the model name to test (e.g., 'gpt-4o-realtime-preview-2024-10-01')"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_api_key_env_var(self) -> str:
|
|
"""Return the environment variable name for the API key (e.g., 'OPENAI_API_KEY')"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_initial_event_type(self) -> Union[str, Tuple[str, ...]]:
|
|
"""Return the expected initial event type(s).
|
|
|
|
May return a single event type (e.g. ``'session.created'``) or a tuple
|
|
of acceptable types when the upstream provider can legitimately emit
|
|
more than one initial event (e.g. xAI's Grok Voice Agent has shipped
|
|
both ``conversation.created`` and ``session.created``).
|
|
"""
|
|
pass
|
|
|
|
def get_skip_reason(self) -> str:
|
|
"""Return the skip reason when API key is missing"""
|
|
return f"No {self.get_api_key_env_var()} provided"
|
|
|
|
def should_skip(self) -> bool:
|
|
"""Check if tests should be skipped due to missing API key"""
|
|
return os.environ.get(self.get_api_key_env_var()) is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_realtime_connection(self):
|
|
"""
|
|
Test basic realtime WebSocket connection.
|
|
Verifies that:
|
|
1. Connection is established successfully
|
|
2. Initial event is received
|
|
3. Messages are properly forwarded
|
|
"""
|
|
litellm._turn_on_debug()
|
|
if self.should_skip():
|
|
pytest.skip(self.get_skip_reason())
|
|
|
|
websocket_client = RealTimeWebSocketClient()
|
|
caught_exception = None
|
|
|
|
print(f"\n{'='*80}")
|
|
print(f"STARTING REALTIME CONNECTION TEST")
|
|
print(f"Model: {self.get_model()}")
|
|
print(f"API Key Env Var: {self.get_api_key_env_var()}")
|
|
print(f"{'='*80}\n")
|
|
|
|
try:
|
|
await litellm._arealtime(
|
|
model=self.get_model(),
|
|
websocket=websocket_client,
|
|
api_key=os.environ.get(self.get_api_key_env_var()),
|
|
timeout=60,
|
|
)
|
|
except websockets.exceptions.ConnectionClosed:
|
|
pass
|
|
except Exception as e:
|
|
print(f"\nException: {type(e).__name__}: {e}\n")
|
|
caught_exception = e
|
|
|
|
# Build debug info
|
|
error_details = []
|
|
error_details.append(f"messages_sent: {len(websocket_client.messages_sent)}")
|
|
error_details.append(
|
|
f"messages_received: {len(websocket_client.messages_received)}"
|
|
)
|
|
error_details.append(f"close_code: {websocket_client.close_code}")
|
|
error_details.append(f"close_reason: {websocket_client.close_reason}")
|
|
if caught_exception:
|
|
error_details.append(
|
|
f"exception: {type(caught_exception).__name__}: {caught_exception}"
|
|
)
|
|
|
|
# Skip on transient connection failures
|
|
if (
|
|
not websocket_client.connection_successful
|
|
and websocket_client.close_code is not None
|
|
):
|
|
pytest.skip(f"Transient connection failure: {'; '.join(error_details)}")
|
|
|
|
# Assertions
|
|
assert (
|
|
websocket_client.connection_successful
|
|
), f"Failed to connect. Debug: {'; '.join(error_details)}"
|
|
assert websocket_client.received_initial_event, f"Did not receive initial event"
|
|
assert len(websocket_client.messages_received) > 0, "No messages received"
|
|
|
|
# Verify initial event
|
|
initial_event = websocket_client.messages_received[0]
|
|
expected_event_type = self.get_initial_event_type()
|
|
if isinstance(expected_event_type, str):
|
|
allowed_event_types: Tuple[str, ...] = (expected_event_type,)
|
|
else:
|
|
allowed_event_types = tuple(expected_event_type)
|
|
assert (
|
|
initial_event["type"] in allowed_event_types
|
|
), f"Expected one of {allowed_event_types}, got {initial_event.get('type')}"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_realtime_with_query_params(self):
|
|
"""
|
|
Test realtime connection with explicit query parameters.
|
|
Verifies that query params are properly passed to the backend.
|
|
"""
|
|
litellm._turn_on_debug()
|
|
if self.should_skip():
|
|
pytest.skip(self.get_skip_reason())
|
|
|
|
from litellm.types.realtime import RealtimeQueryParams
|
|
|
|
websocket_client = RealTimeWebSocketClient()
|
|
caught_exception = None
|
|
|
|
# Strip provider prefix from model name for query params
|
|
model_name = self.get_model()
|
|
if "/" in model_name:
|
|
model_name = model_name.split("/", 1)[1]
|
|
|
|
query_params: RealtimeQueryParams = {"model": model_name}
|
|
|
|
try:
|
|
await litellm._arealtime(
|
|
model=self.get_model(),
|
|
websocket=websocket_client,
|
|
api_key=os.environ.get(self.get_api_key_env_var()),
|
|
query_params=query_params,
|
|
timeout=60,
|
|
)
|
|
except websockets.exceptions.ConnectionClosed:
|
|
pass
|
|
except Exception as e:
|
|
caught_exception = e
|
|
|
|
# Build debug info
|
|
error_details = []
|
|
error_details.append(f"messages_sent: {len(websocket_client.messages_sent)}")
|
|
error_details.append(
|
|
f"messages_received: {len(websocket_client.messages_received)}"
|
|
)
|
|
if caught_exception:
|
|
error_details.append(
|
|
f"exception: {type(caught_exception).__name__}: {caught_exception}"
|
|
)
|
|
|
|
# Skip on transient failures
|
|
if (
|
|
not websocket_client.connection_successful
|
|
and websocket_client.close_code is not None
|
|
):
|
|
pytest.skip(f"Transient connection failure: {'; '.join(error_details)}")
|
|
|
|
# Assertions
|
|
assert (
|
|
websocket_client.connection_successful
|
|
), f"Failed to connect. Debug: {'; '.join(error_details)}"
|
|
assert len(websocket_client.messages_received) > 0, "No messages received"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_user_message(self):
|
|
"""
|
|
Test sending an actual user message and receiving responses.
|
|
This creates a more realistic conversation flow.
|
|
"""
|
|
if self.should_skip():
|
|
pytest.skip(self.get_skip_reason())
|
|
|
|
litellm._turn_on_debug()
|
|
|
|
# Create a custom websocket client that sends a message
|
|
class InteractiveWebSocketClient(RealTimeWebSocketClient):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.sent_user_message = False
|
|
self.response_messages = []
|
|
self.wait_for_responses = True
|
|
|
|
async def receive_text(self):
|
|
"""Enhanced receive that sends a user message after connection"""
|
|
print(f"\n{'='*80}")
|
|
print(f"CLIENT-SIDE RECEIVE HANDLER")
|
|
print(f"{'='*80}\n")
|
|
|
|
# Wait for initial connection
|
|
max_wait = 5.0
|
|
check_interval = 0.1
|
|
waited = 0.0
|
|
|
|
while waited < max_wait:
|
|
if self.connection_successful:
|
|
print(f"Connection established after {waited:.1f}s\n")
|
|
break
|
|
await asyncio.sleep(check_interval)
|
|
waited += check_interval
|
|
|
|
# Step 1: Send a user message after connection is established
|
|
if self.connection_successful and not self.sent_user_message:
|
|
self.sent_user_message = True
|
|
user_msg_data = {
|
|
"type": "conversation.item.create",
|
|
"item": {
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "input_text", "text": "Say hi back to me!"}
|
|
],
|
|
},
|
|
}
|
|
user_msg = json.dumps(user_msg_data)
|
|
|
|
print(f"\n{'='*80}")
|
|
print(f"STEP 1: SENDING USER MESSAGE TO BACKEND")
|
|
print(f"{'='*80}")
|
|
print(json.dumps(user_msg_data, indent=2))
|
|
print(f"{'='*80}\n")
|
|
|
|
return user_msg
|
|
|
|
# Step 2: Trigger the response after user message is acknowledged
|
|
if not hasattr(self, "triggered_response"):
|
|
self.triggered_response = True
|
|
# Wait a bit for the user message to be processed
|
|
await asyncio.sleep(0.5)
|
|
|
|
response_create_data = {"type": "response.create"}
|
|
response_create = json.dumps(response_create_data)
|
|
|
|
print(f"\n{'='*80}")
|
|
print(f"STEP 2: TRIGGERING LLM RESPONSE")
|
|
print(f"{'='*80}")
|
|
print(json.dumps(response_create_data, indent=2))
|
|
print(f"{'='*80}\n")
|
|
|
|
return response_create
|
|
|
|
# Step 3: Wait for LLM responses
|
|
if self.wait_for_responses:
|
|
print(f"\nSTEP 3: Waiting 5 seconds for LLM to respond...\n")
|
|
await asyncio.sleep(5.0)
|
|
self.wait_for_responses = False
|
|
|
|
# Collect response info
|
|
for msg in self.messages_received:
|
|
msg_type = msg.get("type", "unknown")
|
|
if msg_type not in ["conversation.created", "ping"]:
|
|
self.response_messages.append(msg)
|
|
|
|
print(
|
|
f"\nReceived {len(self.response_messages)} response messages (excluding init/ping)\n"
|
|
)
|
|
|
|
print(f"\n{'='*80}")
|
|
print(f"CLOSING CONNECTION")
|
|
print(f"Total messages received: {len(self.messages_received)}")
|
|
print(f"{'='*80}\n")
|
|
raise websockets.exceptions.ConnectionClosed(None, None)
|
|
|
|
websocket_client = InteractiveWebSocketClient()
|
|
caught_exception = None
|
|
|
|
print(f"\n{'='*80}")
|
|
print(f"STARTING INTERACTIVE MESSAGE TEST")
|
|
print(f"Model: {self.get_model()}")
|
|
print(f"Message: 'Say hi back to me!'")
|
|
print(f"{'='*80}\n")
|
|
|
|
try:
|
|
await litellm._arealtime(
|
|
model=self.get_model(),
|
|
websocket=websocket_client,
|
|
api_key=os.environ.get(self.get_api_key_env_var()),
|
|
timeout=60,
|
|
)
|
|
except websockets.exceptions.ConnectionClosed:
|
|
pass
|
|
except Exception as e:
|
|
print(f"\nException: {type(e).__name__}: {e}\n")
|
|
caught_exception = e
|
|
|
|
# Print results
|
|
print(f"\n{'='*80}")
|
|
print(f"TEST RESULTS SUMMARY")
|
|
print(f"{'='*80}")
|
|
print(f"Connection successful: {websocket_client.connection_successful}")
|
|
print(f"User message sent: {websocket_client.sent_user_message}")
|
|
print(f"Total messages received: {len(websocket_client.messages_received)}")
|
|
print(
|
|
f"Response messages (excluding init/ping): {len(websocket_client.response_messages)}"
|
|
)
|
|
|
|
if websocket_client.response_messages:
|
|
print(f"\nResponse Event Types:")
|
|
for i, msg in enumerate(websocket_client.response_messages, 1):
|
|
print(f" {i}. {msg.get('type', 'unknown')}")
|
|
|
|
print(f"{'='*80}\n")
|
|
|
|
# Skip if no responses (might be timing issue)
|
|
if not websocket_client.response_messages:
|
|
pytest.skip("No response messages received (might be timing/network issue)")
|
|
|
|
assert websocket_client.connection_successful, "Failed to establish connection"
|
|
assert websocket_client.sent_user_message, "Failed to send user message"
|
|
|
|
def test_query_params_construction(self):
|
|
"""Test that query params are constructed correctly"""
|
|
from litellm.types.realtime import RealtimeQueryParams
|
|
|
|
# Strip provider prefix from model name
|
|
model_name = self.get_model()
|
|
if "/" in model_name:
|
|
model_name = model_name.split("/", 1)[1]
|
|
|
|
query_params: RealtimeQueryParams = {"model": model_name}
|
|
|
|
assert "model" in query_params
|
|
assert query_params["model"] == model_name
|