scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

Merge branch 'main' into litellm_add_openai_response_compression_api

Browse source
This commit is contained in:
Ishaan Jaff 2026-01-06 16:22:20 +05:30 • committed by GitHub
commit 33d7e26923
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1531 additions and 5656 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
View file

@ -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())

3
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -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 = {}
pass
mode = proxy_general_settings.get("user_mcp_management_mode")
if mode == "view_all":

6887
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

34
tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py
View file

@ -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():

2
tests/llm_translation/test_anthropic_completion.py
View file

@ -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),
],

6
tests/llm_translation/test_databricks.py
View file

@ -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(
{

135
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py
View file

@ -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:

10
tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py
View file

@ -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(

48
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
View file

@ -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()

57
tests/test_litellm/test_model_prices_metadata.py
View file

@ -1,57 +0,0 @@
"""
Tests to validate model_prices_and_context_window.json metadata fields.
Ensures all model entries have required metadata fields:
- display_name: Human-readable name
- model_vendor: Vendor/provider identifier (e.g., "openai", "anthropic", "meta")
"""
import json
import pytest
from pathlib import Path
@pytest.fixture(scope="module")
def model_prices_data():
"""Load the model prices JSON file."""
possible_paths = [
Path(__file__).parent.parent.parent.parent / "model_prices_and_context_window.json",
Path("model_prices_and_context_window.json"),
]
for path in possible_paths:
if path.exists():
with open(path, "r") as f:
return json.load(f)
pytest.fail("Could not find model_prices_and_context_window.json")
class TestModelMetadataPresence:
"""Test that required metadata fields are present."""
def test_all_models_have_display_name(self, model_prices_data):
"""Every model entry should have a display_name."""
missing = []
for model_key, model_data in model_prices_data.items():
if "display_name" not in model_data:
missing.append(model_key)
if missing:
pytest.fail(
f"Missing display_name in {len(missing)} models. "
f"First 10: {missing[:10]}"
)
def test_all_models_have_model_vendor(self, model_prices_data):
"""Every model entry should have a model_vendor."""
missing = []
for model_key, model_data in model_prices_data.items():
if "model_vendor" not in model_data:
missing.append(model_key)
if missing:
pytest.fail(
f"Missing model_vendor in {len(missing)} models. "
f"First 10: {missing[:10]}"
)

3
tests/test_litellm/test_utils.py
View file

@ -507,9 +507,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"additionalProperties": {
"type": "object",
"properties": {
"display_name": {"type": "string"},
"model_vendor": {"type": "string"},
"model_version": {"type": "string"},
"supports_computer_use": {"type": "boolean"},
"cache_creation_input_audio_token_cost": {"type": "number"},
"cache_creation_input_token_cost": {"type": "number"},