diff --git a/tests/llm_translation/base_realtime_tests.py b/tests/llm_translation/base_realtime_tests.py new file mode 100644 index 00000000000..0e95a7420d1 --- /dev/null +++ b/tests/llm_translation/base_realtime_tests.py @@ -0,0 +1,426 @@ +""" +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 +import sys +from abc import ABC, abstractmethod +from typing import Optional + +import pytest +import websockets + +sys.path.insert(0, os.path.abspath("../..")) + +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 sends "session.created", xAI sends "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) -> str: + """Return the expected initial event type (e.g., 'session.created' or 'conversation.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] + assert initial_event["type"] == self.get_initial_event_type(), \ + f"Expected {self.get_initial_event_type()}, 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 diff --git a/tests/llm_translation/test_xai_realtime.py b/tests/llm_translation/test_xai_realtime.py new file mode 100644 index 00000000000..3d4f6e68098 --- /dev/null +++ b/tests/llm_translation/test_xai_realtime.py @@ -0,0 +1,34 @@ +""" +xAI Realtime API E2E Tests + +Tests xAI's Grok Voice Agent API through LiteLLM's realtime interface. +Uses the base test class to ensure consistent behavior across providers. +""" +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from tests.llm_translation.base_realtime_tests import BaseRealtimeTest + + +class TestXAIRealtime(BaseRealtimeTest): + """ + E2E tests for xAI Realtime API. + + xAI's Grok Voice Agent API is OpenAI-compatible but uses: + - Different initial event: "conversation.created" instead of "session.created" + - Different endpoint: wss://api.x.ai/v1/realtime + - Model: grok-4-1-fast-non-reasoning + """ + + def get_model(self) -> str: + return "xai/grok-4-1-fast-non-reasoning" + + def get_api_key_env_var(self) -> str: + return "XAI_API_KEY" + + def get_initial_event_type(self) -> str: + return "conversation.created"