diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index ea8f1b0a97f..5850103132c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -118,7 +118,7 @@ class LassoGuardrail(CustomGuardrail): Falls back to UUID if ULID library is not available. """ if ULID_AVAILABLE and ulid is not None: - return str(ulid.new()) # type: ignore + return str(ulid.ULID()) # type: ignore else: verbose_proxy_logger.debug("ULID library not available, using UUID") return str(uuid.uuid4()) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d4d3359ff91..e64ddf64a43 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -304,12 +304,13 @@ if MCP_AVAILABLE: ## FastAPI Routes def _get_user_mcp_management_mode() -> UserMCPManagementMode: + proxy_general_settings: dict = {} try: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, ) except Exception: - proxy_general_settings = None + pass mode = (proxy_general_settings or {}).get("user_mcp_management_mode") if mode == "view_all": diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c1b6787132a..554e6613179 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17392,7 +17392,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/chat/completions" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17407,7 +17407,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/chat/completions" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17422,7 +17422,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/chat/completions" + "/v1/chat/completions" ], "supports_vision": true }, @@ -17435,7 +17435,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/chat/completions" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17450,7 +17450,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/chat/completions" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17641,8 +17641,8 @@ "max_tokens": 128000, "mode": "chat", "supported_endpoints": [ - "/chat/completions", - "/responses" + "/v1/chat/completions", + "/v1/responses" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17671,8 +17671,8 @@ "max_tokens": 64000, "mode": "chat", "supported_endpoints": [ - "/chat/completions", - "/responses" + "/v1/chat/completions", + "/v1/responses" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17688,7 +17688,7 @@ "max_tokens": 128000, "mode": "responses", "supported_endpoints": [ - "/responses" + "/v1/responses" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17704,8 +17704,8 @@ "max_tokens": 64000, "mode": "chat", "supported_endpoints": [ - "/chat/completions", - "/responses" + "/v1/chat/completions", + "/v1/responses" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -17737,6 +17737,8 @@ "mode": "embedding" }, "gigachat/GigaChat-2-Lite": { + "display_name": "GigaChat 2 Lite", + "model_vendor": "sberbank", "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 128000, @@ -17748,6 +17750,8 @@ "supports_system_messages": true }, "gigachat/GigaChat-2-Max": { + "display_name": "GigaChat 2 Max", + "model_vendor": "sberbank", "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 128000, @@ -17760,6 +17764,8 @@ "supports_vision": true }, "gigachat/GigaChat-2-Pro": { + "display_name": "GigaChat 2 Pro", + "model_vendor": "sberbank", "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 128000, @@ -17772,6 +17778,8 @@ "supports_vision": true }, "gigachat/Embeddings": { + "display_name": "GigaChat Embeddings", + "model_vendor": "sberbank", "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 512, @@ -17781,6 +17789,8 @@ "output_vector_size": 1024 }, "gigachat/Embeddings-2": { + "display_name": "GigaChat Embeddings 2", + "model_vendor": "sberbank", "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 512, @@ -17790,6 +17800,8 @@ "output_vector_size": 1024 }, "gigachat/EmbeddingsGigaR": { + "display_name": "GigaChat Embeddings GigaR", + "model_vendor": "sberbank", "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 4096, diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index 9a96919da87..60d4f479733 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -61,13 +61,19 @@ async def test_bedrock_apply_guardrail_blocked(): guardrailVersion="DRAFT", ) - # Mock the make_bedrock_api_request method + # Mock the make_bedrock_api_request method to raise an exception for blocked content with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock + guardrail, "make_bedrock_api_request", new_callable=AsyncMock ) as mock_api_request: - # Mock a blocked response from Bedrock - mock_response = {"action": "BLOCKED", "reason": "Content violates policy"} - mock_api_request.return_value = mock_response + # Mock the method to raise an HTTPException as it would for blocked content + from fastapi import HTTPException + mock_api_request.side_effect = HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": "", + }, + ) # Test the apply_guardrail method should raise an exception with pytest.raises(Exception) as exc_info: @@ -77,8 +83,9 @@ async def test_bedrock_apply_guardrail_blocked(): input_type="request", ) - assert "Content blocked by Bedrock guardrail" in str(exc_info.value) - assert "Content violates policy" in str(exc_info.value) + # The apply_guardrail method wraps the original exception in a generic Exception + assert "Bedrock guardrail failed:" in str(exc_info.value) + assert "Violated guardrail policy" in str(exc_info.value) @pytest.mark.asyncio @@ -253,7 +260,15 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable with patch.object( guardrail, "make_bedrock_api_request", new_callable=AsyncMock ) as mock_api: - mock_api.return_value = {"action": "BLOCKED", "reason": "policy"} + # Mock the method to raise an HTTPException as it would for blocked content + from fastapi import HTTPException + mock_api.side_effect = HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": "policy", + }, + ) with pytest.raises(Exception, match="policy") as exc_info: await guardrail.apply_guardrail( @@ -265,7 +280,8 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable assert mock_api.called _, kwargs = mock_api.call_args assert kwargs["messages"] == [request_messages[-1]] - assert "Content blocked by Bedrock guardrail" in str(exc_info.value) + # The apply_guardrail method wraps the original exception in a generic Exception + assert "Bedrock guardrail failed:" in str(exc_info.value) def test_bedrock_guardrail_filters_latest_user_message_when_enabled(): diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 7c849650bf6..ab5709cd72d 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -385,7 +385,7 @@ def test_anthropic_tool_use(tool_type, tool_config, message_content): "computer_tool_used, prompt_caching_set, expected_beta_header", [ (True, False, True), - (False, True, True), + (False, True, False), (True, True, True), (False, False, False), ], diff --git a/tests/llm_translation/test_databricks.py b/tests/llm_translation/test_databricks.py index 40fc712f2b7..3013d00288f 100644 --- a/tests/llm_translation/test_databricks.py +++ b/tests/llm_translation/test_databricks.py @@ -15,6 +15,7 @@ import litellm from litellm.exceptions import BadRequestError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper +from litellm._version import version from base_llm_unit_tests import BaseLLMChatTest, BaseAnthropicChatTest try: @@ -725,6 +726,7 @@ def test_embeddings_with_sync_http_handler(monkeypatch): headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", + "User-Agent": f"litellm/{version}", }, data=json.dumps( { @@ -767,6 +769,7 @@ def test_embeddings_with_async_http_handler(monkeypatch): headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", + "User-Agent": f"litellm/{version}", }, data=json.dumps( { @@ -823,6 +826,7 @@ def test_embeddings_uses_databricks_sdk_if_api_key_and_base_not_specified(monkey headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", + "User-Agent": f"litellm/{version}", }, data=json.dumps( { @@ -895,6 +899,7 @@ async def test_databricks_embeddings(sync_mode, monkeypatch): headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", + "User-Agent": f"litellm/{version}", }, data=json.dumps( { @@ -923,6 +928,7 @@ async def test_databricks_embeddings(sync_mode, monkeypatch): headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", + "User-Agent": f"litellm/{version}", }, data=json.dumps( { diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py index 6d6129b17bb..35ed49a84ed 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -2,6 +2,7 @@ Unit tests for Qualifire guardrail integration. """ +import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -139,76 +140,98 @@ class TestQualifireGuardrailEvaluateKwargs: @pytest.mark.asyncio async def test_evaluate_called_with_prompt_injections(self): """Test that evaluate is called with prompt_injections enabled.""" - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + # Mock the qualifire module and its types + mock_qualifire_types = MagicMock() + mock_llm_message = MagicMock() + mock_llm_tool_call = MagicMock() + mock_message_instance = MagicMock() + mock_llm_message.return_value = mock_message_instance + + mock_qualifire_types.LLMMessage = mock_llm_message + mock_qualifire_types.LLMToolCall = mock_llm_tool_call + + with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the client + mock_client = MagicMock() + mock_result = MagicMock() + mock_result.score = 100 + mock_result.status = "completed" + mock_result.evaluationResults = [] + mock_client.evaluate.return_value = mock_result + guardrail._client = mock_client - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output=None, dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert "messages" in call_kwargs + # Verify evaluate was called with correct kwargs + mock_client.evaluate.assert_called_once() + call_kwargs = mock_client.evaluate.call_args[1] + assert call_kwargs["prompt_injections"] is True + assert "messages" in call_kwargs @pytest.mark.asyncio async def test_evaluate_called_with_multiple_checks(self): """Test that evaluate is called with multiple checks enabled.""" - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + # Mock the qualifire module and its types + mock_qualifire_types = MagicMock() + mock_llm_message = MagicMock() + mock_llm_tool_call = MagicMock() + mock_message_instance = MagicMock() + mock_llm_message.return_value = mock_message_instance + + mock_qualifire_types.LLMMessage = mock_llm_message + mock_qualifire_types.LLMToolCall = mock_llm_tool_call + + with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - pii_check=True, - hallucinations_check=True, - assertions=["Output must be valid JSON"], - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + pii_check=True, + hallucinations_check=True, + assertions=["Output must be valid JSON"], + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the client + mock_client = MagicMock() + mock_result = MagicMock() + mock_result.score = 100 + mock_result.status = "completed" + mock_result.evaluationResults = [] + mock_client.evaluate.return_value = mock_result + guardrail._client = mock_client - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output="Test output", dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert call_kwargs["pii_check"] is True - assert call_kwargs["hallucinations_check"] is True - assert call_kwargs["assertions"] == ["Output must be valid JSON"] - assert call_kwargs["output"] == "Test output" + # Verify evaluate was called with correct kwargs + mock_client.evaluate.assert_called_once() + call_kwargs = mock_client.evaluate.call_args[1] + assert call_kwargs["prompt_injections"] is True + assert call_kwargs["pii_check"] is True + assert call_kwargs["hallucinations_check"] is True + assert call_kwargs["assertions"] == ["Output must be valid JSON"] + assert call_kwargs["output"] == "Test output" class TestQualifireGuardrailCheckIfFlagged: diff --git a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py index 011031c1e4f..97c1733a935 100644 --- a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py @@ -40,6 +40,7 @@ class TestKeyManagementEventHooksIndependentOperations: mock_data = MagicMock() mock_data.key_alias = "test-key-alias" mock_data.team_id = None + mock_data.send_invite_email = True mock_response = MagicMock() mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"} @@ -59,6 +60,10 @@ class TestKeyManagementEventHooksIndependentOperations: KeyManagementEventHooks, "_store_virtual_key_in_secret_manager", side_effect=mock_store_secret, + ), patch.object( + KeyManagementEventHooks, + "_is_email_sending_enabled", + return_value=True, ), patch( "litellm.store_audit_logs", False ), patch( @@ -96,6 +101,7 @@ class TestKeyManagementEventHooksIndependentOperations: mock_data = MagicMock() mock_data.key_alias = "test-key-alias" mock_data.team_id = None + mock_data.send_invite_email = True mock_response = MagicMock() mock_response.model_dump.return_value = {"key": "sk-test", "token": "test-token"} @@ -115,6 +121,10 @@ class TestKeyManagementEventHooksIndependentOperations: KeyManagementEventHooks, "_store_virtual_key_in_secret_manager", side_effect=mock_store_secret_raises, + ), patch.object( + KeyManagementEventHooks, + "_is_email_sending_enabled", + return_value=True, ), patch( "litellm.store_audit_logs", False ), patch( 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 8fdfd6897a8..ad4f53dac4b 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 @@ -742,18 +742,16 @@ class TestProxySettingEndpoints: ): """Test updating UI settings with an allowlisted field""" from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy._types import UserAPIKeyAuth - class MockUser: - def __init__(self, user_role): - self.user_role = user_role - - async def mock_admin_auth(): - return MockUser(LitellmUserRoles.PROXY_ADMIN) - - monkeypatch.setattr( - "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.user_api_key_auth", - mock_admin_auth, + # Override the FastAPI dependency with a proper mock + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) mock_prisma = MagicMock() mock_prisma.db.litellm_uisettings.upsert = AsyncMock() @@ -761,7 +759,11 @@ class TestProxySettingEndpoints: payload = {"disable_model_add_for_internal_users": True} - response = client.patch("/update/ui_settings", json=payload) + try: + response = client.patch("/update/ui_settings", json=payload) + finally: + # Clean up the dependency override + app.dependency_overrides.clear() assert response.status_code == 200 data = response.json() @@ -780,18 +782,16 @@ class TestProxySettingEndpoints: ): """Test non-allowlisted UI settings are ignored on update""" from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy._types import UserAPIKeyAuth - class MockUser: - def __init__(self, user_role): - self.user_role = user_role - - async def mock_admin_auth(): - return MockUser(LitellmUserRoles.PROXY_ADMIN) - - monkeypatch.setattr( - "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.user_api_key_auth", - mock_admin_auth, + # Override the FastAPI dependency with a proper mock + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) mock_prisma = MagicMock() mock_prisma.db.litellm_uisettings.upsert = AsyncMock() @@ -802,7 +802,11 @@ class TestProxySettingEndpoints: "unsupported_flag": True, } - response = client.patch("/update/ui_settings", json=payload) + try: + response = client.patch("/update/ui_settings", json=payload) + finally: + # Clean up the dependency override + app.dependency_overrides.clear() assert response.status_code == 200 data = response.json()