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/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py new file mode 100644 index 00000000000..4a9fa3de5fd --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -0,0 +1,203 @@ +import pytest +from unittest.mock import MagicMock, patch +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + UserAPIKeyLabelValues, +) +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_includes_client_ip_user_agent(): + """ + Test that async_post_call_failure_hook includes client_ip and user_agent in UserAPIKeyLabelValues + """ + # Mocking + # Mocking + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + # Initialize attributes manually as __init__ is mocked + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=["client_ip", "user_agent"] + ) + + request_data = { + "model": "gpt-4", + "metadata": { + "requester_ip_address": "127.0.0.1", + "user_agent": "test-agent", + }, + } + user_api_key_dict = UserAPIKeyAuth(token="test_token") + original_exception = Exception("Test exception") + + # Mock prometheus_label_factory to inspect arguments + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) + + # Verification + assert mock_label_factory.call_count >= 1 + + # Check calls + calls = mock_label_factory.call_args_list + found = False + for call in calls: + kwargs = call.kwargs + enum_values = kwargs.get("enum_values") + if isinstance(enum_values, UserAPIKeyLabelValues): + if ( + enum_values.client_ip == "127.0.0.1" + and enum_values.user_agent == "test-agent" + ): + found = True + break + + assert ( + found + ), "UserAPIKeyLabelValues should contain client_ip='127.0.0.1' and user_agent='test-agent'" + + +@pytest.mark.asyncio +async def test_async_post_call_success_hook_includes_client_ip_user_agent(): + """ + Test that async_post_call_success_hook includes client_ip and user_agent in UserAPIKeyLabelValues + """ + # Mocking + # Mocking + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=["client_ip", "user_agent"] + ) + + data = { + "model": "gpt-4", + "metadata": { + "requester_ip_address": "192.168.1.1", + "user_agent": "success-agent", + }, + } + user_api_key_dict = UserAPIKeyAuth(token="test_token") + response = MagicMock() + + # Mock prometheus_label_factory to inspect arguments + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + + await logger.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + + # Verification + assert mock_label_factory.call_count >= 1 + + # Check calls + calls = mock_label_factory.call_args_list + found = False + for call in calls: + kwargs = call.kwargs + enum_values = kwargs.get("enum_values") + if isinstance(enum_values, UserAPIKeyLabelValues): + if ( + enum_values.client_ip == "192.168.1.1" + and enum_values.user_agent == "success-agent" + ): + found = True + break + + assert ( + found + ), "UserAPIKeyLabelValues should contain client_ip='192.168.1.1' and user_agent='success-agent'" + + +def test_set_llm_deployment_failure_metrics_includes_client_ip_user_agent(): + """ + Test that set_llm_deployment_failure_metrics includes client_ip and user_agent in UserAPIKeyLabelValues + """ + # Mocking + # Mocking + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_deployment_failure_responses = MagicMock() + logger.litellm_deployment_total_requests = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=["client_ip", "user_agent"] + ) + logger.set_deployment_partial_outage = MagicMock() + + request_kwargs = { + "model": "gpt-4", + "standard_logging_object": { + "metadata": { + "requester_ip_address": "10.0.0.1", + "user_agent": "failure-deployment", + "user_api_key_team_id": "team_1", + "user_api_key_team_alias": "team_alias_1", + "user_api_key_alias": "key_alias_1", + }, + "model_group": "group_1", + "api_base": "http://api.base", + "model_id": "model_1", + }, + "litellm_params": {}, + "exception": Exception("Deployment failure"), + } + + # Mock prometheus_label_factory to inspect arguments + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + + logger.set_llm_deployment_failure_metrics(request_kwargs=request_kwargs) + + # Verification + assert mock_label_factory.call_count >= 1 + + # Check calls + calls = mock_label_factory.call_args_list + found = False + for call in calls: + kwargs = call.kwargs + enum_values = kwargs.get("enum_values") + if isinstance(enum_values, UserAPIKeyLabelValues): + if ( + enum_values.client_ip == "10.0.0.1" + and enum_values.user_agent == "failure-deployment" + ): + found = True + break + + assert ( + found + ), "UserAPIKeyLabelValues should contain client_ip='10.0.0.1' and user_agent='failure-deployment'" + + +if __name__ == "__main__": + import asyncio + + asyncio.run(test_async_post_call_failure_hook_includes_client_ip_user_agent()) + asyncio.run(test_async_post_call_success_hook_includes_client_ip_user_agent()) + test_set_llm_deployment_failure_metrics_includes_client_ip_user_agent() + print("✅ All client_ip and user_agent tests passed!") diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index c0b863ef6ee..a83bc1df1e1 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -26,15 +26,49 @@ def test_user_email_in_required_metrics(): "litellm_input_tokens_metric", "litellm_output_tokens_metric", "litellm_requests_metric", - "litellm_spend_metric" + "litellm_spend_metric", ] for metric_name in metrics_with_user_email: labels = PrometheusMetricLabels.get_labels(metric_name) - assert user_email_label in labels, f"Metric {metric_name} should contain user_email label" + assert ( + user_email_label in labels + ), f"Metric {metric_name} should contain user_email label" print(f"✅ {metric_name} contains user_email label") +def test_model_id_in_required_metrics(): + """ + Test that model_id label is present in all the metrics that should have it + """ + model_id_label = UserAPIKeyLabelNames.MODEL_ID.value + + # Metrics that should have model_id + metrics_with_model_id = [ + "litellm_proxy_total_requests_metric", + "litellm_proxy_failed_requests_metric", + "litellm_input_tokens_metric", + "litellm_output_tokens_metric", + "litellm_requests_metric", + "litellm_spend_metric", + "litellm_llm_api_latency_metric", + "litellm_remaining_requests_metric", + "litellm_deployment_successful_fallbacks", + "litellm_cache_hits_metric", + "litellm_cache_misses_metric", + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + "litellm_llm_api_failed_requests_metric", + ] + + for metric_name in metrics_with_model_id: + labels = PrometheusMetricLabels.get_labels(metric_name) + assert ( + model_id_label in labels + ), f"Metric {metric_name} should contain model_id label" + print(f"✅ {metric_name} contains model_id label") + + def test_user_email_label_exists(): """Test that the USER_EMAIL label is properly defined""" assert UserAPIKeyLabelNames.USER_EMAIL.value == "user_email" @@ -52,12 +86,14 @@ def test_prometheus_metric_labels_structure(): "litellm_proxy_failed_requests_metric", "litellm_input_tokens_metric", "litellm_output_tokens_metric", - "litellm_spend_metric" + "litellm_spend_metric", ] for metric_name in test_metrics: # Check metric is in DEFINED_PROMETHEUS_METRICS - assert metric_name in get_args(DEFINED_PROMETHEUS_METRICS), f"{metric_name} should be in DEFINED_PROMETHEUS_METRICS" + assert metric_name in get_args( + DEFINED_PROMETHEUS_METRICS + ), f"{metric_name} should be in DEFINED_PROMETHEUS_METRICS" # Check labels can be retrieved labels = PrometheusMetricLabels.get_labels(metric_name) @@ -74,11 +110,11 @@ def test_route_normalization_for_responses_api(): """ Test that route normalization prevents high cardinality in Prometheus metrics for the /v1/responses/{response_id} endpoint. - + Issue: https://github.com/BerriAI/litellm/issues/XXXX Each unique response ID was creating a separate metric line, causing the /metrics endpoint to grow to ~30MB and take ~40 seconds to respond. - + Fix: Routes are normalized to collapse dynamic IDs into placeholders. """ from litellm.proxy.auth.auth_utils import normalize_request_route @@ -91,43 +127,53 @@ def test_route_normalization_for_responses_api(): ("/v1/responses/resp_abc123", "/v1/responses/{response_id}"), ("/v1/responses/litellm_poll_xyz", "/v1/responses/{response_id}"), ] - + for original, expected in responses_routes: normalized = normalize_request_route(original) - assert normalized == expected, \ - f"Failed: {original} -> {normalized} (expected {expected})" - + assert ( + normalized == expected + ), f"Failed: {original} -> {normalized} (expected {expected})" + # Verify cardinality reduction - unique_normalized = set(normalize_request_route(route) for route, _ in responses_routes) - assert len(unique_normalized) == 1, \ - f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}" - - print(f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label") - + unique_normalized = set( + normalize_request_route(route) for route, _ in responses_routes + ) + assert ( + len(unique_normalized) == 1 + ), f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}" + + print( + f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label" + ) + def test_route_normalization_for_sub_routes(): """Test that sub-routes like /cancel and /input_items are normalized correctly""" from litellm.proxy.auth.auth_utils import normalize_request_route - + sub_routes = [ ("/v1/responses/id1/cancel", "/v1/responses/{response_id}/cancel"), ("/v1/responses/id2/cancel", "/v1/responses/{response_id}/cancel"), ("/v1/responses/id3/input_items", "/v1/responses/{response_id}/input_items"), - ("/openai/v1/responses/id4/input_items", "/openai/v1/responses/{response_id}/input_items"), + ( + "/openai/v1/responses/id4/input_items", + "/openai/v1/responses/{response_id}/input_items", + ), ] - + for original, expected in sub_routes: normalized = normalize_request_route(original) - assert normalized == expected, \ - f"Failed: {original} -> {normalized} (expected {expected})" - + assert ( + normalized == expected + ), f"Failed: {original} -> {normalized} (expected {expected})" + print("✅ Sub-routes normalized correctly") def test_route_normalization_preserves_static_routes(): """Test that static routes are not affected by normalization""" from litellm.proxy.auth.auth_utils import normalize_request_route - + static_routes = [ "/chat/completions", "/v1/chat/completions", @@ -137,46 +183,47 @@ def test_route_normalization_preserves_static_routes(): "/v1/models", "/v1/responses", # List endpoint without ID ] - + for route in static_routes: normalized = normalize_request_route(route) - assert normalized == route, \ - f"Static route should not be modified: {route} -> {normalized}" - + assert ( + normalized == route + ), f"Static route should not be modified: {route} -> {normalized}" + print(f"✅ {len(static_routes)} static routes preserved") def test_route_normalization_other_dynamic_apis(): """Test normalization for other OpenAI-compatible APIs with dynamic IDs""" from litellm.proxy.auth.auth_utils import normalize_request_route - + test_cases = [ # Threads API ("/v1/threads/thread_123", "/v1/threads/{thread_id}"), ("/v1/threads/thread_abc/messages", "/v1/threads/{thread_id}/messages"), - ("/v1/threads/thread_abc/runs/run_123", "/v1/threads/{thread_id}/runs/{run_id}"), - + ( + "/v1/threads/thread_abc/runs/run_123", + "/v1/threads/{thread_id}/runs/{run_id}", + ), # Vector Stores API ("/v1/vector_stores/vs_123", "/v1/vector_stores/{vector_store_id}"), ("/v1/vector_stores/vs_123/files", "/v1/vector_stores/{vector_store_id}/files"), - # Assistants API ("/v1/assistants/asst_123", "/v1/assistants/{assistant_id}"), - # Files API ("/v1/files/file_123", "/v1/files/{file_id}"), ("/v1/files/file_123/content", "/v1/files/{file_id}/content"), - # Batches API ("/v1/batches/batch_123", "/v1/batches/{batch_id}"), ("/v1/batches/batch_123/cancel", "/v1/batches/{batch_id}/cancel"), ] - + for original, expected in test_cases: normalized = normalize_request_route(original) - assert normalized == expected, \ - f"Failed: {original} -> {normalized} (expected {expected})" - + assert ( + normalized == expected + ), f"Failed: {original} -> {normalized} (expected {expected})" + print(f"✅ {len(test_cases)} other API routes normalized correctly") @@ -195,26 +242,29 @@ def test_prometheus_metrics_use_normalized_routes(): # Create a mock PrometheusLogger prometheus_logger = MagicMock() - prometheus_logger.get_labels_for_metric = PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) - + prometheus_logger.get_labels_for_metric = ( + PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) + ) + # Test with a normalized route enum_values = UserAPIKeyLabelValues( route="/v1/responses/{response_id}", # Normalized route status_code="200", requested_model="gpt-4", ) - + labels = prometheus_label_factory( supported_enum_labels=prometheus_logger.get_labels_for_metric( metric_name="litellm_proxy_total_requests_metric" ), enum_values=enum_values, ) - + # Verify the route is normalized in labels - assert labels["route"] == "/v1/responses/{response_id}", \ - f"Expected normalized route in labels, got: {labels.get('route')}" - + assert ( + labels["route"] == "/v1/responses/{response_id}" + ), f"Expected normalized route in labels, got: {labels.get('route')}" + print("✅ Prometheus metrics use normalized routes in labels") @@ -227,4 +277,4 @@ if __name__ == "__main__": test_route_normalization_preserves_static_routes() test_route_normalization_other_dynamic_apis() test_prometheus_metrics_use_normalized_routes() - print("\n✅ All prometheus label tests passed!") \ No newline at end of file + print("\n✅ All prometheus label tests passed!") diff --git a/tests/test_litellm/integrations/test_prometheus_missing_metrics.py b/tests/test_litellm/integrations/test_prometheus_missing_metrics.py new file mode 100644 index 00000000000..7fcfb21ed4c --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_missing_metrics.py @@ -0,0 +1,77 @@ +""" +Unit tests for the new Prometheus metrics that were previously missing from validation. + +Tests for: +- litellm_remaining_api_key_requests_for_model +- litellm_remaining_api_key_tokens_for_model +- litellm_callback_logging_failures_metric +""" +from typing import get_args +from litellm.types.integrations.prometheus import ( + DEFINED_PROMETHEUS_METRICS, + PrometheusMetricLabels, + UserAPIKeyLabelNames, +) + + +def test_new_metrics_in_defined_metrics(): + """ + Test that the new metrics are present in DEFINED_PROMETHEUS_METRICS. + """ + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + new_metrics = [ + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + "litellm_callback_logging_failures_metric", + ] + + for metric in new_metrics: + assert ( + metric in defined_metrics + ), f"{metric} should be in DEFINED_PROMETHEUS_METRICS" + + +def test_new_metrics_have_correct_labels(): + """ + Test that the new metrics have the correct labels defined. + """ + # Test API Key limits metrics labels + api_key_metrics = [ + "litellm_remaining_api_key_requests_for_model", + "litellm_remaining_api_key_tokens_for_model", + ] + + expected_api_key_labels = [ + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + for metric in api_key_metrics: + labels = PrometheusMetricLabels.get_labels(metric) + for expected_label in expected_api_key_labels: + assert ( + expected_label in labels + ), f"{metric} should have label {expected_label}" + + # Test Callback failure metric labels + callback_metric = "litellm_callback_logging_failures_metric" + callback_labels = PrometheusMetricLabels.get_labels(callback_metric) + + assert ( + UserAPIKeyLabelNames.CALLBACK_NAME.value in callback_labels + ), f"{callback_metric} should have label {UserAPIKeyLabelNames.CALLBACK_NAME.value}" + + +def test_callback_name_label_definition(): + """ + Test that CALLBACK_NAME is defined correctly in UserAPIKeyLabelNames. + """ + assert UserAPIKeyLabelNames.CALLBACK_NAME.value == "callback_name" + + +if __name__ == "__main__": + test_new_metrics_in_defined_metrics() + test_new_metrics_have_correct_labels() + test_callback_name_label_definition() 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_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index b15d75a4145..9c2939b2da5 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -66,6 +66,41 @@ class LargeImageClient: ) +class StreamingLargeImageClient: + """ + Client that streams a large image to test streaming download protection. + This simulates a huge file without actually creating it all in memory. + """ + + def __init__(self, size_mb=100, include_content_length=False): + self.size_mb = size_mb + self.include_content_length = include_content_length + + def get(self, url, follow_redirects=True): + size_bytes = int(self.size_mb * 1024 * 1024) + headers = {"Content-Type": "image/jpeg"} + if self.include_content_length: + headers["Content-Length"] = str(size_bytes) + + # Create a generator that yields chunks without creating the whole file in memory + def generate_chunks(total_size, chunk_size=8192): + bytes_sent = 0 + while bytes_sent < total_size: + chunk = b"x" * min(chunk_size, total_size - bytes_sent) + bytes_sent += len(chunk) + yield chunk + + # Create response with streaming content + response = Response( + status_code=200, + headers=headers, + request=Request("GET", url), + ) + # Mock the iter_bytes method to return our generator + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) + return response + + def test_image_exceeds_size_limit_with_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected when Content-Length header is present. @@ -83,6 +118,7 @@ def test_image_exceeds_size_limit_with_content_length(monkeypatch): def test_image_exceeds_size_limit_without_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected even without Content-Length header. + This uses the old non-streaming mock for backward compatibility. """ monkeypatch.setattr( litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False) @@ -94,6 +130,29 @@ def test_image_exceeds_size_limit_without_content_length(monkeypatch): assert "exceeds maximum allowed size" in str(excinfo.value) +def test_streaming_download_protects_against_huge_files(monkeypatch): + """ + Test that streaming download aborts early when file exceeds size limit, + preventing memory exhaustion from huge files (e.g., petabyte-sized files). + + This test verifies that the streaming implementation doesn't download the entire + file into memory before checking size. Instead, it should abort as soon as the + limit is exceeded during streaming. + """ + # Simulate a 1GB file - far larger than the 50MB default limit + client = StreamingLargeImageClient(size_mb=1024, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/huge-image.jpg") + + # Verify the error message shows it was caught during streaming + assert "exceeds maximum allowed size" in str(excinfo.value) + + # The error should be raised after downloading just slightly more than the limit + # not after downloading the full 1GB + + class SmallImageClient: """ Client that returns a small valid image. @@ -124,6 +183,26 @@ def test_image_within_size_limit(monkeypatch): assert result.startswith("data:image/jpeg;base64,") +def test_streaming_download_handles_petabyte_file(monkeypatch): + """ + Test that streaming download can handle extremely large file URLs (e.g., petabyte-sized) + without attempting to download the entire file or causing memory exhaustion. + + This simulates what happens if a malicious actor or misconfiguration provides + a URL to an extremely large file. + """ + # Simulate a 1 petabyte file (1,000,000 GB) + # Without streaming protection, this would cause OOM or hang indefinitely + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/petabyte-file.jpg") + + # Should fail fast without downloading anywhere near 1 petabyte + assert "exceeds maximum allowed size" in str(excinfo.value) + + def test_image_size_limit_disabled(monkeypatch): """ Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables all image URL downloads. 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/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 1cb84d32c1d..98b392a353d 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -279,3 +279,188 @@ def test_output_format_with_no_schema(): # Content should remain as string (not converted to list) assert isinstance(last_user_message["content"], str) assert last_user_message["content"] == "Hello" + + +def test_advanced_tool_use_header_translation_for_opus_4_5(): + """ + Test that advanced-tool-use-2025-11-20 header is translated to Bedrock-specific headers + for Claude Opus 4.5. + + Regression test for: Claude Code sends advanced-tool-use header which needs to be + translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for Bedrock + Invoke API on Claude Opus 4.5. + + Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Simulate advanced-tool-use header from Claude Code + headers = { + "anthropic-beta": "advanced-tool-use-2025-11-20" + } + + # Test with Claude Opus 4.5 + result = config.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20250514-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + # Verify advanced-tool-use header was removed + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + assert "advanced-tool-use-2025-11-20" not in beta_headers, \ + "advanced-tool-use header should be removed for Bedrock" + + # Verify Bedrock-specific headers were added + assert "tool-search-tool-2025-10-19" in beta_headers, \ + "tool-search-tool-2025-10-19 should be added for Opus 4.5" + assert "tool-examples-2025-10-29" in beta_headers, \ + "tool-examples-2025-10-29 should be added for Opus 4.5" + + +def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): + """ + Test that advanced-tool-use-2025-11-20 header is filtered out for non-Opus 4.5 models + without adding Bedrock-specific headers. + + The translation to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 should + only happen for Claude Opus 4.5. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Simulate advanced-tool-use header from Claude Code + headers = { + "anthropic-beta": "advanced-tool-use-2025-11-20" + } + + # Test with Claude Sonnet 4.5 (not Opus 4.5) + result = config.transform_anthropic_messages_request( + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + # Verify advanced-tool-use header was removed + beta_headers = result.get("anthropic_beta", []) + assert "advanced-tool-use-2025-11-20" not in beta_headers, \ + "advanced-tool-use header should be removed for Bedrock" + + # Verify Bedrock-specific headers were NOT added (only for Opus 4.5) + assert "tool-search-tool-2025-10-19" not in beta_headers, \ + "tool-search-tool should not be added for non-Opus 4.5 models" + assert "tool-examples-2025-10-29" not in beta_headers, \ + "tool-examples should not be added for non-Opus 4.5 models" + + +def test_advanced_tool_use_header_translation_with_multiple_beta_headers(): + """ + Test that advanced-tool-use header translation works correctly when multiple + beta headers are present. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Multiple beta headers including advanced-tool-use + headers = { + "anthropic-beta": "claude-code-20250219,advanced-tool-use-2025-11-20,interleaved-thinking-2025-05-14" + } + + # Test with Claude Opus 4.5 + result = config.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20250514-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + beta_headers = result.get("anthropic_beta", []) + + # Verify advanced-tool-use was removed + assert "advanced-tool-use-2025-11-20" not in beta_headers + + # Verify Bedrock-specific headers were added + assert "tool-search-tool-2025-10-19" in beta_headers + assert "tool-examples-2025-10-29" in beta_headers + + # Verify other beta headers are preserved + assert "claude-code-20250219" in beta_headers + assert "interleaved-thinking-2025-05-14" in beta_headers + + +def test_opus_4_5_model_detection(): + """ + Test that the _is_claude_opus_4_5 method correctly identifies Opus 4.5 models + with various naming conventions. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + # Test various Opus 4.5 naming patterns + opus_4_5_models = [ + "anthropic.claude-opus-4-5-20250514-v1:0", + "anthropic.claude-opus-4.5-20250514-v1:0", + "anthropic.claude-opus_4_5-20250514-v1:0", + "anthropic.claude-opus_4.5-20250514-v1:0", + "us.anthropic.claude-opus-4-5-20250514-v1:0", + "ANTHROPIC.CLAUDE-OPUS-4-5-20250514-V1:0", # Case insensitive + ] + + for model in opus_4_5_models: + assert config._is_claude_opus_4_5(model), \ + f"Should detect {model} as Opus 4.5" + + # Test non-Opus 4.5 models + non_opus_4_5_models = [ + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-opus-4-20250514-v1:0", # Opus 4, not 4.5 + "anthropic.claude-opus-4-1-20250514-v1:0", # Opus 4.1, not 4.5 + "anthropic.claude-haiku-4-5-20251001-v1:0", + ] + + for model in non_opus_4_5_models: + assert not config._is_claude_opus_4_5(model), \ + f"Should not detect {model} as Opus 4.5" diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index b9324e4966f..e7b6de29b6b 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -390,3 +390,103 @@ class TestAnthropicBetaHeaderSupport: "anthropic_beta SHOULD be added for Anthropic models with cross-region prefix." ) assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] + + def test_messages_advanced_tool_use_translation_opus_4_5(self): + """Test that advanced-tool-use header is translated to Bedrock-specific headers for Opus 4.5. + + Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs + to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for + Bedrock Invoke API on Claude Opus 4.5. + + Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-5-20250514-v1:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers, ( + "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" + ) + + # Bedrock-specific headers should be added for Opus 4.5 + assert "tool-search-tool-2025-10-19" in beta_headers, ( + "tool-search-tool-2025-10-19 should be added for Opus 4.5" + ) + assert "tool-examples-2025-10-29" in beta_headers, ( + "tool-examples-2025-10-29 should be added for Opus 4.5" + ) + + def test_messages_advanced_tool_use_translation_sonnet_4_5(self): + """Test that advanced-tool-use header is translated to Bedrock-specific headers for Sonnet 4.5. + + Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs + to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for + Bedrock Invoke API on Claude Sonnet 4.5. + + Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers, ( + "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" + ) + + # Bedrock-specific headers should be added for Sonnet 4.5 + assert "tool-search-tool-2025-10-19" in beta_headers, ( + "tool-search-tool-2025-10-19 should be added for Sonnet 4.5" + ) + assert "tool-examples-2025-10-29" in beta_headers, ( + "tool-examples-2025-10-29 should be added for Sonnet 4.5" + ) + + def test_messages_advanced_tool_use_filtered_unsupported_model(self): + """Test that advanced-tool-use header is filtered out for models that don't support tool search. + + The translation to Bedrock-specific headers should only happen for models that + support tool search on Bedrock (Opus 4.5, Sonnet 4.5). + For other models, the advanced-tool-use header should just be removed. + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + # Test with Claude 3.5 Sonnet (does NOT support tool search on Bedrock) + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + beta_headers = result.get("anthropic_beta", []) + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers + + # Bedrock-specific headers should NOT be added for unsupported models + assert "tool-search-tool-2025-10-19" not in beta_headers + assert "tool-examples-2025-10-29" not in beta_headers 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..af0591e5f51 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, @@ -28,6 +28,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, + _get_fuzzy_user_object, _get_team_db_check, _virtual_key_max_budget_alert_check, _virtual_key_soft_budget_check, @@ -131,6 +132,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""" @@ -1277,3 +1332,42 @@ async def test_virtual_key_max_budget_alert_check_scenarios( assert ( alert_triggered == expect_alert ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + + +@pytest.mark.asyncio +async def test_get_fuzzy_user_object_case_insensitive_email(): + """Test that _get_fuzzy_user_object uses case-insensitive email lookup""" + # Setup mock Prisma client + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_usertable = MagicMock() + + # Mock user data with mixed case email + test_user = LiteLLM_UserTable( + user_id="test_123", + sso_user_id=None, + user_email="Test@Example.com", # Mixed case in DB + organization_memberships=[], + max_budget=None, + ) + + # Test: SSO ID not found, find by email with different casing + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user) + + # Search with lowercase email (different from DB) + result = await _get_fuzzy_user_object( + prisma_client=mock_prisma, + sso_user_id=None, + user_email="test@example.com", # Lowercase search + ) + + # Verify user was found despite case difference + assert result == test_user + + # Verify the query used case-insensitive mode + mock_prisma.db.litellm_usertable.find_first.assert_called_once() + call_args = mock_prisma.db.litellm_usertable.find_first.call_args + assert call_args.kwargs["where"]["user_email"]["equals"] == "test@example.com" + assert call_args.kwargs["where"]["user_email"]["mode"] == "insensitive" + assert call_args.kwargs["include"] == {"organization_memberships": True} diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 62f9cc33b64..82920ce1d80 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -8,6 +8,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( _get_customer_id_from_standard_headers, get_end_user_id_from_request_body, + get_model_from_request, get_key_model_rpm_limit, get_key_model_tpm_limit, ) @@ -186,3 +187,25 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders: request_body=request_body, request_headers=headers ) assert result == "body-user" + + +def test_get_model_from_request_supports_google_model_names_with_slashes(): + assert ( + get_model_from_request( + request_data={}, + route="/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", + ) + == "bedrock/claude-sonnet-3.7" + ) + assert ( + get_model_from_request( + request_data={}, + route="/models/hosted_vllm/gpt-oss-20b:generateContent", + ) + == "hosted_vllm/gpt-oss-20b" + ) + + +def test_get_model_from_request_vertex_passthrough_still_works(): + route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent" + assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro" 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/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index b8084906fa5..a745ac3de13 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -161,9 +161,11 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route): [ "/v1beta/models/gemini-2.5-flash:countTokens", "/v1beta/models/gemini-2.0-flash:generateContent", + "/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", "/v1beta/models/gemini-1.5-pro:streamGenerateContent", "/models/gemini-2.5-flash:countTokens", "/models/gemini-2.0-flash:generateContent", + "/models/bedrock/claude-sonnet-3.7:generateContent", "/models/gemini-1.5-pro:streamGenerateContent", ], ) @@ -187,9 +189,11 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route): "/v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent", "/v1beta/models/gemini-2.5-flash-exp:countTokens", "/v1beta/models/custom-model-name-123:streamGenerateContent", + "/v1beta/models/bedrock/claude-sonnet-3.7:generateContent", "/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent", "/models/gemini-2.5-flash-exp:countTokens", "/models/custom-model-name-123:streamGenerateContent", + "/models/bedrock/claude-sonnet-3.7:generateContent", ], ) def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(route): diff --git a/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py b/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py new file mode 100644 index 00000000000..1063f59afb6 --- /dev/null +++ b/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py @@ -0,0 +1,75 @@ +""" +Test for interactions endpoint agent parameter handling. + +Tests that the /v1beta/interactions endpoint correctly extracts +the `agent` parameter as a fallback when `model` is not provided. +""" + +import pytest + + +class TestInteractionsAgentParameter: + """Test agent parameter handling in interactions endpoint.""" + + def test_agent_parameter_fallback_logic(self): + """ + Test the core logic: model or agent extraction. + + This tests the fix in endpoints.py line ~267: + model=data.get("model") or data.get("agent") + """ + # Case 1: Only agent provided (Deep Research use case) + data = { + "agent": "deep-research-pro-preview-12-2025", + "input": "Research quantum computing", + "background": True, + } + model = data.get("model") or data.get("agent") + assert model == "deep-research-pro-preview-12-2025" + + # Case 2: Only model provided (normal use case) + data = { + "model": "gemini-2.5-flash", + "input": "Hello world", + } + model = data.get("model") or data.get("agent") + assert model == "gemini-2.5-flash" + + # Case 3: Both provided (model takes precedence) + data = { + "model": "gemini-2.5-flash", + "agent": "deep-research-pro-preview-12-2025", + "input": "Test", + } + model = data.get("model") or data.get("agent") + assert model == "gemini-2.5-flash" + + # Case 4: Neither provided + data = { + "input": "Test", + } + model = data.get("model") or data.get("agent") + assert model is None + + def test_route_type_in_skip_model_routing_list(self): + """ + Test that acreate_interaction is in the list of routes + that skip model-based routing. + + This tests the fix in route_llm_request.py. + """ + # The list of routes that skip model routing for interactions + skip_model_routing_routes = [ + "acreate_interaction", + "aget_interaction", + "adelete_interaction", + "acancel_interaction", + ] + + # acreate_interaction should be in the list (this is the fix) + assert "acreate_interaction" in skip_model_routing_routes + + # All interaction routes should be covered + assert "aget_interaction" in skip_model_routing_routes + assert "adelete_interaction" in skip_model_routing_routes + assert "acancel_interaction" in skip_model_routing_routes 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/guardrails/guardrail_hooks/test_onyx.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py index 9ede649f392..fb7480d263c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py @@ -3,6 +3,7 @@ import sys import uuid from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException from httpx import Request, Response @@ -47,20 +48,129 @@ def test_onyx_guard_config(): del os.environ["ONYX_API_KEY"] +def test_onyx_guard_with_custom_timeout_from_kwargs(): + """Test Onyx guard instantiation with custom timeout passed via kwargs.""" + # Set environment variables for testing + os.environ["ONYX_API_BASE"] = "https://test.onyx.security" + os.environ["ONYX_API_KEY"] = "test-api-key" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + + # Simulate how guardrail is instantiated from config with timeout + guardrail = OnyxGuardrail( + guardrail_name="onyx-guard-custom-timeout", + event_hook="pre_call", + default_on=True, + timeout=45.0, + ) + + # Verify the client was initialized with custom timeout + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 45.0 + assert timeout_param.connect == 5.0 + + # Clean up + if "ONYX_API_BASE" in os.environ: + del os.environ["ONYX_API_BASE"] + if "ONYX_API_KEY" in os.environ: + del os.environ["ONYX_API_KEY"] + + +def test_onyx_guard_with_timeout_none_uses_env_var(): + """Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var. + + When timeout=None is passed (as it would be from config model with default None), + the ONYX_TIMEOUT environment variable should be used. + """ + # Set environment variables for testing + os.environ["ONYX_API_BASE"] = "https://test.onyx.security" + os.environ["ONYX_API_KEY"] = "test-api-key" + os.environ["ONYX_TIMEOUT"] = "60" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + + # Pass timeout=None to simulate config model behavior + guardrail = OnyxGuardrail( + guardrail_name="onyx-guard-env-timeout", + event_hook="pre_call", + default_on=True, + timeout=None, # This triggers env var lookup + ) + + # Verify the client was initialized with timeout from env var + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 60.0 + assert timeout_param.connect == 5.0 + + # Clean up + if "ONYX_API_BASE" in os.environ: + del os.environ["ONYX_API_BASE"] + if "ONYX_API_KEY" in os.environ: + del os.environ["ONYX_API_KEY"] + if "ONYX_TIMEOUT" in os.environ: + del os.environ["ONYX_TIMEOUT"] + + +def test_onyx_guard_with_timeout_none_defaults_to_10(): + """Test Onyx guard with timeout=None and no env var defaults to 10 seconds.""" + # Set environment variables for testing + os.environ["ONYX_API_BASE"] = "https://test.onyx.security" + os.environ["ONYX_API_KEY"] = "test-api-key" + # Ensure ONYX_TIMEOUT is not set + if "ONYX_TIMEOUT" in os.environ: + del os.environ["ONYX_TIMEOUT"] + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + + # Pass timeout=None with no env var - should default to 10.0 + guardrail = OnyxGuardrail( + guardrail_name="onyx-guard-default-timeout", + event_hook="pre_call", + default_on=True, + timeout=None, + ) + + # Verify the client was initialized with default timeout of 10.0 + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 10.0 + assert timeout_param.connect == 5.0 + + # Clean up + if "ONYX_API_BASE" in os.environ: + del os.environ["ONYX_API_BASE"] + if "ONYX_API_KEY" in os.environ: + del os.environ["ONYX_API_KEY"] + + class TestOnyxGuardrail: """Test suite for Onyx Security Guardrail integration.""" def setup_method(self): """Setup test environment.""" # Clean up any existing environment variables - for key in ["ONYX_API_BASE", "ONYX_API_KEY"]: + for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]: if key in os.environ: del os.environ[key] def teardown_method(self): """Clean up test environment.""" # Clean up any environment variables set during tests - for key in ["ONYX_API_BASE", "ONYX_API_KEY"]: + for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]: if key in os.environ: del os.environ[key] @@ -103,6 +213,95 @@ class TestOnyxGuardrail: ): OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call") + def test_initialization_with_default_timeout(self): + """Test that default timeout is 10.0 seconds.""" + os.environ["ONYX_API_KEY"] = "test-api-key" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True + ) + + # Verify the client was initialized with correct timeout + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 10.0 + assert timeout_param.connect == 5.0 + + def test_initialization_with_custom_timeout_parameter(self): + """Test initialization with custom timeout parameter.""" + os.environ["ONYX_API_KEY"] = "test-api-key" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + guardrail = OnyxGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + timeout=30.0, + ) + + # Verify the client was initialized with custom timeout + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 30.0 + assert timeout_param.connect == 5.0 + + def test_initialization_with_timeout_from_env_var(self): + """Test initialization with timeout from ONYX_TIMEOUT environment variable. + + Note: The env var is only used when timeout=None is explicitly passed, + since the default parameter value is 10.0 (not None). + """ + os.environ["ONYX_API_KEY"] = "test-api-key" + os.environ["ONYX_TIMEOUT"] = "25" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + # Must pass timeout=None explicitly to trigger env var lookup + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=None + ) + + # Verify the client was initialized with timeout from env var + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 25.0 + assert timeout_param.connect == 5.0 + + def test_initialization_timeout_parameter_overrides_env_var(self): + """Test that timeout parameter overrides ONYX_TIMEOUT environment variable.""" + os.environ["ONYX_API_KEY"] = "test-api-key" + os.environ["ONYX_TIMEOUT"] = "25" + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = MagicMock() + guardrail = OnyxGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + timeout=15.0, + ) + + # Verify the client was initialized with parameter timeout (not env var) + mock_get_client.assert_called_once() + call_kwargs = mock_get_client.call_args.kwargs + timeout_param = call_kwargs["params"]["timeout"] + assert timeout_param.read == 15.0 + assert timeout_param.connect == 5.0 + @pytest.mark.asyncio async def test_apply_guardrail_request_no_violations(self): """Test apply_guardrail for request with no violations detected.""" @@ -388,6 +587,105 @@ class TestOnyxGuardrail: assert result == inputs + @pytest.mark.asyncio + async def test_apply_guardrail_timeout_error_handling(self): + """Test handling of timeout errors in apply_guardrail (graceful degradation).""" + # Set required API key + os.environ["ONYX_API_KEY"] = "test-api-key" + + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=1.0 + ) + + inputs = GenericGuardrailAPIInputs() + + request_data = { + "proxy_server_request": { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-3.5-turbo", + } + } + + # Test httpx timeout error + with patch.object( + guardrail.async_handler, "post", side_effect=httpx.TimeoutException("Request timed out") + ): + # Should return original inputs on timeout (graceful degradation) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=None, + ) + + assert result == inputs + + @pytest.mark.asyncio + async def test_apply_guardrail_read_timeout_error_handling(self): + """Test handling of read timeout errors in apply_guardrail.""" + # Set required API key + os.environ["ONYX_API_KEY"] = "test-api-key" + + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=5.0 + ) + + inputs = GenericGuardrailAPIInputs() + + request_data = { + "proxy_server_request": { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-3.5-turbo", + } + } + + # Test httpx ReadTimeout error + with patch.object( + guardrail.async_handler, "post", side_effect=httpx.ReadTimeout("Read timed out") + ): + # Should return original inputs on timeout (graceful degradation) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=None, + ) + + assert result == inputs + + @pytest.mark.asyncio + async def test_apply_guardrail_connect_timeout_error_handling(self): + """Test handling of connect timeout errors in apply_guardrail.""" + # Set required API key + os.environ["ONYX_API_KEY"] = "test-api-key" + + guardrail = OnyxGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=5.0 + ) + + inputs = GenericGuardrailAPIInputs() + + request_data = { + "proxy_server_request": { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-3.5-turbo", + } + } + + # Test httpx ConnectTimeout error + with patch.object( + guardrail.async_handler, "post", side_effect=httpx.ConnectTimeout("Connect timed out") + ): + # Should return original inputs on timeout (graceful degradation) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=None, + ) + + assert result == inputs + @pytest.mark.asyncio async def test_apply_guardrail_no_logging_obj(self): """Test apply_guardrail without logging object (uses UUID).""" 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/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a97ed93a85c..467ee3661d1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, + _validate_and_populate_member_user_info, delete_team, router, team_member_add_duplication_check, @@ -5466,3 +5467,122 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) ) and mock_db_client.db.litellm_verificationtoken.find_many.called: # If it was called, that's unexpected for admin users assert False, "API keys should not be fetched for team admin users" + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_both_provided_match(): + """ + Test _validate_and_populate_member_user_info when both user_email and user_id + are provided and they match the same user in the database. + """ + # Create member with both user_email and user_id + member = Member(user_email="test@example.com", user_id="user-123", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object that matches both email and user_id + mock_user = MagicMock() + mock_user.user_id = "user-123" + mock_user.user_email = "test@example.com" + + # Mock get_data to return single user matching email + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify result matches input (both already provided and match) + assert result.user_email == "test@example.com" + assert result.user_id == "user-123" + + # Verify get_data was called with correct parameters + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": "test@example.com"}, + table_name="user", + query_type="find_all", + ) + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_only_email_provided(): + """ + Test _validate_and_populate_member_user_info when only user_email is provided. + Should populate user_id from database. + """ + # Create member with only user_email + member = Member(user_email="test@example.com", user_id=None, role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object from find_first + mock_user_find_first = MagicMock() + mock_user_find_first.user_id = "user-456" + mock_user_find_first.user_email = "test@example.com" + + # Mock find_first to return the user + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=mock_user_find_first + ) + + # Mock get_data to return single user (no duplicates) + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user_find_first]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify user_id was populated + assert result.user_email == "test@example.com" + assert result.user_id == "user-456" + + # Verify find_first was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}} + ) + + # Verify get_data was called to check for duplicates + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": "test@example.com"}, + table_name="user", + query_type="find_all", + ) + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_only_user_id_not_found(): + """ + Test _validate_and_populate_member_user_info when only user_id is provided + but the user doesn't exist in the database. Should allow it to pass with + user_email as None (will be upserted later). + """ + # Create member with only user_id + member = Member(user_email=None, user_id="nonexistent-user", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock find_unique to return None (user not found) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + # Call the function - should NOT raise an exception + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify the result - should return member with user_id set and user_email as None + assert result.user_id == "nonexistent-user" + assert result.user_email is None + assert result.role == "user" + + # Verify find_unique was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with( + where={"user_id": "nonexistent-user"} + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index f983af2d0b2..5e9078ea876 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( GoogleSSOHandler, MicrosoftSSOHandler, SSOAuthenticationHandler, + normalize_email, ) from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, @@ -667,6 +668,85 @@ def test_build_sso_user_update_data_without_role(): assert "user_role" not in update_data +def test_normalize_email(): + """ + Test that normalize_email correctly lowercases email addresses and handles edge cases. + """ + # Test with lowercase email + assert normalize_email("test@example.com") == "test@example.com" + + # Test with uppercase email + assert normalize_email("TEST@EXAMPLE.COM") == "test@example.com" + + # Test with mixed case email + assert normalize_email("Test.User@Example.COM") == "test.user@example.com" + + # Test with None + assert normalize_email(None) is None + + # Test with empty string + assert normalize_email("") == "" + + +def test_build_sso_user_update_data_normalizes_email(): + """ + Test that _build_sso_user_update_data normalizes email addresses to lowercase. + """ + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data + + sso_result = CustomOpenID( + id="test-user-789", + email="Test.User@Example.COM", + display_name="Test User", + provider="microsoft", + team_ids=[], + user_role=None, + ) + + update_data = _build_sso_user_update_data( + result=sso_result, + user_email="Test.User@Example.COM", + user_id="test-user-789", + ) + + # Email should be normalized to lowercase + assert update_data["user_email"] == "test.user@example.com" + assert "user_role" not in update_data + + +def test_generic_response_convertor_normalizes_email(): + """ + Test that generic_response_convertor normalizes email addresses. + """ + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + mock_response = { + "preferred_username": "user123", + "email": "Test.User@Example.COM", + "sub": "Test User", + "first_name": "Test", + "last_name": "User", + "provider": "generic", + } + + # Mock JWT handler + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + # Email should be normalized to lowercase + assert result.email == "test.user@example.com" + assert result.id == "user123" + assert result.display_name == "Test User" + + @pytest.mark.asyncio async def test_upsert_sso_user_updates_role_for_existing_user(): """ 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 760c38f7811..51f4d439da7 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/test_router_silent_experiment.py b/tests/test_litellm/test_router_silent_experiment.py new file mode 100644 index 00000000000..9b82cde13c6 --- /dev/null +++ b/tests/test_litellm/test_router_silent_experiment.py @@ -0,0 +1,157 @@ +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.router import Router + + +@pytest.mark.asyncio +async def test_router_silent_experiment_acompletion(): + """ + Test that silent_model triggers a background acompletion call + and that the silent_model parameter is stripped from both calls. + """ + model_list = [ + { + "model_name": "primary-model", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "fake-key", + "silent_model": "silent-model", + }, + }, + { + "model_name": "silent-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key", + }, + }, + ] + + router = Router(model_list=model_list) + + # Mock litellm.acompletion + mock_acompletion = MagicMock() + # Create a future that resolves to a ModelResponse + mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) + future = asyncio.Future() + future.set_result(mock_response) + mock_acompletion.return_value = future + + with patch("litellm.acompletion", mock_acompletion): + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response.choices[0].message.content == "hello" + + # Give the background task a moment to trigger (it's an asyncio task) + await asyncio.sleep(0.1) + + # Should have 2 calls: one for primary, one for silent + assert mock_acompletion.call_count == 2 + + # Check call arguments + call_args_list = mock_acompletion.call_args_list + + # Verify no silent_model in any call to litellm.acompletion + for call in call_args_list: + args, kwargs = call + assert "silent_model" not in kwargs + if "metadata" in kwargs: + # One call should have is_silent_experiment=True + pass + + # Find the silent call + silent_call = next( + ( + c + for c in call_args_list + if c[1].get("metadata", {}).get("is_silent_experiment") is True + ), + None, + ) + assert silent_call is not None + assert silent_call[1]["model"] == "openai/gpt-4" + + # Find the primary call + primary_call = next( + ( + c + for c in call_args_list + if not c[1].get("metadata", {}).get("is_silent_experiment") + ), + None, + ) + assert primary_call is not None + assert primary_call[1]["model"] == "openai/gpt-3.5-turbo" + + +def test_router_silent_experiment_completion(): + """ + Test that silent_model triggers a background completion call (sync) + and that the silent_model parameter is stripped. + """ + model_list = [ + { + "model_name": "primary-model", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "fake-key", + "silent_model": "silent-model", + }, + }, + { + "model_name": "silent-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key", + }, + }, + ] + + router = Router(model_list=model_list) + + # Mock litellm.completion + mock_completion = MagicMock() + mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) + mock_completion.return_value = mock_response + + with patch("litellm.completion", mock_completion): + response = router.completion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response.choices[0].message.content == "hello" + + # The sync background call uses a thread pool. We might need to wait a bit. + import time + + time.sleep(0.5) + + # Should have 2 calls + assert mock_completion.call_count == 2 + + call_args_list = mock_completion.call_args_list + + # Verify no silent_model in any call + for call in call_args_list: + args, kwargs = call + assert "silent_model" not in kwargs + + # Find the silent call + silent_call = next( + ( + c + for c in call_args_list + if c[1].get("metadata", {}).get("is_silent_experiment") is True + ), + None, + ) + assert silent_call is not None + assert silent_call[1]["model"] == "openai/gpt-4" 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.