import os import sys from unittest.mock import AsyncMock, Mock, patch import pytest sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler handler = BaseLLMHTTPHandler() # Test case 1: fake_stream is True stream = True data = { "stream": True, "model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], } fake_stream = True result_stream, result_data = handler._prepare_fake_stream_request( stream=stream, data=data, fake_stream=fake_stream ) # Verify that stream is set to False assert result_stream is False # Verify that "stream" key is removed from data assert "stream" not in result_data # Verify other data remains unchanged assert result_data["model"] == "gpt-4" assert result_data["messages"] == [{"role": "user", "content": "Hello"}] # Test case 2: fake_stream is False stream = True data = { "stream": True, "model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], } fake_stream = False result_stream, result_data = handler._prepare_fake_stream_request( stream=stream, data=data, fake_stream=fake_stream ) # Verify that stream remains True assert result_stream is True # Verify that data remains unchanged assert "stream" in result_data assert result_data["stream"] is True assert result_data["model"] == "gpt-4" assert result_data["messages"] == [{"role": "user", "content": "Hello"}] # Test case 3: data doesn't have stream key but fake_stream is True stream = True data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]} fake_stream = True result_stream, result_data = handler._prepare_fake_stream_request( stream=stream, data=data, fake_stream=fake_stream ) # Verify that stream is set to False assert result_stream is False # Verify that data remains unchanged (since there was no stream key to remove) assert "stream" not in result_data assert result_data["model"] == "gpt-4" assert result_data["messages"] == [{"role": "user", "content": "Hello"}] @pytest.mark.asyncio async def test_async_anthropic_messages_handler_extra_headers(): """ Test that async_anthropic_messages_handler correctly extracts and merges extra_headers from kwargs with proper priority. """ handler = BaseLLMHTTPHandler() # Mock the config mock_config = Mock() mock_config.validate_anthropic_messages_environment = Mock( return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com") ) mock_config.transform_anthropic_messages_request = Mock( return_value={"model": "claude-3-opus-20240229", "messages": []} ) # Mock the client mock_client = AsyncMock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "id": "msg_123", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], "model": "claude-3-opus-20240229", "stop_reason": "end_turn", } mock_client.post = AsyncMock(return_value=mock_response) # Mock logging object mock_logging_obj = Mock() mock_logging_obj.update_environment_variables = Mock() mock_logging_obj.model_call_details = {} mock_logging_obj.stream = False # Test case 1: Only extra_headers in kwargs kwargs = { "extra_headers": { "X-Custom-Header": "from-kwargs", "X-Auth-Token": "token123", } } with patch( "litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers" ) as mock_provider_headers: mock_provider_headers.return_value = None # Capture what headers are passed to validate_anthropic_messages_environment captured_headers = {} def capture_validate(*args, **kwargs): captured_headers.update(kwargs.get("headers", {})) return ({"x-api-key": "test-key"}, "https://api.anthropic.com") mock_config.validate_anthropic_messages_environment = capture_validate try: await handler.async_anthropic_messages_handler( model="claude-3-opus-20240229", messages=[{"role": "user", "content": "Hello"}], anthropic_messages_provider_config=mock_config, anthropic_messages_optional_request_params={}, custom_llm_provider="anthropic", litellm_params=GenericLiteLLMParams(), logging_obj=mock_logging_obj, client=mock_client, kwargs=kwargs, ) except Exception: pass # We're testing header extraction, not the full flow # Verify extra_headers were extracted and merged assert "X-Custom-Header" in captured_headers assert captured_headers["X-Custom-Header"] == "from-kwargs" assert "X-Auth-Token" in captured_headers assert captured_headers["X-Auth-Token"] == "token123" @pytest.mark.asyncio async def test_async_anthropic_messages_handler_header_priority(): """ Test that async_anthropic_messages_handler respects header priority: forwarded < extra_headers < provider_specific """ handler = BaseLLMHTTPHandler() # Mock the config mock_config = Mock() mock_client = AsyncMock() mock_logging_obj = Mock() mock_logging_obj.update_environment_variables = Mock() mock_logging_obj.model_call_details = {} mock_logging_obj.stream = False # Test with all three header sources kwargs = { "headers": {"X-Priority": "forwarded", "X-Forwarded-Only": "keep"}, "extra_headers": {"X-Priority": "extra", "X-Extra-Only": "also-keep"}, } with patch( "litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers" ) as mock_provider_headers: mock_provider_headers.return_value = { "X-Priority": "provider", "X-Provider-Only": "keep-this-too" } captured_headers = {} def capture_validate(*args, **kwargs): captured_headers.update(kwargs.get("headers", {})) return ({"x-api-key": "test-key"}, "https://api.anthropic.com") mock_config.validate_anthropic_messages_environment = capture_validate mock_config.transform_anthropic_messages_request = Mock( return_value={"model": "claude-3-opus-20240229", "messages": []} ) try: await handler.async_anthropic_messages_handler( model="claude-3-opus-20240229", messages=[{"role": "user", "content": "Hello"}], anthropic_messages_provider_config=mock_config, anthropic_messages_optional_request_params={}, custom_llm_provider="anthropic", litellm_params=GenericLiteLLMParams(), logging_obj=mock_logging_obj, client=mock_client, kwargs=kwargs, ) except Exception: pass # Verify priority: provider_specific should win assert captured_headers["X-Priority"] == "provider" # Verify all unique headers from different sources are present assert captured_headers["X-Forwarded-Only"] == "keep" assert captured_headers["X-Extra-Only"] == "also-keep" assert captured_headers["X-Provider-Only"] == "keep-this-too"