Test email body
" # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) + mock_response = mock.Mock(spec=Response) + mock_response.raise_for_status.return_value = None + mock_response.status_code = 200 mock_response.json.return_value = {"id": "test_email_id"} mock_httpx_client.post.return_value = mock_response @@ -107,7 +125,13 @@ async def test_send_email_missing_api_key(mock_httpx_client): @pytest.mark.asyncio +@respx.mock async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): + # Block all HTTP requests at network level to prevent real API calls + respx.post("https://api.resend.com/emails").mock( + return_value=httpx.Response(200, json={"id": "test_email_id"}) + ) + # Initialize the logger logger = ResendEmailLogger() @@ -118,7 +142,9 @@ async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): html_body = "Test email body
" # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) + mock_response = mock.Mock(spec=Response) + mock_response.raise_for_status.return_value = None + mock_response.status_code = 200 mock_response.json.return_value = {"id": "test_email_id"} mock_httpx_client.post.return_value = mock_response diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index 2b0bb31751c..836b717bd6e 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -2,7 +2,9 @@ import os import sys import unittest.mock as mock +import httpx import pytest +import respx from httpx import Response sys.path.insert(0, os.path.abspath("../../..")) @@ -14,8 +16,26 @@ from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( @pytest.fixture def mock_env_vars(): - with mock.patch.dict(os.environ, {"SENDGRID_API_KEY": "test_api_key"}): + # Store original values + original_api_key = os.environ.get("SENDGRID_API_KEY") + original_sender_email = os.environ.get("SENDGRID_SENDER_EMAIL") + + # Set test API key and remove SENDGRID_SENDER_EMAIL to ensure isolation + os.environ["SENDGRID_API_KEY"] = "test_api_key" + if "SENDGRID_SENDER_EMAIL" in os.environ: + del os.environ["SENDGRID_SENDER_EMAIL"] + + try: yield + finally: + # Restore original values + if original_api_key is not None: + os.environ["SENDGRID_API_KEY"] = original_api_key + elif "SENDGRID_API_KEY" in os.environ: + del os.environ["SENDGRID_API_KEY"] + + if original_sender_email is not None: + os.environ["SENDGRID_SENDER_EMAIL"] = original_sender_email @pytest.fixture @@ -23,14 +43,16 @@ def mock_httpx_client(): with mock.patch( "litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email.get_async_httpx_client" ) as mock_client: - mock_response = mock.AsyncMock(spec=Response) + + mock_response = mock.Mock(spec=Response) mock_response.status_code = 202 mock_response.text = "accepted" + mock_response.raise_for_status.return_value = None mock_async_client = mock.AsyncMock() mock_async_client.post.return_value = mock_response - mock_client.return_value = mock_async_client + mock_client.return_value = mock_async_client yield mock_async_client @@ -62,18 +84,12 @@ async def test_send_email_success(mock_env_vars, mock_httpx_client): @pytest.mark.asyncio -async def test_send_email_missing_api_key(mock_httpx_client): - # Remove the API key from environment before initializing logger +async def test_send_email_missing_api_key(): original_key = os.environ.pop("SENDGRID_API_KEY", None) - + try: logger = SendGridEmailLogger() - # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) - mock_response.status_code = 401 - mock_httpx_client.post.return_value = mock_response - with pytest.raises(ValueError): await logger.send_email( from_email="test@example.com", @@ -81,16 +97,19 @@ async def test_send_email_missing_api_key(mock_httpx_client): subject="Test Subject", html_body="Test email body
", ) - - mock_httpx_client.post.assert_not_called() finally: - # Restore the original key if it existed if original_key is not None: os.environ["SENDGRID_API_KEY"] = original_key @pytest.mark.asyncio +@respx.mock async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): + # Block all HTTP requests at network level to prevent real API calls + respx.post("https://api.sendgrid.com/v3/mail/send").mock( + return_value=httpx.Response(202, text="accepted") + ) + logger = SendGridEmailLogger() from_email = "test@example.com" @@ -98,10 +117,10 @@ async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client): subject = "Test Subject" html_body = "Test email body
" - # Mock the response to avoid making real HTTP requests - mock_response = mock.AsyncMock(spec=Response) + mock_response = mock.Mock(spec=Response) mock_response.status_code = 202 mock_response.text = "accepted" + mock_response.raise_for_status.return_value = None mock_httpx_client.post.return_value = mock_response await logger.send_email( diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py new file mode 100644 index 00000000000..be2084969a5 --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py @@ -0,0 +1,169 @@ +import os +import time +from unittest.mock import AsyncMock + +import pytest +from httpx import Response + +from litellm.integrations.datadog.datadog_cost_management import ( + DatadogCostManagementLogger, +) +from litellm.types.utils import StandardLoggingPayload + + +@pytest.fixture +def clean_env(): + # Save original env + original_api_key = os.environ.get("DD_API_KEY") + original_app_key = os.environ.get("DD_APP_KEY") + original_site = os.environ.get("DD_SITE") + + # Set test env + os.environ["DD_API_KEY"] = "test_api_key" + os.environ["DD_APP_KEY"] = "test_app_key" + os.environ["DD_SITE"] = "test.datadoghq.com" + + yield + + # Restore original env + if original_api_key: + os.environ["DD_API_KEY"] = original_api_key + else: + del os.environ["DD_API_KEY"] + + if original_app_key: + os.environ["DD_APP_KEY"] = original_app_key + else: + del os.environ["DD_APP_KEY"] + + if original_site: + os.environ["DD_SITE"] = original_site + else: + del os.environ["DD_SITE"] + + +@pytest.mark.asyncio +async def test_init(clean_env): + """ + Test initialization sets up clients and url correctly + """ + logger = DatadogCostManagementLogger() + assert logger.dd_api_key == "test_api_key" + assert logger.dd_app_key == "test_app_key" + assert ( + logger.upload_url == "https://api.test.datadoghq.com/api/v2/cost/custom_costs" + ) + + +@pytest.mark.asyncio +async def test_aggregate_costs(clean_env): + """ + Test that costs are correctly aggregated by provider, model, and date + """ + logger = DatadogCostManagementLogger() + + # Mock some log payloads + now = time.time() + day_str = time.strftime("%Y-%m-%d", time.localtime(now)) + + logs = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=now, + metadata={"user_api_key_team_alias": "team-a"}, + ), + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.02, + startTime=now, + metadata={"user_api_key_team_alias": "team-a"}, + ), + StandardLoggingPayload( + custom_llm_provider="anthropic", + model="claude-3", + response_cost=0.05, + startTime=now, + ), + ] + + aggregated = logger._aggregate_costs(logs) + + assert len(aggregated) == 2 + + # Check OpenAI entry + openai_entry = next(e for e in aggregated if e["ProviderName"] == "openai") + assert openai_entry["BilledCost"] == 0.03 + assert openai_entry["ChargeDescription"] == "LLM Usage for gpt-4" + assert openai_entry["ChargePeriodStart"] == day_str + assert openai_entry["Tags"]["team"] == "team-a" + assert "env" in openai_entry["Tags"] + assert "service" in openai_entry["Tags"] + + # Check Anthropic entry + anthropic_entry = next(e for e in aggregated if e["ProviderName"] == "anthropic") + assert anthropic_entry["BilledCost"] == 0.05 + + +@pytest.mark.asyncio +async def test_async_log_success_event(clean_env): + """ + Test that logs are added to queue + """ + logger = DatadogCostManagementLogger(batch_size=10) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": {"response_cost": 0.01}}, + response_obj={}, + start_time=time.time(), + end_time=time.time(), + ) + + assert len(logger.log_queue) == 1 + assert logger.log_queue[0]["response_cost"] == 0.01 + + # Test zero cost ignored + await logger.async_log_success_event( + kwargs={"standard_logging_object": {"response_cost": 0.0}}, + response_obj={}, + start_time=time.time(), + end_time=time.time(), + ) + + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_async_send_batch(clean_env): + """ + Test that batch is aggregated and uploaded + """ + logger = DatadogCostManagementLogger() + logger.async_client = AsyncMock() + logger.async_client.put.return_value = Response(202, json={"status": "ok"}) + + # Add logs directly to queue + logger.log_queue = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + ) + ] + + await logger.async_send_batch() + + # Verify API called + assert logger.async_client.put.called + call_args = logger.async_client.put.call_args + assert call_args[0][0] == "https://api.test.datadoghq.com/api/v2/cost/custom_costs" + + import json + + # Use call_args.kwargs['content'] + content = json.loads(call_args[1]["content"]) + assert content[0]["ProviderName"] == "openai" + assert content[0]["BilledCost"] == 0.01 diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py new file mode 100644 index 00000000000..2bb51e1e1b7 --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py @@ -0,0 +1,62 @@ +import os +from unittest.mock import patch +from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger + + +def test_datadog_llm_obs_agent_configuration(): + """ + Test that DataDog LLM Obs logger correctly configures agent endpoint. + """ + test_env = { + "LITELLM_DD_AGENT_HOST": "localhost", + "LITELLM_DD_LLM_OBS_PORT": "10518", + "DD_API_KEY": "test-api-key", # Optional, but checking if it's preserved + } + + # Ensure DD_SITE is NOT set to verify we don't need it in agent mode + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): # Prevent periodic flush task from running + dd_logger = DataDogLLMObsLogger() + + expected_url = "http://localhost:10518/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + assert dd_logger.DD_API_KEY == "test-api-key" + + +def test_datadog_llm_obs_agent_no_api_key_ok(): + """ + Test that agent mode works WITHOUT DD_API_KEY (agent handles auth). + """ + test_env = { + "LITELLM_DD_AGENT_HOST": "localhost", + # No DD_API_KEY + } + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): + # Should NOT raise exception anymore + dd_logger = DataDogLLMObsLogger() + + assert dd_logger.DD_API_KEY is None + # Default port is 8126 if not set + expected_url = "http://localhost:8126/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + + +def test_datadog_llm_obs_direct_api_configuration(): + """ + Test that direct API configuration still works as expected. + """ + test_env = { + "DD_API_KEY": "direct-api-key", + "DD_SITE": "us5.datadoghq.com", + } + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): + dd_logger = DataDogLLMObsLogger() + + expected_url = "https://api.us5.datadoghq.com/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + assert dd_logger.DD_API_KEY == "direct-api-key" diff --git a/tests/test_litellm/integrations/test_custom_guardrail_recursion.py b/tests/test_litellm/integrations/test_custom_guardrail_recursion.py new file mode 100644 index 00000000000..f05b5848bdf --- /dev/null +++ b/tests/test_litellm/integrations/test_custom_guardrail_recursion.py @@ -0,0 +1,73 @@ +import pytest +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +import json + + +class TestCustomGuardrailRecursion: + """ + Specific tests for the circular reference / RecursionError fix in logging. + """ + + def test_log_guardrail_information_handles_circular_references(self): + """ + Test that add_standard_logging method sanitizes input data containing circular references + instead of crashing. + + This reproduces the Langfuse crash scenario: + Request -> Metadata -> GuardrailResponse -> DebugContext -> Request + """ + guardrail = CustomGuardrail( + guardrail_name="recursion_test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + # 1. Setup Circular Data + request_data = {"user_id": "test_recursive_user"} + metadata = {"session_id": "123"} + request_data["metadata"] = metadata + + # Create the danger: Guardrail Response holding a reference back to request_data + dirty_response = { + "flagged": False, + "debug_context": request_data, # <--- ACCESS TO ROOT (Circular Ref) + } + + # 2. Invoke the logging method + # If the fix is working, this will NOT raise RecursionError + try: + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=dirty_response, + request_data=request_data, + guardrail_status="success", + start_time=1.0, + end_time=2.0, + duration=1.0, + masked_entity_count={}, + event_type=GuardrailEventHooks.pre_call, + ) + except RecursionError: + pytest.fail( + "RecursionError raised! The cyclic reference sanitization failed." + ) + + # 3. Verify the data stored is safe + stored_info = request_data["metadata"][ + "standard_logging_guardrail_information" + ][0] + stored_response = stored_info["guardrail_response"] + + # Check that we can dump it to JSON without crashing (Ultimate proof) + try: + json.dumps(stored_response) + except Exception as e: + pytest.fail(f"Stored data is not JSON serializable: {e}") + + # Check content - keys should be preserved but recursion broken + assert "debug_context" in stored_response + debug_context = stored_response["debug_context"] + + # In a sanitized copy, the nested metadata should be a copy, not the original live dict + assert debug_context["user_id"] == "test_recursive_user" + # The 'metadata' inside 'debug_context' would be where recursion stops or is filtered + assert "metadata" in debug_context diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index a22fe13798f..e87233a52a3 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1138,6 +1138,73 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["name"].endswith(f"_{format_type}") assert block["document"]["format"] == format_type +def test_bedrock_nova_web_search_options_mapping(): + """ + Test that web_search_options is correctly mapped to Nova grounding. + + This follows the LiteLLM pattern for web search where: + - Vertex AI maps web_search_options to {"googleSearch": {}} + - Anthropic maps web_search_options to {"type": "web_search_20250305", ...} + - Nova should map web_search_options to {"systemTool": {"name": "nova_grounding"}} + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + config = AmazonConverseConfig() + + # Test basic mapping for Nova model + result = config._map_web_search_options({}, "amazon.nova-pro-v1:0") + + assert result is not None + system_tool = result.get("systemTool") + assert system_tool is not None + assert system_tool["name"] == "nova_grounding" + + # Test with search_context_size (should be ignored for Nova) + result2 = config._map_web_search_options( + {"search_context_size": "high"}, + "us.amazon.nova-premier-v1:0" + ) + + assert result2 is not None + system_tool2 = result2.get("systemTool") + assert system_tool2 is not None + assert system_tool2["name"] == "nova_grounding" + # Nova doesn't support search_context_size, so it's just ignored + +def test_bedrock_tools_pt_does_not_handle_system_tool(): + """ + Verify that _bedrock_tools_pt does NOT handle system_tool format. + + System tools (nova_grounding) should be added via web_search_options, + not via the tools parameter directly. + """ + + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt + + # Regular function tools should still work + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"} + }, + "required": ["location"] + } + } + } + ] + + result = _bedrock_tools_pt(tools=tools) + + assert len(result) == 1 + tool_spec = result[0].get("toolSpec") + assert tool_spec is not None + assert tool_spec["name"] == "get_weather" def test_convert_to_anthropic_tool_result_image_with_cache_control(): """ @@ -1305,12 +1372,12 @@ def test_convert_to_anthropic_tool_result_image_url_as_http(): assert result["content"][0]["cache_control"]["type"] == "ephemeral" def test_anthropic_messages_pt_server_tool_use_passthrough(): """ - Test that anthropic_messages_pt passes through server_tool_use and + Test that anthropic_messages_pt passes through server_tool_use and tool_search_tool_result blocks in assistant message content. - + These are Anthropic-native content types used for tool search functionality that need to be preserved when reconstructing multi-turn conversations. - + Fixes: https://github.com/BerriAI/litellm/issues/XXXXX """ from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt @@ -1359,15 +1426,15 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): # Verify we have 3 messages (user, assistant, user) assert len(result) == 3 - + # Verify the assistant message content assistant_msg = result[1] assert assistant_msg["role"] == "assistant" assert isinstance(assistant_msg["content"], list) - + # Find the different content block types content_types = [block.get("type") for block in assistant_msg["content"]] - + # Verify server_tool_use block is preserved assert "server_tool_use" in content_types server_tool_use_block = next( @@ -1376,7 +1443,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): assert server_tool_use_block["id"] == "srvtoolu_01ABC123" assert server_tool_use_block["name"] == "tool_search_tool_regex" assert server_tool_use_block["input"] == {"query": ".*time.*"} - + # Verify tool_search_tool_result block is preserved assert "tool_search_tool_result" in content_types tool_result_block = next( @@ -1385,7 +1452,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): assert tool_result_block["tool_use_id"] == "srvtoolu_01ABC123" assert tool_result_block["content"]["type"] == "tool_search_tool_search_result" assert tool_result_block["content"]["tool_references"][0]["tool_name"] == "get_time" - + # Verify text block is also preserved assert "text" in content_types text_block = next( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index e035e193fe1..1f3f558a498 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -787,11 +787,13 @@ def test_get_masked_values(): "presidio_ad_hoc_recognizers": None, "aws_bedrock_runtime_endpoint": None, "presidio_anonymizer_api_base": None, + "vertex_credentials": "{sensitive_api_key}", } masked_values = _get_masked_values( sensitive_object, unmasked_length=4, number_of_asterisks=4 ) assert masked_values["presidio_anonymizer_api_base"] is None + assert masked_values["vertex_credentials"] == "{s****y}" @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 6aadbc058d1..c26d057fbf1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1108,3 +1108,309 @@ def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238(): assert block_type == "tool_use" assert content_block_start["name"] == "Bash" assert content_block_start["id"] == "toolu_bdrk_013xRVejhv3ybmLEGCoZib2b" + + +# ============================================================================ +# Cache Control Transformation Tests +# ============================================================================ + +# Model constant for cache control tests +CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0" +CACHE_CONTROL_NON_ANTHROPIC_MODEL = "gpt-4" + + +def test_should_add_cache_control_for_anthropic_model(): + """Should add cache_control to target for Anthropic Claude models.""" + adapter = LiteLLMAnthropicMessagesAdapter() + cache_control = {"type": "ephemeral"} + + for model in [ + CACHE_CONTROL_BEDROCK_CONVERSE_MODEL, + "anthropic/claude-sonnet-4-5", + "claude-opus-4-5-20251101", + "vertex_ai/claude-3-sonnet@20240229", + ]: + target = {} + adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model) + assert "cache_control" in target + assert target["cache_control"] == cache_control + + +def test_should_not_add_cache_control_for_non_anthropic_model(): + """Should not add cache_control for non-Anthropic models.""" + adapter = LiteLLMAnthropicMessagesAdapter() + cache_control = {"type": "ephemeral"} + + for model in [CACHE_CONTROL_NON_ANTHROPIC_MODEL, "openai/gpt-4-turbo", "gemini-pro"]: + target = {} + adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model) + assert "cache_control" not in target + + +def test_should_not_add_cache_control_when_none(): + """Should not add cache_control when source has None or empty cache_control.""" + adapter = LiteLLMAnthropicMessagesAdapter() + + for source in [{"cache_control": None}, {"cache_control": {}}, {"cache_control": ""}, {}]: + target = {} + adapter._add_cache_control_if_applicable(source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL) + assert "cache_control" not in target + + +def test_should_not_add_cache_control_when_model_none(): + """Should not add cache_control when model is None or empty.""" + adapter = LiteLLMAnthropicMessagesAdapter() + cache_control = {"type": "ephemeral"} + + for model in [None, ""]: + target = {} + adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model) + assert "cache_control" not in target + + +def test_cache_control_preserved_in_text_content_for_claude(): + """Cache control should be preserved in text content for Claude models.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "text", + "text": "This is cached content", + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + assert len(result) == 1 + assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_not_preserved_for_non_claude_model(): + """Cache control should NOT be preserved for non-Claude models.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "text", + "text": "This is cached content", + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_NON_ANTHROPIC_MODEL + ) + + assert len(result) == 1 + assert "cache_control" not in result[0]["content"][0] + + +def test_cache_control_preserved_in_image_content_for_claude(): + """Cache control should be preserved in image content for Claude models.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==", + }, + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + assert len(result) == 1 + assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_preserved_in_document_content_for_claude(): + """Cache control should be preserved in document content for Claude models.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "JVBERi0xLjQKJeLjz9MK", + }, + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + assert len(result) == 1 + assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_preserved_in_tool_result_for_claude(): + """Cache control should be preserved in tool_result for Claude models.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "tool_result", + "tool_use_id": "toolu_01234", + "content": "Tool result content", + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + tool_message = next(msg for msg in result if msg.get("role") == "tool") + assert tool_message["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_not_preserved_in_tool_result_for_non_claude(): + """Cache control should NOT be preserved in tool_result for non-Claude models.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "tool_result", + "tool_use_id": "toolu_01234", + "content": "Tool result content", + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_NON_ANTHROPIC_MODEL + ) + + tool_message = next(msg for msg in result if msg.get("role") == "tool") + assert "cache_control" not in tool_message + + +def test_cache_control_preserved_in_assistant_text_for_claude(): + """Cache control should be preserved in assistant text blocks for Claude models.""" + anthropic_messages = [ + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + { + "type": "text", + "text": "Assistant response", + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + assert len(result) == 1 + assert result[0]["role"] == "assistant" + # When cache_control is present, content should be a list + assert isinstance(result[0]["content"], list) + assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_preserved_in_tool_use_for_claude(): + """Cache control should be preserved in tool_use blocks for Claude models.""" + anthropic_messages = [ + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + { + "type": "tool_use", + "id": "toolu_01234", + "name": "get_weather", + "input": {"location": "Boston"}, + "cache_control": {"type": "ephemeral"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + assert len(result) == 1 + assert "tool_calls" in result[0] + assert result[0]["tool_calls"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_preserved_in_tools_for_claude(): + """Cache control should be preserved in tools for Claude models.""" + tools = [ + { + "name": "get_weather", + "description": "Get weather for a location", + "input_schema": {"type": "object", "properties": {"location": {"type": "string"}}}, + "cache_control": {"type": "ephemeral"}, + } + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_tools_to_openai( + tools=tools, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL + ) + + assert len(result) == 1 + assert result[0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_not_preserved_in_tools_for_non_claude(): + """Cache control should NOT be preserved in tools for non-Claude models.""" + tools = [ + { + "name": "get_weather", + "description": "Get weather for a location", + "input_schema": {"type": "object", "properties": {"location": {"type": "string"}}}, + "cache_control": {"type": "ephemeral"}, + } + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_tools_to_openai( + tools=tools, model=CACHE_CONTROL_NON_ANTHROPIC_MODEL + ) + + assert len(result) == 1 + assert "cache_control" not in result[0] diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py b/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py b/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py new file mode 100644 index 00000000000..af1e1df92fd --- /dev/null +++ b/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py @@ -0,0 +1,218 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.vercel_ai_gateway.embedding.transformation import ( + VercelAIGatewayEmbeddingConfig, +) +from litellm.llms.vercel_ai_gateway.common_utils import VercelAIGatewayException +from litellm.types.utils import EmbeddingResponse + + +def test_vercel_ai_gateway_embedding_get_complete_url(): + """Test URL generation for embeddings endpoint""" + config = VercelAIGatewayEmbeddingConfig() + + # Test with default API base + url = config.get_complete_url( + api_base=None, + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://ai-gateway.vercel.sh/v1/embeddings" + + # Test with custom API base + url = config.get_complete_url( + api_base="https://custom.vercel.sh/v1", + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.vercel.sh/v1/embeddings" + + # Test with trailing slash + url = config.get_complete_url( + api_base="https://custom.vercel.sh/v1/", + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.vercel.sh/v1/embeddings" + + +def test_vercel_ai_gateway_embedding_transform_request(): + """Test request transformation for embeddings""" + config = VercelAIGatewayEmbeddingConfig() + + # Test with string input + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input="Hello world", + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello world"] + + # Test with list input + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input=["Hello", "World"], + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello", "World"] + + # Test stripping vercel_ai_gateway/ prefix + request = config.transform_embedding_request( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input="Hello", + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + + +def test_vercel_ai_gateway_embedding_transform_request_with_dimensions(): + """Test request transformation with dimensions parameter""" + config = VercelAIGatewayEmbeddingConfig() + + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input="Hello world", + optional_params={"dimensions": 768}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello world"] + assert request["dimensions"] == 768 + + +def test_vercel_ai_gateway_embedding_validate_environment(): + """Test header validation and setup""" + config = VercelAIGatewayEmbeddingConfig() + + headers = config.validate_environment( + headers={}, + model="openai/text-embedding-3-small", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_key", + ) + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test_key" + + # Test with existing headers (should merge) + headers = config.validate_environment( + headers={"X-Custom": "value"}, + model="openai/text-embedding-3-small", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_key", + ) + assert headers["X-Custom"] == "value" + assert headers["Authorization"] == "Bearer test_key" + + +def test_vercel_ai_gateway_embedding_get_supported_params(): + """Test supported OpenAI parameters""" + config = VercelAIGatewayEmbeddingConfig() + supported = config.get_supported_openai_params("openai/text-embedding-3-small") + + assert "dimensions" in supported + assert "encoding_format" in supported + assert "timeout" in supported + assert "user" in supported + + +def test_vercel_ai_gateway_embedding_map_openai_params(): + """Test OpenAI parameter mapping""" + config = VercelAIGatewayEmbeddingConfig() + + optional_params = config.map_openai_params( + non_default_params={"dimensions": 768, "encoding_format": "float"}, + optional_params={}, + model="openai/text-embedding-3-small", + drop_params=False, + ) + assert optional_params["dimensions"] == 768 + assert optional_params["encoding_format"] == "float" + + +def test_vercel_ai_gateway_embedding_error_class(): + """Test error class creation""" + config = VercelAIGatewayEmbeddingConfig() + + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, VercelAIGatewayException) + assert error.message == "Test error" + assert error.status_code == 400 + + +def test_vercel_ai_gateway_embedding_transform_response(): + """Test response transformation""" + config = VercelAIGatewayEmbeddingConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.text = '{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}],"model":"openai/text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' + mock_response.json.return_value = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "openai/text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + + mock_logging = MagicMock() + + response = config.transform_embedding_response( + model="openai/text-embedding-3-small", + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=mock_logging, + api_key="test_key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + assert response is not None + mock_logging.post_call.assert_called_once() + + +def test_vercel_ai_gateway_embedding_env_vars(): + """Test environment variable handling""" + config = VercelAIGatewayEmbeddingConfig() + + with patch.dict( + os.environ, + { + "VERCEL_AI_GATEWAY_API_BASE": "https://env.vercel.sh/v1", + }, + ): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://env.vercel.sh/v1/embeddings" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ecdc75ede52..f241b2aa0d5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1885,7 +1885,7 @@ class TestMCPServerManager: # Create mock client that tracks call_tool usage mock_client = AsyncMock() - async def mock_call_tool(params): + async def mock_call_tool(params, host_progress_callback=None): # Return a mock CallToolResult result = MagicMock(spec=CallToolResult) result.content = [{"type": "text", "text": "Tool executed successfully"}] diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 807559207e6..3df0dc881e8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -15,10 +15,10 @@ import pytest import litellm from litellm.proxy._types import ( CallInfo, + Litellm_EntityType, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, - Litellm_EntityType, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -131,6 +131,60 @@ def test_get_key_object_from_ui_hash_key_invalid(): assert key_object is None +def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values): + """Test generating CLI JWT token with default 24-hour expiration""" + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + # Decrypt and verify token contents + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data["user_id"] == "test_user" + assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value + assert token_data["models"] == ["gpt-3.5-turbo"] + assert token_data["max_budget"] == litellm.max_ui_session_budget + + # Verify expiration time is set to 24 hours (default) + assert "expires" in token_data + expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) + assert expires > get_utc_datetime() + assert expires <= get_utc_datetime() + timedelta(hours=24, minutes=1) + assert expires >= get_utc_datetime() + timedelta(hours=23, minutes=59) + + +def test_get_cli_jwt_auth_token_custom_expiration( + valid_sso_user_defined_values, monkeypatch +): + """Test generating CLI JWT token with custom expiration via environment variable""" + # Set custom expiration to 48 hours + monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "48") + + # Reload the constants module to pick up the new env var + import importlib + + from litellm import constants + importlib.reload(constants) + + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + # Decrypt and verify token contents + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + # Verify expiration time is set to 48 hours + assert "expires" in token_data + expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) + assert expires > get_utc_datetime() + timedelta(hours=47, minutes=59) + assert expires <= get_utc_datetime() + timedelta(hours=48, minutes=1) + + + @pytest.mark.asyncio async def test_default_internal_user_params_with_get_user_object(monkeypatch): """Test that default_internal_user_params is used when creating a new user via get_user_object""" diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index e7b27908c14..a0e29e06100 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -626,3 +626,133 @@ async def test_expire_previous_ui_session_tokens_exception_handling(): # Should not raise exception despite database error await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + +@pytest.mark.asyncio +async def test_authenticate_user_admin_login_with_non_ascii_characters(): + """Test admin login with non-ASCII characters in password (issue #19559)""" + master_key = "sk-1234" + ui_username = "admin£test" + ui_password = "sk-1234£pass" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": ui_password, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + ): + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = { + "token": "test-token-123", + "user_id": LITELLM_PROXY_ADMIN_NAME, + } + + with patch( + "litellm.proxy.auth.login_utils.user_update", + new_callable=AsyncMock, + return_value=None, + ) as mock_user_update: + with patch( + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ): + result = await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + assert result.key == "test-token-123" + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + +def test_authenticate_user_non_ascii_direct_comparison(): + """Test that non-ASCII characters can be compared directly (unit test for fix)""" + import secrets + + # This test verifies the fix handles non-ASCII by encoding to bytes + username = "admin£test" + password = "pass£word" + + # This would fail without encoding: + # secrets.compare_digest(username, username) # TypeError! + + # But works with the fix: + result = secrets.compare_digest( + username.encode("utf-8"), username.encode("utf-8") + ) + assert result is True + + # And correctly returns False for different passwords + result = secrets.compare_digest( + password.encode("utf-8"), "different£pass".encode("utf-8") + ) + assert result is False + + +@pytest.mark.asyncio +async def test_authenticate_user_database_login_with_non_ascii_password(): + """Test database user login with non-ASCII characters in password (issue #19559)""" + master_key = "sk-1234" + user_email = "test@example.com" + password_with_special_char = "correct£password" + hashed_password = hash_token(token=password_with_special_char) + + mock_user = MagicMock() + mock_user.user_id = "test-user-123" + mock_user.user_email = user_email + mock_user.password = hashed_password + mock_user.user_role = LitellmUserRoles.INTERNAL_USER + + def mock_find_first(**kwargs): + where = kwargs.get("where", {}) + user_email_filter = where.get("user_email", {}) + if str(user_email_filter.get("equals", "")).lower() == user_email.lower(): + return mock_user + return None + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + side_effect=mock_find_first + ) + + with patch.dict( + os.environ, + { + "DATABASE_URL": "postgresql://test:test@localhost/test", + "UI_USERNAME": "admin", + "UI_PASSWORD": "admin-password", + }, + ): + with patch( + "litellm.proxy.auth.login_utils.expire_previous_ui_session_tokens", + new_callable=AsyncMock, + return_value=None, + ): + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = {"token": "token-123"} + + result = await authenticate_user( + username=user_email, + password=password_with_special_char, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == "test-user-123" + assert result.user_email == user_email diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 61d44e46da5..5c039141928 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -549,6 +549,62 @@ class TestAdditionalParams: ) +class TestModelParameter: + """Test model parameter handling in guardrail requests""" + + @pytest.mark.asyncio + async def test_model_passed_from_inputs( + self, generic_guardrail, mock_request_data_input + ): + """Test that model is passed to the API when provided in inputs""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["test"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["test"], "model": "gpt-4"}, + request_data=mock_request_data_input, + input_type="request", + ) + + # Verify API was called with model + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert json_payload["model"] == "gpt-4" + + @pytest.mark.asyncio + async def test_model_none_when_not_provided( + self, generic_guardrail, mock_request_data_input + ): + """Test that model is None when not provided in inputs""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["test"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["test"]}, # No model in inputs + request_data=mock_request_data_input, + input_type="request", + ) + + # Verify API was called with model=None + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert json_payload["model"] is None + + class TestErrorHandling: """Test error handling scenarios""" diff --git a/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py b/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py new file mode 100644 index 00000000000..3bc111ef142 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py @@ -0,0 +1,273 @@ +""" +Integration tests for async_post_call_streaming_hook. + +Tests verify that the streaming hook can transform streaming responses sent to clients. +""" + +import os +import sys +import pytest +from typing import Any +from unittest.mock import patch, MagicMock + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + + +class StreamingResponseTransformerLogger(CustomLogger): + """Logger that transforms streaming responses""" + + def __init__(self, transform_content: str = None): + self.called = False + self.transform_content = transform_content + self.received_response = None + + async def async_post_call_streaming_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: str, + ) -> Any: + self.called = True + self.received_response = response + if self.transform_content is not None: + return self.transform_content + return None + + +@pytest.mark.asyncio +async def test_streaming_hook_transforms_response(): + """ + Test that async_post_call_streaming_hook can transform streaming responses. + """ + transformer = StreamingResponseTransformerLogger(transform_content="Modified streaming response") + + with patch("litellm.callbacks", [transformer]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + # Create a mock streaming response + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="Original content", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + # Call the hook + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify hook was called + assert transformer.called is True + + # Verify transformed response is returned + assert result == "Modified streaming response" + + +@pytest.mark.asyncio +async def test_streaming_hook_returns_none_keeps_original(): + """ + Test that hook returning None keeps the original response. + """ + + class NoOpLogger(CustomLogger): + def __init__(self): + self.called = False + + async def async_post_call_streaming_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: str, + ): + self.called = True + return None + + logger = NoOpLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="Original content", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Should return original response object + assert result.id == "original-stream" + assert logger.called is True + + +@pytest.mark.asyncio +async def test_streaming_hook_works_with_sse_format(): + """ + Test that hook works with SSE-formatted strings (data: prefix). + This was the only supported format before the fix. + """ + transformer = StreamingResponseTransformerLogger( + transform_content="data: {\"error\": \"custom error\"}\n\n" + ) + + with patch("litellm.callbacks", [transformer]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="Original content", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify SSE-formatted response is returned + assert result == "data: {\"error\": \"custom error\"}\n\n" + + +@pytest.mark.asyncio +async def test_streaming_hook_chains_multiple_callbacks(): + """ + Test that multiple callbacks can chain modifications. + """ + + class AppendLogger(CustomLogger): + def __init__(self, suffix: str): + self.suffix = suffix + self.called = False + + async def async_post_call_streaming_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: str, + ) -> str: + self.called = True + # Note: response here is the complete_response string, not the chunk + return f"[{self.suffix}]" + + callback1 = AppendLogger("CB1") + callback2 = AppendLogger("CB2") + + with patch("litellm.callbacks", [callback1, callback2]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="Hello", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Both callbacks should have been called + assert callback1.called is True + assert callback2.called is True + + # Last callback's result should be used + assert result == "[CB2]" + + +@pytest.mark.asyncio +async def test_streaming_hook_handles_exceptions(): + """ + Test that hook exceptions are propagated. + """ + + class FailingLogger(CustomLogger): + async def async_post_call_streaming_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: str, + ): + raise RuntimeError("Streaming hook crashed!") + + logger = FailingLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="Hello", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + # Exception should be propagated + with pytest.raises(RuntimeError, match="Streaming hook crashed!"): + await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) diff --git a/tests/test_litellm/proxy/hooks/test_post_call_success_hook_integration.py b/tests/test_litellm/proxy/hooks/test_post_call_success_hook_integration.py new file mode 100644 index 00000000000..870286f5382 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_post_call_success_hook_integration.py @@ -0,0 +1,260 @@ +""" +Integration tests for async_post_call_success_hook. + +Tests verify that the success hook can transform responses sent to clients. +This mirrors the behavior of CustomGuardrail hooks and streaming iterator hooks. +""" + +import os +import sys +import pytest +from typing import Any +from unittest.mock import patch, MagicMock + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import ModelResponse, Choices, Message, Usage + + +class ResponseTransformerLogger(CustomLogger): + """Logger that transforms successful responses""" + + def __init__(self, transform_content: str = None): + self.called = False + self.transform_content = transform_content + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + if self.transform_content is not None: + # Create a modified response with custom content + return { + "id": "transformed-response", + "choices": [ + { + "message": {"content": self.transform_content, "role": "assistant"}, + "index": 0, + } + ], + "model": "test-model", + "custom_field": "added_by_hook", + } + return response + + +@pytest.mark.asyncio +async def test_success_hook_transforms_response(): + """ + Test that async_post_call_success_hook can transform successful responses. + """ + transformer = ResponseTransformerLogger(transform_content="Modified response") + + with patch("litellm.callbacks", [transformer]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + # Create a mock response + original_response = ModelResponse( + id="original-response", + choices=[ + Choices( + message=Message(content="Original content", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + # Call the hook + result = await proxy_logging.post_call_success_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify hook was called + assert transformer.called is True + + # Verify transformed response is returned + assert result is not None + assert result["id"] == "transformed-response" + assert result["choices"][0]["message"]["content"] == "Modified response" + assert result["custom_field"] == "added_by_hook" + + +@pytest.mark.asyncio +async def test_success_hook_returns_none_keeps_original(): + """ + Test that hook returning None keeps the original response. + """ + + class NoOpLogger(CustomLogger): + def __init__(self): + self.called = False + + async def async_post_call_success_hook(self, *args, **kwargs): + self.called = True + return None + + logger = NoOpLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponse( + id="original-response", + choices=[ + Choices( + message=Message(content="Original content", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Should return original response + assert result.id == "original-response" + assert logger.called is True + + +@pytest.mark.asyncio +async def test_success_hook_chains_multiple_callbacks(): + """ + Test that multiple callbacks can chain modifications. + """ + + class AddFieldLogger(CustomLogger): + def __init__(self, field_name: str, field_value: Any): + self.field_name = field_name + self.field_value = field_value + self.called = False + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + ) -> Any: + self.called = True + # Convert response to dict if needed + if hasattr(response, "model_dump"): + resp_dict = response.model_dump() + elif hasattr(response, "dict"): + resp_dict = response.dict() + elif isinstance(response, dict): + resp_dict = response.copy() + else: + resp_dict = {} + + resp_dict[self.field_name] = self.field_value + return resp_dict + + callback1 = AddFieldLogger("field1", "value1") + callback2 = AddFieldLogger("field2", "value2") + + with patch("litellm.callbacks", [callback1, callback2]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponse( + id="original-response", + choices=[ + Choices( + message=Message(content="Original content", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.post_call_success_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Both callbacks should have been called + assert callback1.called is True + assert callback2.called is True + + # Both fields should be present (chained modifications) + assert result["field1"] == "value1" + assert result["field2"] == "value2" + + +@pytest.mark.asyncio +async def test_success_hook_handles_exceptions(): + """ + Test that hook exceptions are propagated (not silently swallowed). + """ + + class FailingLogger(CustomLogger): + async def async_post_call_success_hook(self, *args, **kwargs): + raise RuntimeError("Hook crashed!") + + logger = FailingLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponse( + id="original-response", + choices=[ + Choices( + message=Message(content="Original content", role="assistant"), + index=0, + finish_reason="stop", + ) + ], + model="test-model", + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + # Exception should be propagated + with pytest.raises(RuntimeError, match="Hook crashed!"): + await proxy_logging.post_call_success_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 397a6af556f..dc436bac087 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1096,3 +1096,57 @@ async def test_get_users_user_id_partial_match(mocker): assert "user_id" in captured_where_conditions assert "in" in captured_where_conditions["user_id"] assert captured_where_conditions["user_id"]["in"] == ["user1", "user2", "user3"] + + +def test_update_internal_user_params_reset_max_budget_with_none(): + """ + Test that _update_internal_user_params allows setting max_budget to None. + This verifies the fix for unsetting/resetting the budget to unlimited. + """ + + # Case 1: max_budget is explicitly None in the input dictionary + data_json = {"max_budget": None, "user_id": "test_user"} + data = UpdateUserRequest(max_budget=None, user_id="test_user") + + # Call the function + non_default_values = _update_internal_user_params(data_json=data_json, data=data) + + # Assertions + assert "max_budget" in non_default_values + assert non_default_values["max_budget"] is None + assert non_default_values["user_id"] == "test_user" + + +def test_update_internal_user_params_ignores_other_nones(): + """ + Test that other fields are still filtered out if None + """ + # Create test data with other None fields + data_json = {"user_alias": None, "user_id": "test_user", "max_budget": 100.0} + data = UpdateUserRequest(user_alias=None, user_id="test_user", max_budget=100.0) + + # Call the function + non_default_values = _update_internal_user_params(data_json=data_json, data=data) + + # Assertions + assert "user_alias" not in non_default_values + assert non_default_values["max_budget"] == 100.0 + + +def test_generate_request_base_validator(): + """ + Test that GenerateRequestBase validator converts empty string to None for max_budget + """ + from litellm.proxy._types import GenerateRequestBase + + # Test with empty string + req = GenerateRequestBase(max_budget="") + assert req.max_budget is None + + # Test with actual float + req = GenerateRequestBase(max_budget=100.0) + assert req.max_budget == 100.0 + + # Test with None + req = GenerateRequestBase(max_budget=None) + assert req.max_budget is None \ No newline at end of file diff --git a/tests/test_litellm/proxy/test_chat_completion_metadata.py b/tests/test_litellm/proxy/test_chat_completion_metadata.py new file mode 100644 index 00000000000..38dcdc13c50 --- /dev/null +++ b/tests/test_litellm/proxy/test_chat_completion_metadata.py @@ -0,0 +1,154 @@ +import pytest +from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy.proxy_server import chat_completion, completion, embeddings +from litellm.proxy._types import UserAPIKeyAuth +from fastapi import Request, Response + + +@pytest.mark.asyncio +async def test_chat_completion_metadata_population(): + # Setup + request = MagicMock(spec=Request) + # Mock _read_request_body to return a dict + with patch( + "litellm.proxy.proxy_server._read_request_body", new_callable=AsyncMock + ) as mock_read_body: + mock_read_body.return_value = {"model": "gpt-3.5-turbo", "messages": []} + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user_id", team_id="test_team_id", org_id="test_org_id" + ) + + fastapi_response = MagicMock(spec=Response) + + # Mock ProxyBaseLLMRequestProcessing + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + mock_instance = MockProcessor.return_value + mock_instance.base_process_llm_request = AsyncMock( + return_value={"choices": []} + ) + + # Execute + await chat_completion( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify + # Check if ProxyBaseLLMRequestProcessing was initialized with data containing metadata + call_args = MockProcessor.call_args + assert call_args is not None + data_arg = call_args.kwargs.get("data") + assert data_arg is not None + + assert "metadata" in data_arg + assert data_arg["metadata"]["user_api_key_user_id"] == "test_user_id" + assert data_arg["metadata"]["user_api_key_team_id"] == "test_team_id" + assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id" + + +@pytest.mark.asyncio +async def test_embedding_metadata_population(): + """ + Test that the embedding endpoint correctly populates metadata + from UserAPIKeyAuth. + """ + # Setup + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.base_process_llm_request" + ): + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.__init__", + return_value=None, + ) as mock_base_process_init: + # Create a mock UserAPIKeyAuth object + mock_user_auth = MagicMock(spec=UserAPIKeyAuth) + mock_user_auth.user_id = "test_user_id_emb" + mock_user_auth.team_id = "test_team_id_emb" + mock_user_auth.org_id = "test_org_id_emb" + + # Create a mock Request object + mock_request = MagicMock(spec=Request) + mock_request.json = AsyncMock( + return_value={"model": "gpt-3.5-turbo", "input": "hello"} + ) + # Mock _read_request_body to return our data + with patch( + "litellm.proxy.proxy_server._read_request_body", + new=AsyncMock( + return_value={"model": "gpt-3.5-turbo", "input": "hello"} + ), + ): + # Call the endpoint function directly + await embeddings( + request=mock_request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=mock_user_auth, + ) + + # Check if ProxyBaseLLMRequestProcessing was initialized with the correct metadata + mock_base_process_init.assert_called_once() + call_args = mock_base_process_init.call_args + # handle both positional and keyword args for data + if "data" in call_args.kwargs: + data_arg = call_args.kwargs["data"] + else: + data_arg = call_args.args[0] + + assert ( + data_arg["metadata"]["user_api_key_user_id"] == "test_user_id_emb" + ) + assert ( + data_arg["metadata"]["user_api_key_team_id"] == "test_team_id_emb" + ) + assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id_emb" + + +@pytest.mark.asyncio +async def test_completion_metadata_population(): + # Setup + request = MagicMock(spec=Request) + # Mock _read_request_body to return a dict + with patch( + "litellm.proxy.proxy_server._read_request_body", new_callable=AsyncMock + ) as mock_read_body: + mock_read_body.return_value = { + "model": "gpt-3.5-turbo-instruct", + "prompt": "test", + } + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user_id_2", team_id="test_team_id_2", org_id="test_org_id_2" + ) + + fastapi_response = MagicMock(spec=Response) + + # Mock ProxyBaseLLMRequestProcessing + with patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + mock_instance = MockProcessor.return_value + mock_instance.base_process_llm_request = AsyncMock( + return_value={"choices": []} + ) + + # Execute + await completion( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + ) + + # Verify + call_args = MockProcessor.call_args + assert call_args is not None + data_arg = call_args.kwargs.get("data") + assert data_arg is not None + + assert "metadata" in data_arg + assert data_arg["metadata"]["user_api_key_user_id"] == "test_user_id_2" + assert data_arg["metadata"]["user_api_key_team_id"] == "test_team_id_2" + assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id_2" diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index ad4f53dac4b..60bdb7d12cb 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1227,3 +1227,113 @@ class TestProxySettingEndpoints: assert retrieved_role_mappings["provider"] == "google" assert retrieved_role_mappings["group_claim"] == "groups" assert retrieved_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER + + def test_setup_role_mappings_custom_logic_with_env_vars(self, monkeypatch): + """Test the _setup_role_mappings function directly with custom role mapping logic from environment variables""" + import asyncio + import os + from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings + from litellm.proxy._types import LitellmUserRoles + + # Set up environment variables for custom role mappings using valid Python dict format + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", "{'proxy_admin': ['custom-admin-group'], 'internal_user': ['custom-user-group'], 'proxy_admin_viewer': ['custom-viewer-group']}") + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", "custom-groups") + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", "internal_user_viewer") + + # Debug: Print environment variables + print("GENERIC_ROLE_MAPPINGS_ROLES:", os.getenv("GENERIC_ROLE_MAPPINGS_ROLES")) + print("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM:", os.getenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM")) + print("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE:", os.getenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE")) + + # Run the async function + role_mappings = asyncio.run(_setup_role_mappings()) + + # Debug: Print result + print("role_mappings result:", role_mappings) + + # Verify role_mappings is returned correctly from environment variables + assert role_mappings is not None + assert role_mappings.provider == "generic" + assert role_mappings.group_claim == "custom-groups" + assert role_mappings.default_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + assert role_mappings.roles[LitellmUserRoles.PROXY_ADMIN] == ["custom-admin-group"] + assert role_mappings.roles[LitellmUserRoles.INTERNAL_USER] == ["custom-user-group"] + assert role_mappings.roles[LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] == ["custom-viewer-group"] + + def test_setup_role_mappings_custom_logic_with_no_config(self, monkeypatch): + """Test the _setup_role_mappings function returns None when no configuration is available""" + import asyncio + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings + + # Ensure environment variables are not set + monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_ROLES", raising=False) + monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", raising=False) + monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", raising=False) + + # Mock the prisma client to return None (no database record) + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + # Run the async function + role_mappings = asyncio.run(_setup_role_mappings()) + + # Should return None when no configuration is available + assert role_mappings is None + + def test_get_sso_settings_with_env_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", '{"proxy_admin": ["custom-admin-group"], "internal_user": ["custom-user-group"], "proxy_admin_viewer": ["custom-viewer-group"]}') + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", "custom-groups") + monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", "internal_user_viewer") + + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.sso_settings = { + "google_client_id": "test_google_client_id", + "role_mappings": { + "provider": "google", + "group_claim": "db-groups", + "default_role": "proxy_admin", + "roles": { + "proxy_admin": ["db-admin-group"], + }, + }, + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables + ) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + data = response.json() + + values = data["values"] + assert "role_mappings" in values + assert values["role_mappings"] is not None + + # The database values shoeld override the environment variables + assert values["role_mappings"]["provider"] == "google" + assert values["role_mappings"]["group_claim"] == "db-groups" + assert values["role_mappings"]["default_role"] == LitellmUserRoles.PROXY_ADMIN + assert values["role_mappings"]["roles"][LitellmUserRoles.PROXY_ADMIN] == ["db-admin-group"] + + # Verify that the database was checked but environment variables took priority + mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with( + where={"id": "sso_config"} + ) + + # Verify other SSO settings are still correctly returned + assert values["google_client_id"] == "test_google_client_id" + + # Verify field_schema is still present + assert "field_schema" in data + assert "properties" in data["field_schema"] + assert "role_mappings" in data["field_schema"]["properties"] diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/test_litellm/secret_managers/test_secret_managers_main.py index 2e5270b5d70..eaef6956cd5 100644 --- a/tests/test_litellm/secret_managers/test_secret_managers_main.py +++ b/tests/test_litellm/secret_managers/test_secret_managers_main.py @@ -47,11 +47,12 @@ def mock_env(): @patch("litellm.secret_managers.main.oidc_cache") -@patch("litellm.secret_managers.main.HTTPHandler") -def test_oidc_google_success(mock_http_handler, mock_oidc_cache): +@patch("litellm.secret_managers.main._get_oidc_http_handler") +@patch("httpx.Client") # Prevent any real HTTP connections +def test_oidc_google_success(mock_httpx_client, mock_get_http_handler, mock_oidc_cache): mock_oidc_cache.get_cache.return_value = None mock_handler = MockHTTPHandler(timeout=600.0) - mock_http_handler.return_value = mock_handler + mock_get_http_handler.return_value = mock_handler secret_name = "oidc/google/[invalid url, do not cite]" result = get_secret(secret_name) @@ -63,29 +64,31 @@ def test_oidc_google_success(mock_http_handler, mock_oidc_cache): @patch("litellm.secret_managers.main.oidc_cache") -def test_oidc_google_cached(mock_oidc_cache): +@patch("litellm.secret_managers.main._get_oidc_http_handler") +def test_oidc_google_cached(mock_get_http_handler, mock_oidc_cache): mock_oidc_cache.get_cache.return_value = "cached_token" secret_name = "oidc/google/[invalid url, do not cite]" - with patch("litellm.secret_managers.main.HTTPHandler") as mock_http: - result = get_secret(secret_name) + result = get_secret(secret_name) - assert result == "cached_token", f"Expected cached token, got {result}" - mock_oidc_cache.get_cache.assert_called_with(key=secret_name) - mock_http.assert_not_called() + assert result == "cached_token", f"Expected cached token, got {result}" + mock_oidc_cache.get_cache.assert_called_with(key=secret_name) + # Verify HTTP handler was never called since we had a cached token + mock_get_http_handler.assert_not_called() @patch("litellm.secret_managers.main.oidc_cache") -def test_oidc_google_failure(mock_oidc_cache): +@patch("litellm.secret_managers.main._get_oidc_http_handler") +def test_oidc_google_failure(mock_get_http_handler, mock_oidc_cache): mock_handler = MockHTTPHandler(timeout=600.0) mock_handler.status_code = 400 + mock_get_http_handler.return_value = mock_handler + mock_oidc_cache.get_cache.return_value = None + + secret_name = "oidc/google/https://example.com/api" - with patch("litellm.secret_managers.main.HTTPHandler", return_value=mock_handler): - mock_oidc_cache.get_cache.return_value = None - secret_name = "oidc/google/https://example.com/api" - - with pytest.raises(ValueError, match="Google OIDC provider failed"): - get_secret(secret_name) + with pytest.raises(ValueError, match="Google OIDC provider failed"): + get_secret(secret_name) def test_oidc_circleci_success(monkeypatch): @@ -106,13 +109,13 @@ def test_oidc_circleci_failure(monkeypatch): @patch("litellm.secret_managers.main.oidc_cache") -@patch("litellm.secret_managers.main.HTTPHandler") -def test_oidc_github_success(mock_http_handler, mock_oidc_cache, mock_env): +@patch("litellm.secret_managers.main._get_oidc_http_handler") +def test_oidc_github_success(mock_get_http_handler, mock_oidc_cache, mock_env): mock_env["ACTIONS_ID_TOKEN_REQUEST_URL"] = "https://github.com/token" mock_env["ACTIONS_ID_TOKEN_REQUEST_TOKEN"] = "github_token" mock_oidc_cache.get_cache.return_value = None mock_handler = MockHTTPHandler(timeout=600.0) - mock_http_handler.return_value = mock_handler + mock_get_http_handler.return_value = mock_handler secret_name = "oidc/github/github-audience" result = get_secret(secret_name) @@ -142,7 +145,7 @@ def test_oidc_azure_file_success(mock_env, tmp_path): mock_env["AZURE_FEDERATED_TOKEN_FILE"] = str(token_file) secret_name = "oidc/azure/azure-audience" - result = get_secret(secret_name) + result = get_secret(secret_name) assert result == "azure_token" @@ -154,16 +157,22 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider): if "AZURE_FEDERATED_TOKEN_FILE" in os.environ: del os.environ["AZURE_FEDERATED_TOKEN_FILE"] + # Mock the token provider function that gets returned and called mock_token_provider = Mock(return_value="azure_ad_token") mock_get_azure_ad_token_provider.return_value = mock_token_provider - secret_name = "oidc/azure/api://azure-audience" - result = get_secret(secret_name) + + # Also mock the Azure Identity SDK to prevent any real Azure calls + with patch("azure.identity.get_bearer_token_provider") as mock_bearer: + mock_bearer.return_value = mock_token_provider + + secret_name = "oidc/azure/api://azure-audience" + result = get_secret(secret_name) - assert result == "azure_ad_token" - mock_get_azure_ad_token_provider.assert_called_once_with( - azure_scope="api://azure-audience" - ) - mock_token_provider.assert_called_once_with() + assert result == "azure_ad_token" + mock_get_azure_ad_token_provider.assert_called_once_with( + azure_scope="api://azure-audience" + ) + mock_token_provider.assert_called_once_with() def test_oidc_file_success(tmp_path): diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py index 0a2a62b6c97..620c0734980 100644 --- a/tests/test_litellm/test_gpt_image_cost_calculator.py +++ b/tests/test_litellm/test_gpt_image_cost_calculator.py @@ -19,10 +19,13 @@ import pytest import litellm from litellm.types.utils import ( + CompletionTokensDetailsWrapper, ImageResponse, ImageObject, ImageUsage, ImageUsageInputTokensDetails, + PromptTokensDetailsWrapper, + Usage, ) @@ -202,6 +205,71 @@ class TestGPTImageCostRouting: assert cost >= 0 +class TestGPTImage15OutputImageTokens: + """ + Test for GitHub issue #19508: + Image usage calculation does not include image tokens in gpt-image-1.5 + + gpt-image-1.5 returns output_tokens_details with separate image_tokens and text_tokens, + and these must be correctly included in cost calculation. + """ + + def test_gpt_image_15_output_image_tokens_cost(self): + """ + Test that output image tokens are correctly included in cost calculation. + + This tests the fix for issue #19508 where output_tokens_details.image_tokens + were not being included in the cost calculation, causing costs to be + underreported (e.g., $0.046 instead of $0.14). + """ + # Simulate gpt-image-1.5 response with output_tokens_details + # This is what the API returns and what convert_to_image_response transforms + usage = Usage( + prompt_tokens=169, + completion_tokens=4599, + total_tokens=4768, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=169, + image_tokens=0, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=439, + image_tokens=4160, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(b64_json="test")], + ) + image_response.usage = usage + image_response._hidden_params = {"custom_llm_provider": "openai"} + + cost = litellm.completion_cost( + completion_response=image_response, + model="gpt-image-1.5", + call_type="image_generation", + custom_llm_provider="openai", + ) + + # gpt-image-1.5 pricing: + # - input_cost_per_token: 5e-06 ($5/1M for text input) + # - output_cost_per_token: 1e-05 ($10/1M for text output) + # - output_cost_per_image_token: 3.2e-05 ($32/1M for image output) + # + # Expected cost: + # Input text: 169 * $5/1M = $0.000845 + # Output text: 439 * $10/1M = $0.00439 + # Output image: 4160 * $32/1M = $0.13312 + # Total: $0.138355 + expected_cost = 169 * 5e-06 + 439 * 1e-05 + 4160 * 3.2e-05 + + assert abs(cost - expected_cost) < 1e-6, ( + f"Expected {expected_cost}, got {cost}. " + f"Image tokens may not be included in cost calculation." + ) + + class TestCompletionCostIntegration: """Test the full completion_cost integration for gpt-image-1""" diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index d1185c72a29..80fd9f61298 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -477,6 +477,12 @@ async def test_openai_env_base( respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch ): "This tests OpenAI env variables are honored, including legacy OPENAI_API_BASE" + # Clear cache to ensure no cached clients from previous tests interfere + # This prevents cache pollution where a previous test cached a client with + # aiohttp transport, which would bypass respx mocks + if hasattr(litellm, "in_memory_llm_clients_cache"): + litellm.in_memory_llm_clients_cache.flush_cache() + # Ensure aiohttp transport is disabled to use httpx which respx can mock litellm.disable_aiohttp_transport = True diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/test_litellm/test_router_per_deployment_num_retries.py index 4021ca28073..154ba579e4e 100644 --- a/tests/test_litellm/test_router_per_deployment_num_retries.py +++ b/tests/test_litellm/test_router_per_deployment_num_retries.py @@ -32,17 +32,17 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + # Create a mock exception without num_retries class MockException(Exception): pass - + exc = MockException("test error") assert not hasattr(exc, "num_retries") or exc.num_retries is None - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was set from deployment assert exc.num_retries == 5 @@ -66,16 +66,16 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + # Create an exception that already has num_retries class MockException(Exception): num_retries = 10 # Already set - + exc = MockException("test error") - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was NOT overridden assert exc.num_retries == 10 @@ -99,15 +99,15 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + class MockException(Exception): pass - + exc = MockException("test error") - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was not set (deployment has no num_retries) assert not hasattr(exc, "num_retries") or exc.num_retries is None @@ -155,3 +155,36 @@ class TestPerDeploymentNumRetries: kwargs = {} router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) assert kwargs["num_retries"] == 7 # Uses global + + def test_set_deployment_num_retries_with_string_value(self): + """ + Test that _set_deployment_num_retries_on_exception handles string values + from environment variables correctly. + GitHub Issue: #19481 + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + "num_retries": "6", # String value (as from env var) + }, + }, + ], + num_retries=0, # Global setting + ) + + deployment = router.model_list[0] + + class MockException(Exception): + pass + + exc = MockException("test error") + + # Call the helper + router._set_deployment_num_retries_on_exception(exc, deployment) + + # Verify num_retries was converted from string to int + assert exc.num_retries == 6 diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py index 3c4308dd3aa..a3af476bc71 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -119,8 +119,12 @@ def test_add_vector_store_to_registry(): +@respx.mock def test_search_uses_registry_credentials(): """search() should pull credentials from vector_store_registry when available""" + # Block all HTTP requests at the network level to prevent real API calls + respx.route().mock(return_value=httpx.Response(200, json={"object": "list", "data": []})) + vector_store = LiteLLM_ManagedVectorStore( vector_store_id="vs1", custom_llm_provider="bedrock", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowPrompts.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowPrompts.ts new file mode 100644 index 00000000000..801fbdbb99d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowPrompts.ts @@ -0,0 +1,35 @@ +// hooks/useDisableShowPrompts.ts +import { useSyncExternalStore } from "react"; +import { getLocalStorageItem } from "@/utils/localStorageUtils"; +import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; + +function subscribe(callback: () => void) { + const onStorage = (e: StorageEvent) => { + if (e.key === "disableShowPrompts") { + callback(); + } + }; + + const onCustom = (e: Event) => { + const { key } = (e as CustomEvent).detail; + if (key === "disableShowPrompts") { + callback(); + } + }; + + window.addEventListener("storage", onStorage); + window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); + + return () => { + window.removeEventListener("storage", onStorage); + window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); + }; +} + +function getSnapshot() { + return getLocalStorageItem("disableShowPrompts") === "true"; +} + +export function useDisableShowPrompts() { + return useSyncExternalStore(subscribe, getSnapshot); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 97837ff8e0a..97e4c799e72 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -56,8 +56,10 @@ export default function Layout({ children }: { children: React.ReactNode }) { userRole={userRole} premiumUser={premiumUser} proxySettings={undefined} - setProxySettings={() => {}} + setProxySettings={() => { }} accessToken={accessToken} + isDarkMode={false} + toggleDarkMode={() => { }} />+ Got it, we will not ask again. Reactivate this at any time in the User Menu. +
+{description}
- +
+
Click or drag files to this area to upload
++ Support for single or bulk upload. Supported formats: PDF, TXT, DOCX, MD +
++ Vector Store ID: {ingestResults[0]?.vector_store_id} +
++ Documents Ingested: {ingestResults.length} +
+You can use vector stores to store and retrieve LLM embeddings..
+You can use vector stores to store and retrieve LLM embeddings.