From d143df8591af75f8ea3361236de71a7fb1531a83 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 4 Mar 2026 18:27:06 +0530 Subject: [PATCH] Add tests for resp onses websocket --- .../test_responses_websocket_all_providers.py | 973 ++++++++++++++++++ 1 file changed, 973 insertions(+) create mode 100644 tests/test_litellm/responses/test_responses_websocket_all_providers.py diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py new file mode 100644 index 00000000000..0d83b9f88de --- /dev/null +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -0,0 +1,973 @@ +""" +Unit tests to verify that all providers support Responses API WebSocket mode. + +Tests that: +1. All providers with ResponsesAPIConfig support websocket mode +2. Providers with native websocket support use direct connection +3. Providers without native websocket support use ManagedResponsesWebSocketHandler +""" + +import pytest + +from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig +from litellm.llms.databricks.responses.transformation import ( + DatabricksResponsesAPIConfig, +) +from litellm.llms.github_copilot.responses.transformation import ( + GithubCopilotResponsesAPIConfig, +) +from litellm.llms.hosted_vllm.responses.transformation import ( + HostedVLLMResponsesAPIConfig, +) +from litellm.llms.litellm_proxy.responses.transformation import ( + LiteLLMProxyResponsesAPIConfig, +) +from litellm.llms.manus.responses.transformation import ManusResponsesAPIConfig +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.openrouter.responses.transformation import ( + OpenRouterResponsesAPIConfig, +) +from litellm.llms.perplexity.responses.transformation import PerplexityResponsesConfig +from litellm.llms.volcengine.responses.transformation import ( + VolcEngineResponsesAPIConfig, +) +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig + + +class TestResponsesAPIWebSocketSupport: + """Test that all providers have websocket support configured correctly""" + + def test_openai_supports_native_websocket(self): + """OpenAI should support native websocket""" + config = OpenAIResponsesAPIConfig() + assert ( + config.supports_native_websocket() is True + ), "OpenAI should support native websocket" + + def test_azure_supports_native_websocket(self): + """Azure should support native websocket (inherits from OpenAI)""" + config = AzureOpenAIResponsesAPIConfig() + assert ( + config.supports_native_websocket() is True + ), "Azure should support native websocket" + + def test_xai_uses_managed_websocket(self): + """XAI should use managed websocket handler""" + config = XAIResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "XAI should use managed websocket handler" + + def test_github_copilot_uses_managed_websocket(self): + """GitHub Copilot should use managed websocket handler""" + config = GithubCopilotResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "GitHub Copilot should use managed websocket handler" + + def test_chatgpt_uses_managed_websocket(self): + """ChatGPT should use managed websocket handler""" + config = ChatGPTResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "ChatGPT should use managed websocket handler" + + def test_litellm_proxy_uses_managed_websocket(self): + """LiteLLM Proxy should use managed websocket handler""" + config = LiteLLMProxyResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "LiteLLM Proxy should use managed websocket handler" + + def test_volcengine_uses_managed_websocket(self): + """VolcEngine should use managed websocket handler""" + config = VolcEngineResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "VolcEngine should use managed websocket handler" + + def test_manus_uses_managed_websocket(self): + """Manus should use managed websocket handler""" + config = ManusResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "Manus should use managed websocket handler" + + def test_perplexity_uses_managed_websocket(self): + """Perplexity should use managed websocket handler""" + config = PerplexityResponsesConfig() + assert ( + config.supports_native_websocket() is False + ), "Perplexity should use managed websocket handler" + + def test_databricks_uses_managed_websocket(self): + """Databricks should use managed websocket handler""" + config = DatabricksResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "Databricks should use managed websocket handler" + + def test_openrouter_uses_managed_websocket(self): + """OpenRouter should use managed websocket handler""" + config = OpenRouterResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "OpenRouter should use managed websocket handler" + + def test_hosted_vllm_uses_managed_websocket(self): + """Hosted vLLM should use managed websocket handler""" + config = HostedVLLMResponsesAPIConfig() + assert ( + config.supports_native_websocket() is False + ), "Hosted vLLM should use managed websocket handler" + + +class TestManagedWebSocketHandlerIntegration: + """Test that ManagedResponsesWebSocketHandler is properly integrated""" + + @pytest.mark.asyncio + async def test_managed_handler_instantiation(self): + """Test that ManagedResponsesWebSocketHandler can be instantiated""" + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + mock_websocket = MagicMock() + mock_logging_obj = Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ) + + handler = ManagedResponsesWebSocketHandler( + websocket=mock_websocket, + model="test-model", + logging_obj=mock_logging_obj, + user_api_key_dict=None, + litellm_metadata={}, + api_key="test-key", + api_base="https://api.example.com", + timeout=30.0, + custom_llm_provider="test_provider", + ) + + assert handler.model == "test-model" + assert handler.api_key == "test-key" + assert handler.api_base == "https://api.example.com" + assert handler.timeout == 30.0 + assert handler.custom_llm_provider == "test_provider" + + +class TestChunkTransformation: + """Test chunk serialization and transformation for WebSocket streaming""" + + def test_serialize_chunk_with_dict(self): + """Test serialization of dict chunks""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.created", + "response": {"id": "resp_456", "status": "in_progress"}, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.created" in serialized + assert "resp_456" in serialized + + def test_serialize_chunk_handles_invalid_json(self): + """Test that chunks with circular references are handled""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + # Create object with circular reference + obj = {"a": 1} + obj["self"] = obj # type: ignore + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(obj) + assert serialized is None + + def test_extract_output_messages_with_text_content(self): + """Test extraction of output messages with text content""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello world"}], + } + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["type"] == "message" + assert messages[0]["role"] == "assistant" + assert messages[0]["content"][0]["text"] == "Hello world" + + def test_extract_output_messages_with_multiple_content_parts(self): + """Test extraction with multiple content parts""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Part 1. "}, + {"type": "output_text", "text": "Part 2."}, + ], + } + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == "Part 1. Part 2." + + def test_extract_output_messages_with_function_calls(self): + """Test that function calls are preserved""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "function_call", + "id": "call_123", + "name": "get_weather", + "arguments": '{"location": "Paris"}', + } + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["type"] == "function_call" + assert messages[0]["id"] == "call_123" + assert messages[0]["name"] == "get_weather" + + def test_extract_output_messages_filters_empty_text(self): + """Test that messages with empty text are filtered out""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": ""}], + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Valid text"}], + }, + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == "Valid text" + + def test_extract_output_messages_handles_non_dict_items(self): + """Test that non-dict items are skipped""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + "invalid_string", + None, + 123, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Valid"}], + }, + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == "Valid" + + def test_input_to_messages_with_string(self): + """Test conversion of string input to messages""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + messages = ManagedResponsesWebSocketHandler._input_to_messages("Hello world") + assert len(messages) == 1 + assert messages[0]["type"] == "message" + assert messages[0]["role"] == "user" + assert messages[0]["content"][0]["type"] == "input_text" + assert messages[0]["content"][0]["text"] == "Hello world" + + def test_input_to_messages_with_list(self): + """Test conversion of list input to messages""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + input_list = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Question"}], + } + ] + + messages = ManagedResponsesWebSocketHandler._input_to_messages(input_list) + assert len(messages) == 1 + assert messages[0]["type"] == "message" + assert messages[0]["content"][0]["text"] == "Question" + + def test_input_to_messages_filters_non_dict_items(self): + """Test that non-dict items in list input are filtered""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + input_list = [ + "invalid_string", + None, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Valid"}], + }, + ] + + messages = ManagedResponsesWebSocketHandler._input_to_messages(input_list) + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == "Valid" + + def test_input_to_messages_handles_empty_input(self): + """Test that empty input returns empty list""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + assert ManagedResponsesWebSocketHandler._input_to_messages(None) == [] + assert ManagedResponsesWebSocketHandler._input_to_messages([]) == [] + assert ManagedResponsesWebSocketHandler._input_to_messages({}) == [] + + +class TestWebSocketEventTypes: + """Test that all WebSocket event types are properly handled with dict-based chunks""" + + def test_serialize_response_created_event_dict(self): + """Test serialization of response.created event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.created", + "response_id": "resp_123", + "response": { + "id": "resp_123", + "object": "response", + "status": "in_progress", + "created_at": 1234567890, + }, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.created" in serialized + assert "resp_123" in serialized + + def test_serialize_response_in_progress_event_dict(self): + """Test serialization of response.in_progress event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = {"type": "response.in_progress", "response_id": "resp_123"} + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.in_progress" in serialized + + def test_serialize_output_item_added_event_dict(self): + """Test serialization of response.output_item.added event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.output_item.added", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "item": {"type": "message", "role": "assistant"}, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.output_item.added" in serialized + assert "msg_456" in serialized + + def test_serialize_output_text_delta_event_dict(self): + """Test serialization of response.output_text.delta event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.output_text.delta", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "content_index": 0, + "delta": "Hello", + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.output_text.delta" in serialized + assert "Hello" in serialized + + def test_serialize_output_text_done_event_dict(self): + """Test serialization of response.output_text.done event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.output_text.done", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "content_index": 0, + "text": "Hello world", + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.output_text.done" in serialized + assert "Hello world" in serialized + + def test_serialize_content_part_done_event_dict(self): + """Test serialization of response.content_part.done event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.content_part.done", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "Complete text"}, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.content_part.done" in serialized + + def test_serialize_output_item_done_event_dict(self): + """Test serialization of response.output_item.done event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.output_item.done", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "item": {"type": "message", "role": "assistant", "status": "completed"}, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.output_item.done" in serialized + assert "msg_456" in serialized + + def test_serialize_response_completed_event_dict(self): + """Test serialization of response.completed event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.completed", + "response_id": "resp_123", + "response": { + "id": "resp_123", + "status": "completed", + "output": [ + { + "type": "message", + "content": [{"type": "output_text", "text": "Done"}], + } + ], + }, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.completed" in serialized + assert "resp_123" in serialized + + def test_serialize_response_failed_event_dict(self): + """Test serialization of response.failed event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.failed", + "response_id": "resp_123", + "response": { + "id": "resp_123", + "status": "failed", + "status_details": {"error": {"message": "Rate limit exceeded"}}, + }, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.failed" in serialized + assert "Rate limit exceeded" in serialized + + def test_serialize_response_incomplete_event_dict(self): + """Test serialization of response.incomplete event as dict""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.incomplete", + "response_id": "resp_123", + "response": { + "id": "resp_123", + "status": "incomplete", + "status_details": {"reason": "max_output_tokens"}, + }, + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.incomplete" in serialized + assert "max_output_tokens" in serialized + + +class TestMultiTurnSessionHistory: + """Test multi-turn conversation handling via session history""" + + def test_extract_output_messages_preserves_multiple_messages(self): + """Test that multiple output messages are all preserved""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "First message"}], + }, + { + "type": "function_call", + "id": "call_123", + "name": "get_weather", + "arguments": "{}", + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Second message"}], + }, + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 3 + assert messages[0]["content"][0]["text"] == "First message" + assert messages[1]["type"] == "function_call" + assert messages[2]["content"][0]["text"] == "Second message" + + def test_input_to_messages_with_mixed_content_types(self): + """Test input conversion with mixed content types""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + input_list = [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Question"}, + {"type": "input_image", "image_url": "https://example.com/img.png"}, + ], + } + ] + + messages = ManagedResponsesWebSocketHandler._input_to_messages(input_list) + assert len(messages) == 1 + assert len(messages[0]["content"]) == 2 + assert messages[0]["content"][0]["type"] == "input_text" + assert messages[0]["content"][1]["type"] == "input_image" + + def test_extract_output_messages_with_mixed_text_types(self): + """Test that both 'output_text' and 'text' types are extracted""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Part 1"}, + {"type": "text", "text": "Part 2"}, + ], + } + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == "Part 1Part 2" + + def test_extract_response_id_from_completed_event(self): + """Test extraction of response ID from completed event""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": {"id": "resp_abc123", "status": "completed"}, + } + + response_id = ManagedResponsesWebSocketHandler._extract_response_id( + completed_event + ) + assert response_id == "resp_abc123" + + def test_extract_response_id_handles_missing_response(self): + """Test that missing response dict returns None""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = {"type": "response.completed"} + + response_id = ManagedResponsesWebSocketHandler._extract_response_id( + completed_event + ) + assert response_id is None + + +class TestWebSocketErrorHandling: + """Test error handling in WebSocket mode""" + + @pytest.mark.asyncio + async def test_managed_handler_handles_invalid_json(self): + """Test that invalid JSON in response.create is handled gracefully""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.recv = AsyncMock(return_value="invalid json {{{") + + mock_logging_obj = Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ) + + handler = ManagedResponsesWebSocketHandler( + websocket=mock_websocket, + model="test-model", + logging_obj=mock_logging_obj, + ) + + # Process invalid JSON + await handler._process_response_create("invalid json {{{") + + # Should have sent an error event + mock_websocket.send_text.assert_called_once() + error_event = mock_websocket.send_text.call_args[0][0] + assert "error" in error_event + assert "Invalid JSON" in error_event + + +class TestWebSocketChunkTypes: + """Test handling of different chunk types from streaming responses""" + + def test_serialize_function_call_chunk(self): + """Test serialization of function call chunks""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.function_call.added", + "response_id": "resp_123", + "item_id": "call_456", + "output_index": 0, + "call_id": "call_456", + "name": "get_weather", + "arguments": "", + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.function_call.added" in serialized + assert "get_weather" in serialized + + def test_serialize_function_call_arguments_delta(self): + """Test serialization of function call arguments delta""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.function_call_arguments.delta", + "response_id": "resp_123", + "item_id": "call_456", + "output_index": 0, + "call_id": "call_456", + "delta": '{"location"', + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.function_call_arguments.delta" in serialized + assert "location" in serialized + + def test_serialize_function_call_arguments_done(self): + """Test serialization of function call arguments done""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.function_call_arguments.done", + "response_id": "resp_123", + "item_id": "call_456", + "output_index": 0, + "call_id": "call_456", + "arguments": '{"location": "Paris"}', + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.function_call_arguments.done" in serialized + assert "Paris" in serialized + + def test_serialize_reasoning_content_delta(self): + """Test serialization of reasoning content delta""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.reasoning_content.delta", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "content_index": 0, + "delta": "Thinking step 1...", + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.reasoning_content.delta" in serialized + assert "Thinking step 1" in serialized + + def test_serialize_reasoning_content_done(self): + """Test serialization of reasoning content done""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + chunk = { + "type": "response.reasoning_content.done", + "response_id": "resp_123", + "item_id": "msg_456", + "output_index": 0, + "content_index": 0, + "reasoning_content": "Complete reasoning...", + } + + serialized = ManagedResponsesWebSocketHandler._serialize_chunk(chunk) + assert serialized is not None + assert "response.reasoning_content.done" in serialized + assert "Complete reasoning" in serialized + + def test_extract_output_messages_preserves_multiple_messages(self): + """Test that multiple output messages are all preserved""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "First message"}], + }, + { + "type": "function_call", + "id": "call_123", + "name": "get_weather", + "arguments": "{}", + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Second message"}], + }, + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 3 + assert messages[0]["content"][0]["text"] == "First message" + assert messages[1]["type"] == "function_call" + assert messages[2]["content"][0]["text"] == "Second message" + + def test_input_to_messages_with_mixed_content_types(self): + """Test input conversion with mixed content types""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + input_list = [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Question"}, + {"type": "input_image", "image_url": "https://example.com/img.png"}, + ], + } + ] + + messages = ManagedResponsesWebSocketHandler._input_to_messages(input_list) + assert len(messages) == 1 + assert len(messages[0]["content"]) == 2 + assert messages[0]["content"][0]["type"] == "input_text" + assert messages[0]["content"][1]["type"] == "input_image" + + def test_extract_output_messages_with_mixed_text_types(self): + """Test that both 'output_text' and 'text' types are extracted""" + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Part 1"}, + {"type": "text", "text": "Part 2"}, + ], + } + ], + }, + } + + messages = ManagedResponsesWebSocketHandler._extract_output_messages( + completed_event + ) + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == "Part 1Part 2"