scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35
litellm/tests/proxy_unit_tests/test_proxy_server.py
Yassin Kortam c6b2f111a6
fix(team): make team member add atomic to prevent concurrent-add member loss (#34185) 
_add_team_members_to_team reconciled membership by reading the complete_team_data
snapshot captured at the start of team_member_add, appending in memory, and
writing the whole members_with_roles array back. Two concurrent /team/member_add
calls for the same team read the same snapshot, so the last write wins and one
member is silently lost. This affects every concurrent team member add, including
the SCIM group PATCH op:add path that routes through team_member_add

Reconcile members_with_roles inside a transaction that locks the team row with
SELECT ... FOR UPDATE before re-reading the current membership, so concurrent
writers serialize on the row lock and each appends onto the other's committed
result. The interactive transaction is exposed through a thin PrismaClient.tx()
passthrough and the locked read is encapsulated in
TeamRepository.get_members_with_roles_locked, and the SCIM group PATCH applies
membership as deltas so concurrent adds are not clobbered
2026-07-22 21:47:22 +00:00

3054 lines
102 KiB
Python
Raw Blame History

import os
import sys
import traceback
from unittest import mock
from dotenv import load_dotenv
import litellm.proxy
import litellm.proxy.proxy_server
load_dotenv()
import io
import json
import os
# this file is to test litellm/proxy
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import logging
import pytest
import litellm
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
# Configure logging
logging.basicConfig(
level=logging.DEBUG, # Set the desired logging level
format="%(asctime)s - %(levelname)s - %(message)s",
)
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import FastAPI
# test /chat/completion request to the proxy
from fastapi.testclient import TestClient
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.proxy_server import ( # Replace with the actual module where your FastAPI router is defined
app,
initialize,
save_worker_config,
)
from litellm.proxy.utils import ProxyLogging
# Your bearer token
token = "sk-1234"
headers = {"Authorization": f"Bearer {token}"}
example_completion_result = {
"choices": [
{
"message": {
"content": "Whispers of the wind carry dreams to me.",
"role": "assistant",
}
}
],
}
example_embedding_result = {
"object": "list",
"data": [
{
"object": "embedding",
"index": 0,
"embedding": [
-0.006929283495992422,
-0.005336422007530928,
-4.547132266452536e-05,
-0.024047505110502243,
-0.006929283495992422,
-0.005336422007530928,
-4.547132266452536e-05,
-0.024047505110502243,
-0.006929283495992422,
-0.005336422007530928,
-4.547132266452536e-05,
-0.024047505110502243,
],
}
],
"model": "text-embedding-3-small",
"usage": {"prompt_tokens": 5, "total_tokens": 5},
}
example_image_generation_result = {
"created": 1589478378,
"data": [{"url": "https://..."}, {"url": "https://..."}],
}
def mock_patch_acompletion():
return mock.patch(
"litellm.proxy.proxy_server.llm_router.acompletion",
return_value=example_completion_result,
)
def mock_patch_aembedding():
return mock.patch(
"litellm.proxy.proxy_server.llm_router.aembedding",
return_value=example_embedding_result,
)
def mock_patch_aimage_generation():
return mock.patch(
"litellm.proxy.proxy_server.llm_router.aimage_generation",
return_value=example_image_generation_result,
)
@pytest.fixture(scope="function")
def fake_env_vars(monkeypatch):
# Set some fake environment variables
monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key")
monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base")
monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base")
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key")
monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base")
monkeypatch.setenv("REDIS_HOST", "localhost")
@pytest.fixture(scope="function")
def client_no_auth(fake_env_vars):
# Assuming litellm.proxy.proxy_server is an object
from litellm.proxy.proxy_server import cleanup_router_config_variables
cleanup_router_config_variables()
filepath = os.path.dirname(os.path.abspath(__file__))
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
# initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables
asyncio.run(initialize(config=config_fp, debug=True))
return TestClient(app)
@mock_patch_acompletion()
def test_chat_completion(mock_acompletion, client_no_auth):
global headers
try:
# Your test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi"},
],
"max_tokens": 10,
}
print("testing proxy server with chat completions")
response = client_no_auth.post("/v1/chat/completions", json=test_data)
mock_acompletion.assert_called_once_with(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "hi"},
],
max_tokens=10,
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
specific_deployment=True,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
print(f"response - {response.text}")
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
def test_chat_completion_malformed_messages_returns_400(client_no_auth):
"""
Test that malformed messages (strings instead of dicts) return 400 instead of 500.
This test verifies that when a client sends messages as raw strings instead of
{role, content} objects, LiteLLM returns a 400 invalid_request_error instead
of a 500 Internal Server Error.
"""
global headers
try:
# Test data with malformed messages (string instead of dict)
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
"hi how are you"
], # Invalid: should be [{"role": "user", "content": "hi how are you"}]
}
print("testing proxy server with malformed messages")
response = client_no_auth.post(
"/v1/chat/completions", json=test_data, headers=headers
)
print(f"response status: {response.status_code}")
print(f"response text: {response.text}")
# Should return 400, not 500
assert (
response.status_code == 400
), f"Expected 400, got {response.status_code}. Response: {response.text}"
# Verify error format
result = response.json()
assert "error" in result, "Response should contain 'error' key"
error = result["error"]
# Verify error type and message
assert (
error.get("type") == "invalid_request_error" or error.get("type") is None
), f"Expected invalid_request_error or None, got {error.get('type')}"
assert (
error.get("code") == "400" or error.get("code") == 400
), f"Expected code 400, got {error.get('code')}"
# Error message should indicate invalid request format
error_message = error.get("message", "")
assert len(error_message) > 0, "Error message should not be empty"
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
def test_get_settings_request_timeout(client_no_auth):
"""
When no timeout is set, it should use the litellm.request_timeout value
"""
# Set a known value for litellm.request_timeout
import litellm
# Make a GET request to /settings
response = client_no_auth.get("/settings")
# Check if the request was successful
assert response.status_code == 200
# Parse the JSON response
settings = response.json()
print("settings", settings)
assert settings["litellm.request_timeout"] == litellm.request_timeout
@pytest.mark.parametrize(
"litellm_key_header_name",
["x-litellm-key", None],
)
def test_add_headers_to_request(litellm_key_header_name):
from fastapi import Request
from starlette.datastructures import URL
import json
from litellm.proxy.litellm_pre_call_utils import (
clean_headers,
LiteLLMProxyRequestSetup,
)
headers = {
"Authorization": "Bearer 1234",
"X-Custom-Header": "Custom-Value",
"X-Stainless-Header": "Stainless-Value",
"anthropic-beta": "beta-value",
}
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
request._body = json.dumps({"model": "gpt-3.5-turbo"}).encode("utf-8")
request_headers = clean_headers(headers, litellm_key_header_name)
forwarded_headers = LiteLLMProxyRequestSetup._get_forwardable_headers(
request_headers
)
assert forwarded_headers == {
"X-Custom-Header": "Custom-Value",
"anthropic-beta": "beta-value",
}
@pytest.mark.parametrize(
"litellm_key_header_name",
["x-litellm-key", None],
)
@pytest.mark.parametrize(
"forward_headers",
[True, False],
)
@mock_patch_acompletion()
def test_chat_completion_forward_headers(
mock_acompletion, client_no_auth, litellm_key_header_name, forward_headers
):
global headers
try:
if forward_headers:
gs = getattr(litellm.proxy.proxy_server, "general_settings")
gs["forward_client_headers_to_llm_api"] = True
setattr(litellm.proxy.proxy_server, "general_settings", gs)
if litellm_key_header_name is not None:
gs = getattr(litellm.proxy.proxy_server, "general_settings")
gs["litellm_key_header_name"] = litellm_key_header_name
setattr(litellm.proxy.proxy_server, "general_settings", gs)
# Your test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi"},
],
"max_tokens": 10,
}
headers_to_forward = {
"X-Custom-Header": "Custom-Value",
"X-Another-Header": "Another-Value",
}
if litellm_key_header_name is not None:
headers_to_not_forward = {litellm_key_header_name: "Bearer 1234"}
else:
headers_to_not_forward = {"Authorization": "Bearer 1234"}
received_headers = {**headers_to_forward, **headers_to_not_forward}
print("testing proxy server with chat completions")
response = client_no_auth.post(
"/v1/chat/completions", json=test_data, headers=received_headers
)
if not forward_headers:
assert "headers" not in mock_acompletion.call_args.kwargs
else:
assert mock_acompletion.call_args.kwargs["headers"] == {
"x-custom-header": "Custom-Value",
"x-another-header": "Another-Value",
}
print(f"response - {response.text}")
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@pytest.mark.parametrize("forward_llm_auth_headers", [True, False])
@mock_patch_acompletion()
def test_chat_completion_forward_llm_provider_auth_headers(
mock_acompletion, client_no_auth, forward_llm_auth_headers
):
"""
Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded
when forward_llm_provider_auth_headers=True.
This allows clients to send their own LLM provider API keys through the proxy.
"""
try:
# Configure general settings
gs = getattr(litellm.proxy.proxy_server, "general_settings")
gs["forward_client_headers_to_llm_api"] = True
gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers
setattr(litellm.proxy.proxy_server, "general_settings", gs)
# Test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hello"},
],
"max_tokens": 10,
}
# Headers including LLM provider auth
request_headers = {
"Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped)
"x-api-key": "sk-ant-api03-test-anthropic-key", # Anthropic API key
"x-goog-api-key": "google-api-key-123", # Google API key
"X-Custom-Header": "custom-value", # Custom header (should be forwarded)
}
# Make request
response = client_no_auth.post(
"/v1/chat/completions", json=test_data, headers=request_headers
)
assert response.status_code == 200
# Check forwarded headers
forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {})
if forward_llm_auth_headers:
# LLM provider auth headers should be forwarded
assert "x-api-key" in forwarded_headers
assert forwarded_headers["x-api-key"] == "sk-ant-api03-test-anthropic-key"
assert "x-goog-api-key" in forwarded_headers
assert forwarded_headers["x-goog-api-key"] == "google-api-key-123"
else:
# LLM provider auth headers should be stripped
assert "x-api-key" not in forwarded_headers
assert "x-goog-api-key" not in forwarded_headers
# Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True)
assert "x-custom-header" in forwarded_headers
assert forwarded_headers["x-custom-header"] == "custom-value"
# Proxy Authorization should never be forwarded
assert "authorization" not in forwarded_headers
print(
f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}"
)
print(f" Forwarded headers: {list(forwarded_headers.keys())}")
except Exception as e:
pytest.fail(
f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}"
)
finally:
# Clean up
gs = getattr(litellm.proxy.proxy_server, "general_settings")
gs.pop("forward_llm_provider_auth_headers", None)
setattr(litellm.proxy.proxy_server, "general_settings", gs)
@mock_patch_acompletion()
@pytest.mark.asyncio
async def test_team_disable_guardrails(mock_acompletion, client_no_auth):
"""
If team not allowed to turn on/off guardrails
Raise 403 forbidden error, if request is made by team on `/key/generate` or `/chat/completions`.
"""
import asyncio
import json
import time
from fastapi import HTTPException, Request
from starlette.datastructures import URL
from litellm.proxy._types import (
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
_team_id = "1234"
user_key = "sk-12345678"
valid_token = UserAPIKeyAuth(
team_id=_team_id,
team_blocked=True,
token=hash_token(user_key),
last_refreshed_at=time.time(),
)
await asyncio.sleep(1)
team_obj = LiteLLM_TeamTableCachedObj(
team_id=_team_id,
blocked=False,
last_refreshed_at=time.time(),
metadata={"guardrails": {"modify_guardrails": False}},
)
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
body = {"metadata": {"guardrails": {"hide_secrets": False}}}
json_bytes = json.dumps(body).encode("utf-8")
request._body = json_bytes
try:
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
pytest.fail("Expected to raise 403 forbidden error.")
except ProxyException as e:
assert e.code == str(403)
from test_custom_callback_input import CompletionCustomHandler
@mock_patch_acompletion()
def test_custom_logger_failure_handler(mock_acompletion, client_no_auth):
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
rpm_limit = 0
mock_api_key = "sk-my-test-key"
cache_value = UserAPIKeyAuth(token=hash_token(mock_api_key), rpm_limit=rpm_limit)
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
mock_logger = CustomLogger()
mock_logger_unit_tests = CompletionCustomHandler()
proxy_logging_obj: ProxyLogging = getattr(
litellm.proxy.proxy_server, "proxy_logging_obj"
)
litellm.callbacks = [mock_logger, mock_logger_unit_tests]
proxy_logging_obj._init_litellm_callbacks(llm_router=None)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
with patch.object(
mock_logger, "async_log_failure_event", new=AsyncMock()
) as mock_failed_alert:
# Your test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi"},
],
"max_tokens": 10,
}
print("testing proxy server with chat completions")
response = client_no_auth.post(
"/v1/chat/completions",
json=test_data,
headers={"Authorization": "Bearer {}".format(mock_api_key)},
)
assert response.status_code == 429
# confirm async_log_failure_event is called
mock_failed_alert.assert_called()
assert len(mock_logger_unit_tests.errors) == 0
@mock_patch_acompletion()
def test_engines_model_chat_completions(mock_acompletion, client_no_auth):
global headers
try:
# Your test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi"},
],
"max_tokens": 10,
}
print("testing proxy server with chat completions")
response = client_no_auth.post(
"/engines/gpt-3.5-turbo/chat/completions", json=test_data
)
mock_acompletion.assert_called_once_with(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "hi"},
],
max_tokens=10,
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
specific_deployment=True,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
print(f"response - {response.text}")
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@mock_patch_acompletion()
def test_chat_completion_azure(mock_acompletion, client_no_auth):
global headers
try:
# Your test data
test_data = {
"model": "azure/gpt-4.1-mini",
"messages": [
{"role": "user", "content": "write 1 sentence poem"},
],
"max_tokens": 10,
}
print("testing proxy server with Azure Request /chat/completions")
response = client_no_auth.post("/v1/chat/completions", json=test_data)
mock_acompletion.assert_called_once_with(
model="azure/gpt-4.1-mini",
messages=[
{"role": "user", "content": "write 1 sentence poem"},
],
max_tokens=10,
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
specific_deployment=True,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
assert len(result["choices"][0]["message"]["content"]) > 0
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test
# test_chat_completion_azure()
@mock_patch_acompletion()
def test_openai_deployments_model_chat_completions_azure(
mock_acompletion, client_no_auth
):
global headers
try:
# Your test data
test_data = {
"model": "azure/gpt-4.1-mini",
"messages": [
{"role": "user", "content": "write 1 sentence poem"},
],
"max_tokens": 10,
}
url = "/openai/deployments/azure/gpt-4.1-mini/chat/completions"
print(f"testing proxy server with Azure Request {url}")
response = client_no_auth.post(url, json=test_data)
mock_acompletion.assert_called_once_with(
model="azure/gpt-4.1-mini",
messages=[
{"role": "user", "content": "write 1 sentence poem"},
],
max_tokens=10,
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
specific_deployment=True,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
assert len(result["choices"][0]["message"]["content"]) > 0
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test
# test_openai_deployments_model_chat_completions_azure()
### EMBEDDING
@mock_patch_aembedding()
def test_embedding(mock_aembedding, client_no_auth):
global headers
from litellm.proxy.proxy_server import user_custom_auth
try:
test_data = {
"model": "azure/text-embedding-ada-002",
"input": ["good morning from litellm"],
}
async def _pre_call_hook_side_effect(**kwargs):
data = kwargs["data"]
metadata = {**(data.get("metadata") or {}), "source": "unit-test"}
data["metadata"] = metadata
proxy_request = {**(data.get("proxy_server_request") or {})}
proxy_request["path"] = "/v1/embeddings"
data["proxy_server_request"] = proxy_request
return data
async def _post_call_success_side_effect(**kwargs):
return kwargs["response"]
with (
patch.object(
litellm.proxy.proxy_server.proxy_logging_obj,
"pre_call_hook",
new=AsyncMock(side_effect=_pre_call_hook_side_effect),
) as mock_pre_call_hook,
patch.object(
litellm.proxy.proxy_server.proxy_logging_obj,
"during_call_hook",
new=AsyncMock(return_value=None),
) as mock_during_hook,
patch.object(
litellm.proxy.proxy_server.proxy_logging_obj,
"post_call_success_hook",
new=AsyncMock(side_effect=_post_call_success_side_effect),
),
):
response = client_no_auth.post("/v1/embeddings", json=test_data)
mock_aembedding.assert_called_once_with(
model="azure/text-embedding-ada-002",
input=["good morning from litellm"],
specific_deployment=True,
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
assert response.status_code == 200
result = response.json()
print(len(result["data"][0]["embedding"]))
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
call_metadata = mock_aembedding.call_args.kwargs["metadata"]
assert call_metadata.get("source") == "unit-test"
pre_call_kwargs = mock_pre_call_hook.await_args_list[0].kwargs
assert (
pre_call_kwargs.get("call_type") == "aembedding"
), f"expected pre_call_hook to receive call_type='aembedding', got {pre_call_kwargs.get('call_type')}"
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@mock_patch_aembedding()
def test_bedrock_embedding(mock_aembedding, client_no_auth):
global headers
from litellm.proxy.proxy_server import user_custom_auth
try:
test_data = {
"model": "amazon-embeddings",
"input": ["good morning from litellm"],
}
response = client_no_auth.post("/v1/embeddings", json=test_data)
mock_aembedding.assert_called_once_with(
model="amazon-embeddings",
input=["good morning from litellm"],
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
assert response.status_code == 200
print(response.status_code, response.text)
result = response.json()
print(len(result["data"][0]["embedding"]))
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@pytest.mark.skip(reason="AWS Suspended Account")
def test_sagemaker_embedding(client_no_auth):
global headers
from litellm.proxy.proxy_server import user_custom_auth
try:
test_data = {
"model": "GPT-J 6B - Sagemaker Text Embedding (Internal)",
"input": ["good morning from litellm"],
}
response = client_no_auth.post("/v1/embeddings", json=test_data)
assert response.status_code == 200
result = response.json()
print(len(result["data"][0]["embedding"]))
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test
# test_embedding()
#### IMAGE GENERATION
@mock_patch_aimage_generation()
def test_img_gen(mock_aimage_generation, client_no_auth):
global headers
from litellm.proxy.proxy_server import user_custom_auth
try:
test_data = {
"model": "dall-e-3",
"prompt": "A cute baby sea otter",
"n": 1,
"size": "1024x1024",
"imageConfig": {"aspectRatio": "9:16", "imageSize": "1K"},
}
response = client_no_auth.post("/v1/images/generations", json=test_data)
mock_aimage_generation.assert_called_once_with(
model="dall-e-3",
prompt="A cute baby sea otter",
n=1,
size="1024x1024",
imageConfig={"aspectRatio": "9:16", "imageSize": "1K"},
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
assert response.status_code == 200
result = response.json()
print(len(result["data"][0]["url"]))
assert len(result["data"][0]["url"]) > 10
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
#### ADDITIONAL
@pytest.mark.skip(reason="test via docker tests. Requires prisma client.")
def test_add_new_model(client_no_auth):
global headers
try:
test_data = {
"model_name": "test_openai_models",
"litellm_params": {
"model": "gpt-3.5-turbo",
},
"model_info": {"description": "this is a test openai model"},
}
client_no_auth.post("/model/new", json=test_data, headers=headers)
response = client_no_auth.get("/model/info", headers=headers)
assert response.status_code == 200
result = response.json()
print(f"response: {result}")
model_info = None
for m in result["data"]:
if m["model_name"] == "test_openai_models":
model_info = m["model_info"]
assert model_info["description"] == "this is a test openai model"
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")
@pytest.mark.xdist_group("proxy_heavy")
def test_health(client_no_auth):
global headers
import logging
import time
from litellm._logging import verbose_logger, verbose_proxy_logger
verbose_proxy_logger.setLevel(logging.DEBUG)
try:
response = client_no_auth.get("/health")
assert response.status_code == 200
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# test_add_new_model()
from litellm.integrations.custom_logger import CustomLogger
class MyCustomHandler(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
print(f"Pre-API Call")
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Success")
assert kwargs["user"] == "proxy-user"
assert kwargs["model"] == "gpt-3.5-turbo"
assert kwargs["max_tokens"] == 10
customHandler = MyCustomHandler()
@mock_patch_acompletion()
def test_chat_completion_optional_params(mock_acompletion, client_no_auth):
# [PROXY: PROD TEST] - DO NOT DELETE
# This tests if all the /chat/completion params are passed to litellm
try:
# Your test data
litellm.set_verbose = True
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi"},
],
"max_tokens": 10,
"user": "proxy-user",
}
litellm.callbacks = [customHandler]
print("testing proxy server: optional params")
response = client_no_auth.post("/v1/chat/completions", json=test_data)
mock_acompletion.assert_called_once_with(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "hi"},
],
max_tokens=10,
user="proxy-user",
litellm_call_id=mock.ANY,
litellm_logging_obj=mock.ANY,
request_timeout=mock.ANY,
specific_deployment=True,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,
)
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
# Run the test
# test_chat_completion_optional_params()
# Test Reading config.yaml file
from litellm.proxy.proxy_server import ProxyConfig
@pytest.mark.skip(reason="local variable conflicts. needs to be refactored.")
@mock.patch("litellm.proxy.proxy_server.litellm.Cache")
def test_load_router_config(mock_cache, fake_env_vars):
mock_cache.return_value.cache.__dict__ = {"redis_client": None}
mock_cache.return_value.supported_call_types = [
"completion",
"acompletion",
"embedding",
"aembedding",
"atranscription",
"transcription",
]
try:
import asyncio
print("testing reading config")
# this is a basic config.yaml with only a model
filepath = os.path.dirname(os.path.abspath(__file__))
proxy_config = ProxyConfig()
result = asyncio.run(
proxy_config.load_config(
router=None,
config_file_path=f"{filepath}/example_config_yaml/simple_config.yaml",
)
)
print(result)
assert len(result[1]) == 1
# this is a load balancing config yaml
result = asyncio.run(
proxy_config.load_config(
router=None,
config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml",
)
)
print(result)
assert len(result[1]) == 2
# config with general settings - custom callbacks
result = asyncio.run(
proxy_config.load_config(
router=None,
config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml",
)
)
print(result)
assert len(result[1]) == 2
# tests for litellm.cache set from config
print("testing reading proxy config for cache")
litellm.cache = None
asyncio.run(
proxy_config.load_config(
router=None,
config_file_path=f"{filepath}/example_config_yaml/cache_no_params.yaml",
)
)
assert litellm.cache is not None
assert "redis_client" in vars(
litellm.cache.cache
) # it should default to redis on proxy
assert litellm.cache.supported_call_types == [
"completion",
"acompletion",
"embedding",
"aembedding",
"atranscription",
"transcription",
] # init with all call types
litellm.disable_cache()
print("testing reading proxy config for cache with params")
mock_cache.return_value.supported_call_types = [
"embedding",
"aembedding",
]
asyncio.run(
proxy_config.load_config(
router=None,
config_file_path=f"{filepath}/example_config_yaml/cache_with_params.yaml",
)
)
assert litellm.cache is not None
print(litellm.cache)
print(litellm.cache.supported_call_types)
print(vars(litellm.cache.cache))
assert "redis_client" in vars(
litellm.cache.cache
) # it should default to redis on proxy
assert litellm.cache.supported_call_types == [
"embedding",
"aembedding",
] # init with all call types
except Exception as e:
pytest.fail(
f"Proxy: Got exception reading config: {str(e)}\n{traceback.format_exc()}"
)
# test_load_router_config()
@pytest.mark.asyncio
async def test_team_update_redis():
"""
Tests if team update, updates the redis cache if set
"""
from litellm.caching.caching import DualCache, RedisCache
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
from litellm.proxy.auth.auth_checks import _cache_team_object
proxy_logging_obj: ProxyLogging = getattr(
litellm.proxy.proxy_server, "proxy_logging_obj"
)
redis_cache = RedisCache(host="localhost")
with patch.object(
redis_cache,
"async_set_cache",
new=AsyncMock(),
) as mock_client:
await _cache_team_object(
team_id="1234",
team_table=LiteLLM_TeamTableCachedObj(team_id="1234"),
user_api_key_cache=DualCache(redis_cache=redis_cache),
proxy_logging_obj=proxy_logging_obj,
)
mock_client.assert_called()
@pytest.mark.asyncio
async def test_get_team_redis(client_no_auth):
"""
Tests if get_team_object gets value from redis cache, if set
"""
from litellm.caching.caching import DualCache, RedisCache
from litellm.proxy.auth.auth_checks import get_team_object
proxy_logging_obj: ProxyLogging = getattr(
litellm.proxy.proxy_server, "proxy_logging_obj"
)
redis_cache = RedisCache()
from fastapi import HTTPException
with patch.object(
redis_cache,
"async_get_cache",
new=AsyncMock(),
) as mock_client:
try:
await get_team_object(
team_id="1234",
user_api_key_cache=DualCache(redis_cache=redis_cache),
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
prisma_client=AsyncMock(),
)
except HTTPException:
pass
mock_client.assert_called_once()
import random
from litellm._uuid import uuid
from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch
from litellm.proxy._types import (
LitellmUserRoles,
NewUserRequest,
TeamMemberAddRequest,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
from test_key_generate_prisma import prisma_client
@pytest.fixture
def mock_prisma_client():
client = MagicMock()
client.connect = AsyncMock()
client.disconnect = AsyncMock()
return client
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
@pytest.mark.parametrize(
"user_role",
[LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.PROXY_ADMIN.value],
)
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_create_user_default_budget(prisma_client, user_role):
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm, "max_internal_user_budget", 10)
setattr(litellm, "internal_user_budget_duration", "5m")
await litellm.proxy.proxy_server.prisma_client.connect()
user = f"ishaan {uuid.uuid4().hex}"
request = NewUserRequest(
user_id=user, user_role=user_role
) # create a key with no budget
with patch.object(
litellm.proxy.proxy_server.prisma_client, "insert_data", new=AsyncMock()
) as mock_client:
await new_user(
request,
)
mock_client.assert_called()
print(f"mock_client.call_args: {mock_client.call_args}")
print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs))
if user_role == LitellmUserRoles.INTERNAL_USER.value:
assert (
mock_client.call_args.kwargs["data"]["max_budget"]
== litellm.max_internal_user_budget
)
assert (
mock_client.call_args.kwargs["data"]["budget_duration"]
== litellm.internal_user_budget_duration
)
else:
assert mock_client.call_args.kwargs["data"]["max_budget"] is None
assert mock_client.call_args.kwargs["data"]["budget_duration"] is None
@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_create_team_member_add(prisma_client, new_member_method):
import time
from fastapi import Request
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm, "max_internal_user_budget", 10)
setattr(litellm, "internal_user_budget_duration", "5m")
await litellm.proxy.proxy_server.prisma_client.connect()
user = f"ishaan {uuid.uuid4().hex}"
_team_id = "litellm-test-client-id-new"
team_obj = LiteLLM_TeamTableCachedObj(
team_id=_team_id,
blocked=False,
last_refreshed_at=time.time(),
metadata={"guardrails": {"modify_guardrails": False}},
)
# user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
if new_member_method == "user_id":
data = {
"team_id": _team_id,
"member": [{"role": "user", "user_id": user}],
}
elif new_member_method == "user_email":
data = {
"team_id": _team_id,
"member": [{"role": "user", "user_email": user}],
}
team_member_add_request = TeamMemberAddRequest(**data)
with (
patch(
"litellm.proxy.proxy_server.prisma_client.db.litellm_usertable",
new_callable=AsyncMock,
) as mock_litellm_usertable,
patch(
"litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache",
new=AsyncMock(return_value=team_obj),
) as mock_team_obj,
patch(
"litellm.proxy.proxy_server.prisma_client.get_data",
new=AsyncMock(return_value=[]),
) as mock_get_data,
):
mock_client = AsyncMock(
return_value=LiteLLM_UserTable(
user_id="1234", max_budget=100, user_email="1234"
)
)
mock_litellm_usertable.upsert = mock_client
mock_litellm_usertable.find_many = AsyncMock(return_value=None)
# Mock find_first for user_email validation (returns None for new users)
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
# Mock find_unique for user_id validation (returns None for new users)
mock_litellm_usertable.find_unique = AsyncMock(return_value=None)
team_mock_client = AsyncMock()
original_val = getattr(
litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable"
)
litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = team_mock_client
team_mock_client.update = AsyncMock(
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
tx_mock = AsyncMock()
tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
tx_mock.litellm_teamtable = team_mock_client
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx_mock)
tx_cm.__aexit__ = AsyncMock(return_value=None)
original_tx = litellm.proxy.proxy_server.prisma_client.tx
litellm.proxy.proxy_server.prisma_client.tx = MagicMock(
return_value=tx_cm
)
print(f"team_member_add_request={team_member_add_request}")
await team_member_add(
data=team_member_add_request,
user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin"),
)
mock_client.assert_called()
print(f"mock_client.call_args: {mock_client.call_args}")
print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs))
assert (
mock_client.call_args.kwargs["data"]["create"]["max_budget"]
== litellm.max_internal_user_budget
)
assert (
mock_client.call_args.kwargs["data"]["create"]["budget_duration"]
== litellm.internal_user_budget_duration
)
litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val
litellm.proxy.proxy_server.prisma_client.tx = original_tx
@pytest.mark.parametrize("team_member_role", ["admin", "user"])
@pytest.mark.parametrize("team_route", ["/team/member_add", "/team/member_delete"])
@pytest.mark.asyncio
async def test_create_team_member_add_team_admin_user_api_key_auth(
prisma_client, team_member_role, team_route
):
import time
from fastapi import Request
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, Member
from litellm.proxy.proxy_server import (
ProxyException,
hash_token,
user_api_key_auth,
user_api_key_cache,
)
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm, "max_internal_user_budget", 10)
setattr(litellm, "internal_user_budget_duration", "5m")
user = f"ishaan {uuid.uuid4().hex}"
_team_id = "litellm-test-client-id-new"
user_key = "sk-12345678"
valid_token = UserAPIKeyAuth(
team_id=_team_id,
token=hash_token(user_key),
team_member=Member(role=team_member_role, user_id=user),
last_refreshed_at=time.time(),
)
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
team_obj = LiteLLM_TeamTableCachedObj(
team_id=_team_id,
blocked=False,
last_refreshed_at=time.time(),
metadata={"guardrails": {"modify_guardrails": False}},
)
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
## TEST IF TEAM ADMIN ALLOWED TO CALL /MEMBER_ADD ENDPOINT
import json
from starlette.datastructures import URL
request = Request(scope={"type": "http"})
request._url = URL(url=team_route)
body = {}
json_bytes = json.dumps(body).encode("utf-8")
request._body = json_bytes
## ALLOWED BY USER_API_KEY_AUTH
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
@pytest.mark.parametrize("user_role", ["admin", "user"])
@pytest.mark.asyncio
async def test_create_team_member_add_team_admin(
prisma_client, new_member_method, user_role
):
"""
Relevant issue - https://github.com/BerriAI/litellm/issues/5300
Allow team admins to:
- Add and remove team members
- raise error if team member not an existing 'internal_user'
"""
import time
from fastapi import Request
from litellm.proxy._types import (
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
Member,
)
from litellm.proxy.proxy_server import (
HTTPException,
ProxyException,
hash_token,
user_api_key_auth,
user_api_key_cache,
)
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm, "max_internal_user_budget", 10)
setattr(litellm, "internal_user_budget_duration", "5m")
user = f"ishaan {uuid.uuid4().hex}"
_team_id = "litellm-test-client-id-new"
user_key = "sk-12345678"
team_admin = f"krrish {uuid.uuid4().hex}"
valid_token = UserAPIKeyAuth(
team_id=_team_id,
user_id=team_admin,
token=hash_token(user_key),
last_refreshed_at=time.time(),
)
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
team_obj = LiteLLM_TeamTableCachedObj(
team_id=_team_id,
blocked=False,
last_refreshed_at=time.time(),
members_with_roles=[Member(role=user_role, user_id=team_admin)],
metadata={"guardrails": {"modify_guardrails": False}},
)
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
if new_member_method == "user_id":
data = {
"team_id": _team_id,
"member": [{"role": "user", "user_id": user}],
}
elif new_member_method == "user_email":
data = {
"team_id": _team_id,
"member": [{"role": "user", "user_email": user}],
}
team_member_add_request = TeamMemberAddRequest(**data)
with (
patch(
"litellm.proxy.proxy_server.prisma_client.db.litellm_usertable",
new_callable=AsyncMock,
) as mock_litellm_usertable,
patch(
"litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache",
new=AsyncMock(return_value=team_obj),
) as mock_team_obj,
patch(
"litellm.proxy.proxy_server.prisma_client.get_data",
new=AsyncMock(return_value=[]),
) as mock_get_data,
):
mock_client = AsyncMock(
return_value=LiteLLM_UserTable(
user_id="1234", max_budget=100, user_email="1234"
)
)
mock_litellm_usertable.upsert = mock_client
mock_litellm_usertable.find_many = AsyncMock(return_value=None)
# Mock find_first for user_email validation (returns None for new users)
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
# Mock find_unique for user_id validation (returns None for new users)
mock_litellm_usertable.find_unique = AsyncMock(return_value=None)
team_mock_client = AsyncMock()
team_mock_client.update = AsyncMock(
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
tx_mock = AsyncMock()
tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
tx_mock.litellm_teamtable = team_mock_client
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx_mock)
tx_cm.__aexit__ = AsyncMock(return_value=None)
with (
patch.object(
litellm.proxy.proxy_server.prisma_client.db,
"litellm_teamtable",
team_mock_client,
),
patch.object(
litellm.proxy.proxy_server.prisma_client,
"tx",
MagicMock(return_value=tx_cm),
),
):
try:
await team_member_add(
data=team_member_add_request,
user_api_key_dict=valid_token,
)
except HTTPException as e:
if user_role == "user":
assert e.status_code == 403
return
else:
raise e
mock_client.assert_called()
assert (
mock_client.call_args.kwargs["data"]["create"]["max_budget"]
== litellm.max_internal_user_budget
)
assert (
mock_client.call_args.kwargs["data"]["create"]["budget_duration"]
== litellm.internal_user_budget_duration
)
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_user_info_team_list(prisma_client):
"""Assert user_info for admin calls team_list function"""
from litellm.proxy._types import LiteLLM_UserTable
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info
with patch(
"litellm.proxy.management_endpoints.team_endpoints.list_team",
new_callable=AsyncMock,
) as mock_client:
prisma_client.get_data = AsyncMock(
return_value=LiteLLM_UserTable(
user_role="proxy_admin",
user_id="default_user_id",
max_budget=None,
user_email="",
)
)
try:
await user_info(
request=MagicMock(),
user_id=None,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234", user_id="default_user_id"
),
)
except Exception:
pass
mock_client.assert_called()
@pytest.mark.skip(reason="Local test")
@pytest.mark.asyncio
async def test_add_callback_via_key(prisma_client):
"""
Test if callback specified in key, is used.
"""
global headers
import json
from fastapi import HTTPException, Request, Response
from starlette.datastructures import URL
from litellm.proxy.proxy_server import chat_completion
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
litellm.set_verbose = True
try:
# Your test data
test_data = {
"model": "azure/gpt-4.1-mini",
"messages": [
{"role": "user", "content": "write 1 sentence poem"},
],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
}
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
request._url = URL(url="/chat/completions")
json_bytes = json.dumps(test_data).encode("utf-8")
request._body = json_bytes
with patch.object(
litellm.litellm_core_utils.litellm_logging,
"LangFuseLogger",
new=MagicMock(),
) as mock_client:
resp = await chat_completion(
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "langfuse", # 'otel', 'langfuse', 'lunary'
"callback_type": "success", # set, if required by integration - future improvement, have logging tools work for success + failure by default
"callback_vars": {
"langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY",
"langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY",
"langfuse_host": "https://us.cloud.langfuse.com",
},
}
],
}
),
)
print(resp)
mock_client.assert_called()
mock_client.return_value.log_event.assert_called()
args, kwargs = mock_client.return_value.log_event.call_args
kwargs = kwargs["kwargs"]
assert "user_api_key_metadata" in kwargs["litellm_params"]["metadata"]
assert (
"logging"
in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"]
)
checked_keys = False
for item in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"][
"logging"
]:
for k, v in item["callback_vars"].items():
print("k={}, v={}".format(k, v))
if "key" in k:
assert "os.environ" in v
checked_keys = True
assert checked_keys
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"callback_type, expected_success_callbacks, expected_failure_callbacks",
[
("success", ["langfuse"], []),
("failure", [], ["langfuse"]),
("success_and_failure", ["langfuse"], ["langfuse"]),
],
)
async def test_add_callback_via_key_litellm_pre_call_utils(
mock_prisma_client,
callback_type,
expected_success_callbacks,
expected_failure_callbacks,
):
import json
from fastapi import HTTPException, Request, Response
from starlette.datastructures import URL
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
request._url = URL(url="/chat/completions")
test_data = {
"model": "azure/gpt-4.1-mini",
"messages": [
{"role": "user", "content": "write 1 sentence poem"},
],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
}
json_bytes = json.dumps(test_data).encode("utf-8")
request._body = json_bytes
data = {
"data": {
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
},
"request": request,
"user_api_key_dict": UserAPIKeyAuth(
token=None,
key_name=None,
key_alias=None,
spend=0.0,
max_budget=None,
expires=None,
models=[],
aliases={},
config={},
user_id=None,
team_id=None,
max_parallel_requests=None,
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "langfuse",
"callback_type": callback_type,
"callback_vars": {
"langfuse_public_key": "my-mock-public-key",
"langfuse_secret_key": "my-mock-secret-key",
"langfuse_host": "https://us.cloud.langfuse.com",
},
}
],
},
tpm_limit=None,
rpm_limit=None,
budget_duration=None,
budget_reset_at=None,
allowed_cache_controls=[],
permissions={},
model_spend={},
model_max_budget={},
soft_budget_cooldown=False,
litellm_budget_table=None,
org_id=None,
team_spend=None,
team_alias=None,
team_tpm_limit=None,
team_rpm_limit=None,
team_max_budget=None,
team_models=[],
team_blocked=False,
soft_budget=None,
team_model_aliases=None,
team_member_spend=None,
team_metadata=None,
end_user_id=None,
end_user_tpm_limit=None,
end_user_rpm_limit=None,
end_user_max_budget=None,
last_refreshed_at=None,
api_key=None,
user_role=None,
allowed_model_region=None,
parent_otel_span=None,
),
"proxy_config": proxy_config,
"general_settings": {},
"version": "0.0.0",
}
new_data = await add_litellm_data_to_request(**data)
print("NEW DATA: {}".format(new_data))
assert "langfuse_public_key" in new_data
assert new_data["langfuse_public_key"] == "my-mock-public-key"
assert "langfuse_secret_key" in new_data
assert new_data["langfuse_secret_key"] == "my-mock-secret-key"
if expected_success_callbacks:
assert "success_callback" in new_data
assert new_data["success_callback"] == expected_success_callbacks
if expected_failure_callbacks:
assert "failure_callback" in new_data
assert new_data["failure_callback"] == expected_failure_callbacks
@pytest.mark.asyncio
@pytest.mark.parametrize(
"disable_fallbacks_set",
[
True,
False,
],
)
async def test_disable_fallbacks_by_key(disable_fallbacks_set):
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
key_metadata = {"disable_fallbacks": disable_fallbacks_set}
existing_data = {
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
}
data = LiteLLMProxyRequestSetup.add_key_level_controls(
key_metadata=key_metadata,
data=existing_data,
_metadata_variable_name="metadata",
)
assert data["disable_fallbacks"] == disable_fallbacks_set
@pytest.mark.asyncio
@pytest.mark.parametrize(
"callback_type, expected_success_callbacks, expected_failure_callbacks",
[
("success", ["gcs_bucket"], []),
("failure", [], ["gcs_bucket"]),
("success_and_failure", ["gcs_bucket"], ["gcs_bucket"]),
],
)
async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket(
mock_prisma_client,
callback_type,
expected_success_callbacks,
expected_failure_callbacks,
):
import json
from fastapi import HTTPException, Request, Response
from starlette.datastructures import URL
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
request._url = URL(url="/chat/completions")
test_data = {
"model": "azure/gpt-4.1-mini",
"messages": [
{"role": "user", "content": "write 1 sentence poem"},
],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
}
json_bytes = json.dumps(test_data).encode("utf-8")
request._body = json_bytes
data = {
"data": {
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
},
"request": request,
"user_api_key_dict": UserAPIKeyAuth(
token=None,
key_name=None,
key_alias=None,
spend=0.0,
max_budget=None,
expires=None,
models=[],
aliases={},
config={},
user_id=None,
team_id=None,
max_parallel_requests=None,
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "gcs_bucket",
"callback_type": callback_type,
"callback_vars": {
"gcs_bucket_name": "key-logging-project1",
"gcs_path_service_account": "pathrise-convert-1606954137718-a956eef1a2a8.json",
},
}
],
},
tpm_limit=None,
rpm_limit=None,
budget_duration=None,
budget_reset_at=None,
allowed_cache_controls=[],
permissions={},
model_spend={},
model_max_budget={},
soft_budget_cooldown=False,
litellm_budget_table=None,
org_id=None,
team_spend=None,
team_alias=None,
team_tpm_limit=None,
team_rpm_limit=None,
team_max_budget=None,
team_models=[],
team_blocked=False,
soft_budget=None,
team_model_aliases=None,
team_member_spend=None,
team_metadata=None,
end_user_id=None,
end_user_tpm_limit=None,
end_user_rpm_limit=None,
end_user_max_budget=None,
last_refreshed_at=None,
api_key=None,
user_role=None,
allowed_model_region=None,
parent_otel_span=None,
),
"proxy_config": proxy_config,
"general_settings": {},
"version": "0.0.0",
}
new_data = await add_litellm_data_to_request(**data)
print("NEW DATA: {}".format(new_data))
assert "gcs_bucket_name" in new_data
assert new_data["gcs_bucket_name"] == "key-logging-project1"
assert "gcs_path_service_account" in new_data
assert (
new_data["gcs_path_service_account"]
== "pathrise-convert-1606954137718-a956eef1a2a8.json"
)
if expected_success_callbacks:
assert "success_callback" in new_data
assert new_data["success_callback"] == expected_success_callbacks
if expected_failure_callbacks:
assert "failure_callback" in new_data
assert new_data["failure_callback"] == expected_failure_callbacks
@pytest.mark.asyncio
@pytest.mark.parametrize(
"callback_type, expected_success_callbacks, expected_failure_callbacks",
[
("success", ["langsmith"], []),
("failure", [], ["langsmith"]),
("success_and_failure", ["langsmith"], ["langsmith"]),
],
)
async def test_add_callback_via_key_litellm_pre_call_utils_langsmith(
mock_prisma_client,
callback_type,
expected_success_callbacks,
expected_failure_callbacks,
):
import json
from fastapi import HTTPException, Request, Response
from starlette.datastructures import URL
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
request._url = URL(url="/chat/completions")
test_data = {
"model": "azure/gpt-4.1-mini",
"messages": [
{"role": "user", "content": "write 1 sentence poem"},
],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
}
json_bytes = json.dumps(test_data).encode("utf-8")
request._body = json_bytes
data = {
"data": {
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "write 1 sentence poem"}],
"max_tokens": 10,
"mock_response": "Hello world",
"api_key": "my-fake-key",
},
"request": request,
"user_api_key_dict": UserAPIKeyAuth(
token=None,
key_name=None,
key_alias=None,
spend=0.0,
max_budget=None,
expires=None,
models=[],
aliases={},
config={},
user_id=None,
team_id=None,
max_parallel_requests=None,
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "langsmith",
"callback_type": callback_type,
"callback_vars": {
"langsmith_api_key": "ls-1234",
"langsmith_project": "pr-brief-resemblance-72",
"langsmith_base_url": "https://api.smith.langchain.com",
},
}
],
},
tpm_limit=None,
rpm_limit=None,
budget_duration=None,
budget_reset_at=None,
allowed_cache_controls=[],
permissions={},
model_spend={},
model_max_budget={},
soft_budget_cooldown=False,
litellm_budget_table=None,
org_id=None,
team_spend=None,
team_alias=None,
team_tpm_limit=None,
team_rpm_limit=None,
team_max_budget=None,
team_models=[],
team_blocked=False,
soft_budget=None,
team_model_aliases=None,
team_member_spend=None,
team_metadata=None,
end_user_id=None,
end_user_tpm_limit=None,
end_user_rpm_limit=None,
end_user_max_budget=None,
last_refreshed_at=None,
api_key=None,
user_role=None,
allowed_model_region=None,
parent_otel_span=None,
),
"proxy_config": proxy_config,
"general_settings": {},
"version": "0.0.0",
}
new_data = await add_litellm_data_to_request(**data)
print("NEW DATA: {}".format(new_data))
assert "langsmith_api_key" in new_data
assert new_data["langsmith_api_key"] == "ls-1234"
assert "langsmith_project" in new_data
assert new_data["langsmith_project"] == "pr-brief-resemblance-72"
assert "langsmith_base_url" in new_data
assert new_data["langsmith_base_url"] == "https://api.smith.langchain.com"
if expected_success_callbacks:
assert "success_callback" in new_data
assert new_data["success_callback"] == expected_success_callbacks
if expected_failure_callbacks:
assert "failure_callback" in new_data
assert new_data["failure_callback"] == expected_failure_callbacks
@pytest.mark.skipif(
not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"),
reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.",
)
@pytest.mark.asyncio
async def test_gemini_pass_through_endpoint():
from starlette.datastructures import URL
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
Request,
Response,
gemini_proxy_route,
)
body = b"""
{
"contents": [{
"parts":[{
"text": "The quick brown fox jumps over the lazy dog."
}]
}]
}
"""
# Construct the scope dictionary
scope = {
"type": "http",
"method": "POST",
"path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens",
"query_string": b"key=sk-1234",
"headers": [
(b"content-type", b"application/json"),
],
}
# Create a new Request object
async def async_receive():
return {"type": "http.request", "body": body, "more_body": False}
request = Request(
scope=scope,
receive=async_receive,
)
resp = await gemini_proxy_route(
endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234",
request=request,
fastapi_response=Response(),
)
print(resp.body)
@pytest.mark.parametrize("hidden", [True, False])
@pytest.mark.asyncio
async def test_model_info_alias_without_prisma(hidden):
from litellm.proxy.proxy_server import model_info_v1
_model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo"},
}
]
model_alias = "gpt-4"
router = litellm.Router(
model_list=_model_list,
model_group_alias={
model_alias: {
"model": "gpt-3.5-turbo",
"hidden": hidden,
}
},
)
setattr(litellm.proxy.proxy_server, "llm_router", router)
setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
setattr(litellm.proxy.proxy_server, "prisma_client", None)
resp = await model_info_v1(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
models = resp["data"]
alias_found = any(
m["model_name"] == model_alias
for m in models
)
assert alias_found is (not hidden)
@pytest.mark.parametrize("hidden", [True, False])
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_proxy_model_group_alias_checks(prisma_client, hidden):
"""
Check if model group alias is returned on
`/v1/models`
`/v1/model/info`
`/v1/model_group/info`
"""
import json
from fastapi import HTTPException, Request, Response
from starlette.datastructures import URL
from litellm.proxy.proxy_server import model_group_info, model_info_v1, model_list
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
_model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo"},
}
]
model_alias = "gpt-4"
router = litellm.Router(
model_list=_model_list,
model_group_alias={model_alias: {"model": "gpt-3.5-turbo", "hidden": hidden}},
)
setattr(litellm.proxy.proxy_server, "llm_router", router)
setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
request._url = URL(url="/v1/models")
resp = await model_list(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
if hidden:
assert len(resp["data"]) == 1
else:
assert len(resp["data"]) == 2
print(resp)
resp = await model_info_v1(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
models = resp["data"]
is_model_alias_in_list = False
for item in models:
if model_alias == item["model_name"]:
is_model_alias_in_list = True
if hidden:
assert is_model_alias_in_list is False
else:
assert is_model_alias_in_list
resp = await model_group_info(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
print(f"resp: {resp}")
models = resp["data"]
is_model_alias_in_list = False
print(f"model_alias: {model_alias}, models: {models}")
for item in models:
if model_alias == item.model_group:
is_model_alias_in_list = True
if hidden:
assert is_model_alias_in_list is False
else:
assert is_model_alias_in_list, f"models: {models}"
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_proxy_model_group_info_rerank(prisma_client):
"""
Check if rerank model is returned on the following endpoints
`/v1/models`
`/v1/model/info`
`/v1/model_group/info`
"""
import json
from fastapi import HTTPException, Request, Response
from starlette.datastructures import URL
from litellm.proxy.proxy_server import model_group_info, model_info_v1, model_list
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
_model_list = [
{
"model_name": "rerank-english-v3.0",
"litellm_params": {"model": "cohere/rerank-english-v3.0"},
"model_info": {
"mode": "rerank",
},
}
]
router = litellm.Router(model_list=_model_list)
setattr(litellm.proxy.proxy_server, "llm_router", router)
setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
request = Request(scope={"type": "http", "method": "POST", "headers": {}})
request._url = URL(url="/v1/models")
resp = await model_list(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
assert len(resp["data"]) == 1
print(resp)
resp = await model_info_v1(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
models = resp["data"]
assert models[0]["model_info"]["mode"] == "rerank"
resp = await model_group_info(
user_api_key_dict=UserAPIKeyAuth(models=[]),
)
print(resp)
models = resp["data"]
assert models[0].mode == "rerank"
# @pytest.mark.asyncio
# async def test_proxy_team_member_add(prisma_client):
# """
# Add 10 people to a team. Confirm all 10 are added.
# """
# from litellm.proxy.management_endpoints.team_endpoints import (
# team_member_add,
# new_team,
# )
# from litellm.proxy._types import TeamMemberAddRequest, Member, NewTeamRequest
# setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
# setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
# try:
# async def test():
# await litellm.proxy.proxy_server.prisma_client.connect()
# from litellm.proxy.proxy_server import user_api_key_cache
# user_api_key_dict = UserAPIKeyAuth(
# user_role=LitellmUserRoles.PROXY_ADMIN,
# api_key="sk-1234",
# user_id="1234",
# )
# new_team()
# for _ in range(10):
# request = TeamMemberAddRequest(
# team_id="1234",
# member=Member(
# user_id="1234",
# user_role=LitellmUserRoles.INTERNAL_USER,
# ),
# )
# key = await team_member_add(
# request, user_api_key_dict=user_api_key_dict
# )
# print(key)
# user_id = key.user_id
# # check /user/info to verify user_role was set correctly
# new_user_info = await user_info(
# user_id=user_id, user_api_key_dict=user_api_key_dict
# )
# new_user_info = new_user_info.user_info
# print("new_user_info=", new_user_info)
# assert new_user_info["user_role"] == LitellmUserRoles.INTERNAL_USER
# assert new_user_info["user_id"] == user_id
# generated_key = key.key
# bearer_token = "Bearer " + generated_key
# assert generated_key not in user_api_key_cache.in_memory_cache.cache_dict
# value_from_prisma = await prisma_client.get_data(
# token=generated_key,
# )
# print("token from prisma", value_from_prisma)
# request = Request(
# {
# "type": "http",
# "route": api_route,
# "path": api_route.path,
# "headers": [("Authorization", bearer_token)],
# }
# )
# # use generated key to auth in
# result = await user_api_key_auth(request=request, api_key=bearer_token)
# print("result from user auth with new key", result)
# asyncio.run(test())
# except Exception as e:
# pytest.fail(f"An exception occurred - {str(e)}")
@pytest.mark.asyncio
async def test_proxy_server_prisma_setup():
from litellm.proxy.proxy_server import ProxyStartupEvent, proxy_state
from litellm.proxy.utils import ProxyLogging
from litellm.caching import DualCache
user_api_key_cache = DualCache()
with patch.object(
litellm.proxy.proxy_server, "PrismaClient", new=MagicMock()
) as mock_prisma_client:
mock_client = mock_prisma_client.return_value # This is the mocked instance
mock_client.connect = AsyncMock() # Mock the connect method
mock_client.check_view_exists = AsyncMock() # Mock the check_view_exists method
mock_client.health_check = AsyncMock() # Mock the health_check method
mock_client._set_spend_logs_row_count_in_proxy_state = (
AsyncMock()
) # Mock the _set_spend_logs_row_count_in_proxy_state method
mock_client.start_db_health_watchdog_task = AsyncMock()
# Mock the db attribute with start_token_refresh_task for RDS IAM token refresh
mock_db = MagicMock()
mock_db.start_token_refresh_task = AsyncMock()
mock_client.db = mock_db
await ProxyStartupEvent._setup_prisma_client(
database_url=os.getenv("DATABASE_URL"),
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
user_api_key_cache=user_api_key_cache,
)
# Verify our mocked methods were called
mock_client.connect.assert_called_once()
mock_client.check_view_exists.assert_called_once()
# Note: This is REALLY IMPORTANT to check that the health check is called
# This is how we ensure the DB is ready before proceeding
mock_client.health_check.assert_called_once()
# check that the spend logs row count is set in proxy state
mock_client._set_spend_logs_row_count_in_proxy_state.assert_called_once()
assert proxy_state.get_proxy_state_variable("spend_logs_row_count") is not None
@pytest.mark.asyncio
async def test_proxy_server_prisma_setup_invalid_db():
"""
PROD TEST: Test that proxy server startup fails when it's unable to connect to the database
Think 2-3 times before editing / deleting this test, it's important for PROD
"""
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
from litellm.caching import DualCache
user_api_key_cache = DualCache()
invalid_db_url = "postgresql://invalid:invalid@localhost:5432/nonexistent"
_old_db_url = os.getenv("DATABASE_URL")
os.environ["DATABASE_URL"] = invalid_db_url
with pytest.raises(Exception) as exc_info:
await ProxyStartupEvent._setup_prisma_client(
database_url=invalid_db_url,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
user_api_key_cache=user_api_key_cache,
)
print("GOT EXCEPTION=", exc_info)
assert "httpx.ConnectError" in str(exc_info.value)
# # Verify the error message indicates a database connection issue
# assert any(x in str(exc_info.value).lower() for x in ["database", "connection", "authentication"])
if _old_db_url:
os.environ["DATABASE_URL"] = _old_db_url
@pytest.mark.asyncio
async def test_get_ui_settings_spend_logs_threshold():
"""
Test that get_ui_settings correctly sets DISABLE_EXPENSIVE_DB_QUERIES based on spend_logs_row_count threshold
"""
from litellm.proxy.management_endpoints.ui_sso import get_ui_settings
from litellm.proxy.proxy_server import proxy_state
from fastapi import Request
from litellm.constants import MAX_SPENDLOG_ROWS_TO_QUERY
# Create a mock request
mock_request = Request(
scope={
"type": "http",
"headers": [],
"method": "GET",
"scheme": "http",
"server": ("testserver", 80),
"path": "/sso/get/ui_settings",
"query_string": b"",
}
)
# Test case 1: When spend_logs_row_count > MAX_SPENDLOG_ROWS_TO_QUERY
proxy_state.set_proxy_state_variable(
"spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY + 1
)
response = await get_ui_settings(mock_request)
print("response from get_ui_settings", json.dumps(response, indent=4))
assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is True
assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY + 1
# Test case 2: When spend_logs_row_count < MAX_SPENDLOG_ROWS_TO_QUERY
proxy_state.set_proxy_state_variable(
"spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY - 1
)
response = await get_ui_settings(mock_request)
print("response from get_ui_settings", json.dumps(response, indent=4))
assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is False
assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY - 1
# Test case 3: Edge case - exactly MAX_SPENDLOG_ROWS_TO_QUERY
proxy_state.set_proxy_state_variable(
"spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY
)
response = await get_ui_settings(mock_request)
print("response from get_ui_settings", json.dumps(response, indent=4))
assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is False
assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY
# Clean up
proxy_state.set_proxy_state_variable("spend_logs_row_count", 0)
@pytest.mark.asyncio
async def test_run_background_health_check_reflects_llm_model_list(monkeypatch):
"""
Test that _run_background_health_check reflects changes to llm_model_list in each health check iteration.
"""
import litellm.proxy.proxy_server as proxy_server
import copy
test_model_list_1 = [{"model_name": "model-a"}]
test_model_list_2 = [{"model_name": "model-b"}]
called_model_lists = []
async def fake_perform_health_check(model_list, details, max_concurrency=None):
called_model_lists.append(copy.deepcopy(model_list))
return (["healthy"], ["unhealthy"], {})
monkeypatch.setattr(proxy_server, "health_check_interval", 1)
monkeypatch.setattr(proxy_server, "health_check_details", None)
monkeypatch.setattr(
proxy_server, "llm_model_list", copy.deepcopy(test_model_list_1)
)
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
monkeypatch.setattr(proxy_server, "health_check_results", {})
async def fake_sleep(interval):
raise asyncio.CancelledError()
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
try:
await proxy_server._run_background_health_check()
except asyncio.CancelledError:
pass
monkeypatch.setattr(
proxy_server, "llm_model_list", copy.deepcopy(test_model_list_2)
)
try:
await proxy_server._run_background_health_check()
except asyncio.CancelledError:
pass
assert len(called_model_lists) >= 2
assert called_model_lists[0] == test_model_list_1
assert called_model_lists[1] == test_model_list_2
@pytest.mark.asyncio
async def test_background_health_check_skip_disabled_models(monkeypatch):
"""Ensure models with disable_background_health_check are skipped."""
import litellm.proxy.proxy_server as proxy_server
import copy
test_model_list = [
{"model_name": "model-a"},
{
"model_name": "model-b",
"model_info": {"disable_background_health_check": True},
},
]
called_model_lists = []
async def fake_perform_health_check(
model_list, details, max_concurrency=None, **kwargs
):
called_model_lists.append(copy.deepcopy(model_list))
return (["healthy"], [], {})
monkeypatch.setattr(proxy_server, "health_check_interval", 1)
monkeypatch.setattr(proxy_server, "health_check_details", None)
monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list))
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
monkeypatch.setattr(proxy_server, "health_check_results", {})
async def fake_sleep(interval):
raise asyncio.CancelledError()
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
try:
await proxy_server._run_background_health_check()
except asyncio.CancelledError:
pass
assert called_model_lists == [[{"model_name": "model-a"}]]
@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_legacy_three_arg_stub(
monkeypatch,
):
"""Monkeypatched perform_health_check with only base kwargs should still run."""
import litellm.proxy.proxy_server as proxy_server
async def fake_perform_health_check(model_list, details, max_concurrency=None):
return ([], [], {})
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
result = await proxy_server._run_direct_health_check_with_instrumentation(
[{"model_name": "m"}],
True,
1,
{"enabled": True, "source": "test", "cycle_id": "c1"},
)
assert result == ([], [], {})
@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_accepts_instrumentation_only(
monkeypatch,
):
"""Stub that accepts instrumentation_context but not health_check filter kwargs."""
import litellm.proxy.proxy_server as proxy_server
seen: list = []
async def fake_perform_health_check(
model_list, details, max_concurrency=None, instrumentation_context=None
):
seen.append(instrumentation_context)
return ([], [], {})
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
await proxy_server._run_direct_health_check_with_instrumentation(
[],
False,
2,
{"enabled": True, "source": "test", "cycle_id": "c2"},
)
assert len(seen) == 1
assert seen[0]["cycle_id"] == "c2"
@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_accepts_filter_only(
monkeypatch,
):
"""Stub that accepts health_check_skip_disabled_background_models but not instrumentation."""
import litellm.proxy.proxy_server as proxy_server
seen: list = []
async def fake_perform_health_check(
model_list,
details,
max_concurrency=None,
health_check_skip_disabled_background_models=False,
):
seen.append(health_check_skip_disabled_background_models)
return ([], [], {})
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
await proxy_server._run_direct_health_check_with_instrumentation(
[],
True,
None,
{"enabled": False},
)
assert len(seen) == 1
assert seen[0] is False
@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_non_kw_typeerror_reraises(
monkeypatch,
):
import litellm.proxy.proxy_server as proxy_server
async def fake_perform_health_check(**kwargs):
raise TypeError("unsupported operand type(s)")
monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
with pytest.raises(TypeError, match="unsupported operand"):
await proxy_server._run_direct_health_check_with_instrumentation(
[],
True,
1,
{},
)
def test_get_timeout_from_request():
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
headers = {
"x-litellm-timeout": "90",
}
timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
assert timeout == 90
headers = {
"x-litellm-timeout": "90.5",
}
timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
assert timeout == 90.5
@pytest.mark.parametrize(
"ui_exists, ui_has_content",
[
(True, True), # UI path exists and has content
(True, False), # UI path exists but is empty
(False, False), # UI path doesn't exist
],
)
def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content):
"""
Test the non-root Docker UI path detection logic.
Tests that when LITELLM_NON_ROOT is set to "true":
- If UI path exists and has content, it should be used
- If UI path doesn't exist or is empty, proper error logging occurs
"""
import tempfile
import shutil
from unittest.mock import MagicMock
# Create a temporary directory to act as /tmp/litellm_ui
test_ui_path = tmp_path / "litellm_ui"
if ui_exists:
test_ui_path.mkdir(parents=True, exist_ok=True)
if ui_has_content:
# Create some dummy files to simulate built UI
(test_ui_path / "index.html").write_text("<html></html>")
(test_ui_path / "app.js").write_text("console.log('test');")
# Mock the environment variable and os.path operations
monkeypatch.setenv("LITELLM_NON_ROOT", "true")
# Create a mock logger to capture log messages
mock_logger = MagicMock()
# We need to reimport or reload the relevant code section
# Since this is module-level code, we'll test the logic directly
ui_path = None
non_root_ui_path = str(test_ui_path)
# Simulate the logic from proxy_server.py lines 909-920
if os.getenv("LITELLM_NON_ROOT", "").lower() == "true":
if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path):
mock_logger.info(
f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
)
mock_logger.info(
f"UI files found: {len(os.listdir(non_root_ui_path))} items"
)
ui_path = non_root_ui_path
else:
mock_logger.error(
f"UI not found at {non_root_ui_path}. UI will not be available."
)
mock_logger.error(
f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}"
)
# Verify behavior based on test parameters
if ui_exists and ui_has_content:
# UI should be found and used
assert ui_path == non_root_ui_path
assert mock_logger.info.call_count == 2
mock_logger.info.assert_any_call(
f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
)
# Verify the second info call mentions the number of items
info_calls = [call[0][0] for call in mock_logger.info.call_args_list]
assert any("UI files found:" in call and "items" in call for call in info_calls)
assert mock_logger.error.call_count == 0
else:
# UI should not be found, error should be logged
assert ui_path is None
assert mock_logger.error.call_count == 2
mock_logger.error.assert_any_call(
f"UI not found at {non_root_ui_path}. UI will not be available."
)
# Verify the second error call has path existence info
error_calls = [call[0][0] for call in mock_logger.error.call_args_list]
assert any("Path exists:" in call for call in error_calls)
assert mock_logger.info.call_count == 0
@pytest.mark.asyncio
async def test_get_config_callbacks_with_all_types(client_no_auth):
"""
Test that /get/config/callbacks returns all three callback types:
- success_callback with type="success"
- failure_callback with type="failure"
- callbacks (success_and_failure) with type="success_and_failure"
"""
from litellm.proxy.proxy_server import ProxyConfig
# Create a mock config with all three callback types
mock_config_data = {
"litellm_settings": {
"success_callback": ["langfuse", "braintrust"],
"failure_callback": ["sentry"],
"callbacks": ["otel", "langsmith"],
},
"environment_variables": {
"LANGFUSE_PUBLIC_KEY": "test-public-key",
"LANGFUSE_SECRET_KEY": "test-secret-key",
"LANGFUSE_HOST": "https://test.langfuse.com",
"BRAINTRUST_API_KEY": "test-braintrust-key",
"OTEL_EXPORTER": "otlp",
"OTEL_ENDPOINT": "http://localhost:4317",
"LANGSMITH_API_KEY": "test-langsmith-key",
},
"general_settings": {},
}
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
with patch.object(
proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
):
response = client_no_auth.get("/get/config/callbacks")
assert response.status_code == 200
result = response.json()
# Verify response structure
assert "status" in result
assert result["status"] == "success"
assert "callbacks" in result
callbacks = result["callbacks"]
# Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure)
assert len(callbacks) == 5
# Group callbacks by type
success_callbacks = [cb for cb in callbacks if cb.get("type") == "success"]
failure_callbacks = [cb for cb in callbacks if cb.get("type") == "failure"]
success_and_failure_callbacks = [
cb for cb in callbacks if cb.get("type") == "success_and_failure"
]
# Verify all callbacks have required fields
for callback in callbacks:
assert "name" in callback
assert "variables" in callback
assert "type" in callback
assert callback["type"] in ["success", "failure", "success_and_failure"]
# Verify success callbacks
assert len(success_callbacks) == 2
success_names = [cb["name"] for cb in success_callbacks]
assert "langfuse" in success_names
assert "braintrust" in success_names
# Verify failure callbacks
assert len(failure_callbacks) == 1
assert failure_callbacks[0]["name"] == "sentry"
# Verify success_and_failure callbacks
assert len(success_and_failure_callbacks) == 2
success_and_failure_names = [cb["name"] for cb in success_and_failure_callbacks]
assert "otel" in success_and_failure_names
assert "langsmith" in success_and_failure_names
@pytest.mark.asyncio
async def test_get_config_callbacks_environment_variables(client_no_auth):
"""
Test that /get/config/callbacks correctly includes environment variables
for each callback type. Under ``client_no_auth`` the resolved role is
not ``PROXY_ADMIN``, so values matched by the redaction helper come back
as ``"REDACTED"`` and other values pass through verbatim.
"""
from litellm.proxy.proxy_server import ProxyConfig
# Create a mock config with callbacks and their env vars
mock_config_data = {
"litellm_settings": {
"success_callback": ["langfuse"],
"failure_callback": [],
"callbacks": ["otel"],
},
"environment_variables": {
"LANGFUSE_PUBLIC_KEY": "test-public-key",
"LANGFUSE_SECRET_KEY": "test-secret-key",
"LANGFUSE_HOST": "https://cloud.langfuse.com",
"OTEL_EXPORTER": "otlp",
"OTEL_ENDPOINT": "http://localhost:4317",
"OTEL_HEADERS": "key=value",
},
"general_settings": {},
}
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
with patch.object(
proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
):
response = client_no_auth.get("/get/config/callbacks")
assert response.status_code == 200
result = response.json()
callbacks = result["callbacks"]
# Find langfuse callback (success type)
langfuse_callback = next(
(cb for cb in callbacks if cb["name"] == "langfuse"), None
)
assert langfuse_callback is not None
assert langfuse_callback["type"] == "success"
assert "variables" in langfuse_callback
langfuse_vars = langfuse_callback["variables"]
assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars
assert langfuse_vars["LANGFUSE_PUBLIC_KEY"] == "REDACTED"
assert "LANGFUSE_SECRET_KEY" in langfuse_vars
assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "REDACTED"
assert "LANGFUSE_HOST" in langfuse_vars
assert langfuse_vars["LANGFUSE_HOST"] == "https://cloud.langfuse.com"
# Find otel callback (success_and_failure type)
otel_callback = next((cb for cb in callbacks if cb["name"] == "otel"), None)
assert otel_callback is not None
assert otel_callback["type"] == "success_and_failure"
assert "variables" in otel_callback
otel_vars = otel_callback["variables"]
assert "OTEL_EXPORTER" in otel_vars
assert otel_vars["OTEL_EXPORTER"] == "otlp"
assert "OTEL_ENDPOINT" in otel_vars
assert otel_vars["OTEL_ENDPOINT"] == "http://localhost:4317"
assert "OTEL_HEADERS" in otel_vars
assert otel_vars["OTEL_HEADERS"] == "REDACTED"
@pytest.mark.asyncio
async def test_update_config_success_callback_normalization():
"""
Ensure success_callback values are normalized to lowercase when updating config.
This prevents delete_callback (which searches lowercase) from failing on mixed case inputs like 'SQS'.
"""
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ConfigYAML
setattr(proxy_server, "proxy_logging_obj", MagicMock())
existing_litellm_settings = {"success_callback": ["langfuse"]}
class FakeRow:
def __init__(self, name, value):
self.param_name = name
self.param_value = value
upserted = {}
async def fake_find_first(where=None):
if where and where.get("param_name") == "litellm_settings":
return FakeRow("litellm_settings", existing_litellm_settings)
return None
async def fake_upsert(where=None, data=None):
upserted[where["param_name"]] = json.loads(data["update"]["param_value"])
class MockPrisma:
def __init__(self):
self.db = MagicMock()
self.db.litellm_config = MagicMock()
self.db.litellm_config.find_first = AsyncMock(side_effect=fake_find_first)
self.db.litellm_config.upsert = AsyncMock(side_effect=fake_upsert)
setattr(proxy_server, "prisma_client", MockPrisma())
class MockProxyConfig:
async def add_deployment(self, prisma_client=None, proxy_logging_obj=None):
return None
setattr(proxy_server, "proxy_config", MockProxyConfig())
config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]})
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test"
)
await proxy_server.update_config(config_update, user_api_key_dict=admin_user)
assert (
"litellm_settings" in upserted
), "litellm_config.upsert was not called for litellm_settings"
callbacks = upserted["litellm_settings"]["success_callback"]
# Deduped and normalized
assert "sqs" in callbacks
assert "SQS" not in callbacks
assert "sQs" not in callbacks
# Existing callback should still be present
assert "langfuse" in callbacks
@pytest.mark.parametrize(
"data",
[
{
"model": {
"model_name": "azure/gpt-4.1-mini",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"base_model": "gpt-4.1-mini"},
},
"expected": "gpt-4.1-mini",
},
{
"model": {
"model_name": "openai/gpt-4.1-mini",
"litellm_params": {"model": "openai/gpt-4.1-mini"},
},
"expected": "openai/gpt-4.1-mini",
},
{
"model": {
"model_name": "openai/gpt-4.1-mini",
"litellm_params": {"model": "openai/gpt-4.1-mini"},
"model_info": {"base_model": "gpt-4.1-mini"},
},
"expected": "gpt-4.1-mini",
},
{
"model": {
"model_name": "claude-sonnet-4-5-20250929",
"litellm_params": {"model": "anthropic/claude-sonnet-4-5@20250929"},
"model_info": {"base_model": "anthropic/claude-sonnet-4-5-20250929"},
},
"expected": "anthropic/claude-sonnet-4-5-20250929",
},
{
"model": {
"model_name": "gemini-2.5-flash-001",
"litellm_params": {"model": "gemini/gemini-2.5-flash@001"},
"model_info": {"base_model": "gemini-2.5-flash-001"},
},
"expected": "gemini-2.5-flash-001",
},
],
)
def test_get_litellm_model_info(data):
from litellm.proxy.proxy_server import get_litellm_model_info
model = data["model"]
get_info_mock = MagicMock()
with mock.patch(
"litellm.get_model_info",
new=get_info_mock,
):
get_litellm_model_info(model=model)
get_info_mock.assert_called_once_with(data["expected"])
Reference in a new issue View git blame Copy permalink