tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 14:24:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

test: test

Browse Source
This commit is contained in:
Krrish Dholakia
2026-03-28 19:17:38 -07:00
parent a41ba7bb6a
commit bc829d51f2
56 changed files with 1192 additions and 907 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py
+3 -3
View File
@@ -169,9 +169,9 @@ async def test_prometheus_metric_tracking():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"model_info": {"id": "azure-model-id"},
},
tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py
+1 -1
View File
@@ -72,7 +72,7 @@ async def test_enterprise_custom_auth_returns_string():
auth_obj = await _user_api_key_auth_builder(
request=request,
api_key="my-custom-key",
azure_api_key_header="",
AZURE_AI_API_KEY_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
tests/image_gen_tests/test_image_edits.py
+142 -111
View File
@@ -23,10 +23,11 @@ from litellm.types.utils import StandardLoggingPayload
# Configure pytest marks to avoid warnings
pytestmark = pytest.mark.asyncio
class TestCustomLogger(CustomLogger):
def __init__(self):
self.standard_logging_payload: Optional[StandardLoggingPayload] = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.standard_logging_payload = kwargs.get("standard_logging_object", None)
pass
@@ -80,12 +81,12 @@ class BaseLLMImageEditTest(ABC):
result = self.image_edit_function(**call_args)
else:
result = await self.async_image_edit_function(**call_args)
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -97,6 +98,7 @@ class BaseLLMImageEditTest(ABC):
except litellm.ContentPolicyViolationError as e:
pass
# Get the current directory of the file being run
pwd = os.path.dirname(os.path.realpath(__file__))
@@ -107,6 +109,7 @@ TEST_IMAGES = [
SINGLE_TEST_IMAGE = open(os.path.join(pwd, "ishaan_github.png"), "rb")
def get_test_images_as_bytesio():
"""Helper function to get test images as BytesIO objects"""
bytesio_images = []
@@ -129,6 +132,7 @@ class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest):
"image": TEST_IMAGES,
}
class TestOpenAIImageEditDallE2(BaseLLMImageEditTest):
"""
Concrete implementation of BaseLLMImageEditTest for OpenAI DALL-E-2 image edits.
@@ -155,7 +159,7 @@ class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest):
"model": "azure_ai/flux.2-pro",
"image": SINGLE_TEST_IMAGE,
"api_base": "https://litellm-ci-cd-prod.services.ai.azure.com",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": "preview",
}
@@ -187,7 +191,7 @@ async def test_openai_image_edit_litellm_router():
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -199,17 +203,19 @@ async def test_openai_image_edit_litellm_router():
except litellm.ContentPolicyViolationError as e:
pass
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_openai_image_edit_with_bytesio():
"""Test image editing using BytesIO objects instead of file readers"""
from litellm import image_edit, aimage_edit
litellm._turn_on_debug()
try:
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Get images as BytesIO objects
bytesio_images = get_test_images_as_bytesio()
@@ -222,7 +228,7 @@ async def test_openai_image_edit_with_bytesio():
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -239,7 +245,7 @@ async def test_openai_image_edit_with_bytesio():
async def test_azure_image_edit_litellm_sdk():
"""Test Azure image edit with mocked httpx request to validate request body and URL"""
from litellm import image_edit, aimage_edit
# Mock response for Azure image edit
mock_response = {
"created": 1589478378,
@@ -247,7 +253,7 @@ async def test_azure_image_edit_litellm_sdk():
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
],
}
class MockResponse:
@@ -267,16 +273,16 @@ async def test_azure_image_edit_litellm_sdk():
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
test_api_base = "https://ai-api-gw-uae-north.openai.azure.com"
test_api_key = "test-api-key"
test_api_version = "2025-04-01-preview"
result = await aimage_edit(
prompt=prompt,
model="azure/gpt-image-1",
@@ -285,41 +291,54 @@ async def test_azure_image_edit_litellm_sdk():
api_version=test_api_version,
image=TEST_IMAGES,
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Check the URL
call_args = mock_post.call_args
expected_url = f"{test_api_base}/openai/deployments/gpt-image-1/images/edits?api-version={test_api_version}"
actual_url = call_args.args[0] if call_args.args else call_args.kwargs.get('url')
actual_url = (
call_args.args[0] if call_args.args else call_args.kwargs.get("url")
)
print(f"Expected URL: {expected_url}")
print(f"Actual URL: {actual_url}")
assert actual_url == expected_url, f"URL mismatch. Expected: {expected_url}, Got: {actual_url}"
assert (
actual_url == expected_url
), f"URL mismatch. Expected: {expected_url}, Got: {actual_url}"
# Check the request body
if 'data' in call_args.kwargs:
if "data" in call_args.kwargs:
# For multipart form data, check the data parameter
form_data = call_args.kwargs['data']
print("Form data keys:", list(form_data.keys()) if hasattr(form_data, 'keys') else "Not a dict")
form_data = call_args.kwargs["data"]
print(
"Form data keys:",
list(form_data.keys()) if hasattr(form_data, "keys") else "Not a dict",
)
# Validate that model and prompt are in the form data
assert 'model' in form_data, "model should be in form data"
assert 'prompt' in form_data, "prompt should be in form data"
assert form_data['model'] == 'gpt-image-1', f"Expected model 'gpt-image-1', got {form_data['model']}"
assert prompt.strip() in form_data['prompt'], f"Expected prompt to contain '{prompt.strip()}'"
assert "model" in form_data, "model should be in form data"
assert "prompt" in form_data, "prompt should be in form data"
assert (
form_data["model"] == "gpt-image-1"
), f"Expected model 'gpt-image-1', got {form_data['model']}"
assert (
prompt.strip() in form_data["prompt"]
), f"Expected prompt to contain '{prompt.strip()}'"
# Check headers
headers = call_args.kwargs.get('headers', {})
headers = call_args.kwargs.get("headers", {})
print("Request headers:", headers)
assert 'Authorization' in headers, "Authorization header should be present"
assert headers['Authorization'].startswith('Bearer '), "Authorization should be Bearer token"
assert "Authorization" in headers, "Authorization header should be present"
assert headers["Authorization"].startswith(
"Bearer "
), "Authorization should be Bearer token"
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -330,15 +349,15 @@ async def test_azure_image_edit_litellm_sdk():
f.write(image_bytes)
@pytest.mark.asyncio
async def test_openai_image_edit_cost_tracking():
"""Test OpenAI image edit cost tracking with custom logger"""
from litellm import image_edit, aimage_edit
test_custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
@@ -350,12 +369,9 @@ async def test_openai_image_edit_cost_tracking():
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {
"image_tokens": 50,
"text_tokens": 50
},
"output_tokens": 1000
}
"input_tokens_details": {"image_tokens": 50, "text_tokens": 50},
"output_tokens": 1000,
},
}
class MockResponse:
@@ -375,26 +391,25 @@ async def test_openai_image_edit_cost_tracking():
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
result = await aimage_edit(
prompt=prompt,
model="openai/gpt-image-1",
image=TEST_IMAGES,
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -403,30 +418,36 @@ async def test_openai_image_edit_cost_tracking():
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
await asyncio.sleep(5)
print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str))
print(
"standard logging payload",
json.dumps(
test_custom_logger.standard_logging_payload, indent=4, default=str
),
)
# check model
assert test_custom_logger.standard_logging_payload["model"] == "gpt-image-1"
assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "openai"
assert (
test_custom_logger.standard_logging_payload["custom_llm_provider"]
== "openai"
)
# check response_cost
assert test_custom_logger.standard_logging_payload["response_cost"] is not None
assert test_custom_logger.standard_logging_payload["response_cost"] > 0
@pytest.mark.asyncio
async def test_azure_image_edit_cost_tracking():
"""Test Azure image edit cost tracking with custom logger"""
from litellm import image_edit, aimage_edit
test_custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
@@ -438,12 +459,9 @@ async def test_azure_image_edit_cost_tracking():
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {
"image_tokens": 50,
"text_tokens": 50
},
"output_tokens": 1000
}
"input_tokens_details": {"image_tokens": 50, "text_tokens": 50},
"output_tokens": 1000,
},
}
class MockResponse:
@@ -463,27 +481,26 @@ async def test_azure_image_edit_cost_tracking():
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
result = await aimage_edit(
prompt=prompt,
model="azure/CUSTOM_AZURE_DEPLOYMENT_NAME",
base_model="azure/gpt-image-1",
image=TEST_IMAGES,
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -492,14 +509,24 @@ async def test_azure_image_edit_cost_tracking():
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
await asyncio.sleep(5)
print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str))
print(
"standard logging payload",
json.dumps(
test_custom_logger.standard_logging_payload, indent=4, default=str
),
)
# check model
assert test_custom_logger.standard_logging_payload["model"] == "CUSTOM_AZURE_DEPLOYMENT_NAME"
assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "azure"
assert (
test_custom_logger.standard_logging_payload["model"]
== "CUSTOM_AZURE_DEPLOYMENT_NAME"
)
assert (
test_custom_logger.standard_logging_payload["custom_llm_provider"]
== "azure"
)
# check response_cost
assert test_custom_logger.standard_logging_payload["response_cost"] is not None
@@ -511,6 +538,7 @@ async def test_azure_image_edit_cost_tracking():
async def test_recraft_image_edit_api():
from litellm import aimage_edit
import requests
litellm._turn_on_debug()
global TEST_IMAGES
try:
@@ -526,10 +554,10 @@ async def test_recraft_image_edit_api():
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_url = result.data[0].url
# download the image
image_bytes = requests.get(image_url).content
with open("test_image_edit.png", "wb") as f:
@@ -545,51 +573,55 @@ def test_recraft_image_edit_config():
from litellm.llms.recraft.image_edit.transformation import RecraftImageEditConfig
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
config = RecraftImageEditConfig()
# Test supported OpenAI params
supported_params = config.get_supported_openai_params("recraftv3")
expected_params = ["n", "response_format", "style"]
assert supported_params == expected_params
# Test parameter mapping (reuses OpenAI logic with filtering)
image_edit_params = ImageEditOptionalRequestParams({
"n": 2,
"response_format": "b64_json",
"style": "realistic_image",
"size": "1024x1024", # Should be dropped
"quality": "high" # Should be dropped
})
mapped_params = config.map_openai_params(image_edit_params, "recraftv3", drop_params=True)
image_edit_params = ImageEditOptionalRequestParams(
{
"n": 2,
"response_format": "b64_json",
"style": "realistic_image",
"size": "1024x1024", # Should be dropped
"quality": "high", # Should be dropped
}
)
mapped_params = config.map_openai_params(
image_edit_params, "recraftv3", drop_params=True
)
# Should only contain supported params
assert mapped_params["n"] == 2
assert mapped_params["response_format"] == "b64_json"
assert mapped_params["style"] == "realistic_image"
assert "size" not in mapped_params # Should be dropped
assert "quality" not in mapped_params # Should be dropped
# Test request transformation (reuses OpenAI file handling)
mock_image = b"fake_image_data"
prompt = "winter landscape"
litellm_params = GenericLiteLLMParams(api_key="test_key")
data, files = config.transform_image_edit_request(
model="recraftv3",
prompt=prompt,
image=mock_image,
image_edit_optional_request_params={"strength": 0.7, "n": 1},
litellm_params=litellm_params,
headers={}
headers={},
)
# Check data structure (like OpenAI but with Recraft additions)
assert data["prompt"] == prompt
assert data["strength"] == 0.7 # Recraft-specific parameter
assert data["model"] == "recraftv3"
# Check file structure (reuses OpenAI logic)
assert len(files) == 1
assert files[0][0] == "image" # Field name (not image[] like OpenAI)
@@ -603,11 +635,12 @@ def test_recraft_image_edit_config():
async def test_multiple_vs_single_image_edit(sync_mode):
"""Test that both single and multiple image editing work correctly"""
from litellm import image_edit, aimage_edit
litellm._turn_on_debug()
try:
prompt = "Add a soft blue tint to the image(s)"
# Test single image
if sync_mode:
single_result = image_edit(
@@ -621,10 +654,10 @@ async def test_multiple_vs_single_image_edit(sync_mode):
model="gpt-image-1",
image=SINGLE_TEST_IMAGE,
)
print("Single image result:", single_result)
ImageResponse.model_validate(single_result)
# Test multiple images
if sync_mode:
multiple_result = image_edit(
@@ -638,10 +671,10 @@ async def test_multiple_vs_single_image_edit(sync_mode):
model="gpt-image-1",
image=TEST_IMAGES,
)
print("Multiple images result:", multiple_result)
ImageResponse.model_validate(multiple_result)
# Both should return valid responses
assert single_result is not None
assert multiple_result is not None
@@ -649,7 +682,7 @@ async def test_multiple_vs_single_image_edit(sync_mode):
assert multiple_result.data is not None
assert len(single_result.data) > 0
assert len(multiple_result.data) > 0
except litellm.ContentPolicyViolationError as e:
pytest.skip(f"Content policy violation: {e}")
@@ -659,36 +692,37 @@ async def test_multiple_vs_single_image_edit(sync_mode):
async def test_multiple_image_edit_with_different_formats():
"""Test multiple images editing with different file formats and types"""
from litellm import aimage_edit
litellm._turn_on_debug()
try:
prompt = "Create a cohesive artistic style across all images"
# Test with mixed BytesIO and file objects
mixed_images = [
SINGLE_TEST_IMAGE, # File object
get_test_images_as_bytesio()[1] # BytesIO object
get_test_images_as_bytesio()[1], # BytesIO object
]
result = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=mixed_images,
)
print("Mixed format images result:", result)
ImageResponse.model_validate(result)
assert result is not None
assert result.data is not None
assert len(result.data) > 0
# Save result if available
if result.data and result.data[0].b64_json:
image_bytes = base64.b64decode(result.data[0].b64_json)
with open("test_multiple_image_edit_mixed.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pytest.skip(f"Content policy violation: {e}")
@@ -698,7 +732,7 @@ async def test_multiple_image_edit_with_different_formats():
async def test_image_edit_array_handling():
"""Test that the image parameter correctly handles both single items and arrays"""
from litellm import aimage_edit
# Mock response
mock_response = {
"created": 1589478378,
@@ -706,7 +740,7 @@ async def test_image_edit_array_handling():
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
],
}
class MockResponse:
@@ -723,29 +757,26 @@ async def test_image_edit_array_handling():
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = MockResponse(mock_response, 200)
prompt = "Test prompt"
# Test 1: Single image (should be converted to list internally)
result1 = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=SINGLE_TEST_IMAGE,
)
# Test 2: Multiple images (already a list)
result2 = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=TEST_IMAGES,
)
# Both valid calls should succeed
ImageResponse.model_validate(result1)
ImageResponse.model_validate(result2)
# Verify that both calls were made to the API
assert mock_post.call_count == 2
tests/image_gen_tests/test_image_generation.py
+25 -16
View File
@@ -121,6 +121,7 @@ class TestVertexImageGeneration(BaseImageGenTest):
class TestVertexAIGeminiImageGeneration(BaseImageGenTest):
"""Test Gemini image generation models (Nano Banana)"""
def get_base_image_generation_call_args(self) -> dict:
# comment this when running locally
load_vertex_ai_credentials()
@@ -212,7 +213,9 @@ class TestAimlImageGeneration(BaseImageGenTest):
custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [custom_logger]
base_image_generation_call_args = self.get_base_image_generation_call_args()
base_image_generation_call_args = (
self.get_base_image_generation_call_args()
)
litellm.set_verbose = True
# Pass dummy api_key so validate_environment passes; HTTP is mocked
response = await litellm.aimage_generation(
@@ -229,7 +232,9 @@ class TestAimlImageGeneration(BaseImageGenTest):
# print("response_cost", response._hidden_params["response_cost"])
logged_standard_logging_payload = custom_logger.standard_logging_payload
print("logged_standard_logging_payload", logged_standard_logging_payload)
print(
"logged_standard_logging_payload", logged_standard_logging_payload
)
assert logged_standard_logging_payload is not None
assert logged_standard_logging_payload["response_cost"] is not None
assert logged_standard_logging_payload["response_cost"] > 0
@@ -244,7 +249,9 @@ class TestAimlImageGeneration(BaseImageGenTest):
response_dict["usage"] = dict(response_dict["usage"])
print("response usage=", response_dict.get("usage"))
assert response.data is not None # type guard for iteration (base fails here if None)
assert (
response.data is not None
) # type guard for iteration (base fails here if None)
for d in response.data:
assert isinstance(d, Image)
print("data in response.data", d)
@@ -266,25 +273,27 @@ class TestGoogleImageGen(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "gemini/imagen-4.0-generate-001"}
@pytest.mark.skip(reason="Runwayml image generation API only tested locally")
class TestRunwaymlImageGeneration(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "runwayml/gen4_image"}
class TestAzureOpenAIDalle3(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {
"model": "azure/dall-e-3",
"api_version": "2024-02-01",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"metadata": {
"model_info": {
"base_model": "azure/dall-e-3",
}
},
}
## AZURE AI DALL-E 3 is deprecated and new deployments cannot be made
# class TestAzureOpenAIDalle3(BaseImageGenTest):
# def get_base_image_generation_call_args(self) -> dict:
# return {
# "model": "azure/dall-e-3",
# "api_version": "2024-02-01",
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "metadata": {
# "model_info": {
# "base_model": "azure/dall-e-3",
# }
# },
# }
@pytest.mark.skip(reason="model EOL")
tests/litellm_utils_tests/test_health_check.py
+16 -11
View File
@@ -21,9 +21,9 @@ async def test_azure_health_check():
model_params={
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
}
)
print(f"response: {response}")
@@ -51,9 +51,9 @@ async def test_azure_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
},
input=["test for litellm"],
mode="embedding",
@@ -83,7 +83,9 @@ async def test_openai_img_gen_health_check():
# asyncio.run(test_openai_img_gen_health_check())
@pytest.mark.skip(reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)")
@pytest.mark.skip(
reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)"
)
@pytest.mark.asyncio
async def test_azure_img_gen_health_check():
"""
@@ -98,8 +100,8 @@ async def test_azure_img_gen_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/dall-e-3",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
},
mode="image_generation",
prompt="cute baby sea otter",
@@ -500,7 +502,9 @@ async def test_perform_health_check_filters_by_model_id():
async def mock_perform_health_check(m_list, details=True, **kwargs):
captured_list.append(m_list)
return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], []
return [
{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}
], []
with patch(
"litellm.proxy.health_check._perform_health_check",
@@ -657,7 +661,8 @@ async def test_health_check_creates_only_bounded_initial_tasks():
return real_create_task(coro)
with patch("litellm.ahealth_check", side_effect=mock_health_check), patch(
"litellm.proxy.health_check.asyncio.create_task", side_effect=tracked_create_task
"litellm.proxy.health_check.asyncio.create_task",
side_effect=tracked_create_task,
):
perform_task = real_create_task(
_perform_health_check(model_list, max_concurrency=2)
tests/llm_responses_api_testing/test_azure_responses_api.py
+34 -31
View File
@@ -25,11 +25,11 @@ class TestAzureResponsesAPITest(BaseResponsesAPITest):
return {
"model": "azure/gpt-4.1-mini",
"truncation": "auto",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": "2025-03-01-preview",
}
def get_advanced_model_for_shell_tool(self) -> Optional[str]:
"""If specified, overrides the model used by test_responses_api_shell_tool_streaming_sees_shell_output (e.g. openai/gpt-5.2 for shell support)."""
return "azure/gpt-5-mini"
@@ -45,8 +45,8 @@ async def test_azure_responses_api_preview_api_version():
model="azure/gpt-5-mini",
truncation="auto",
api_version="preview",
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
input="Hello, can you tell me a short joke?",
)
@@ -108,7 +108,9 @@ async def test_azure_responses_api_status_error():
"role": "assistant",
"type": "message",
"status": "completed",
"content": [{"type": "output_text", "text": "Here's an interesting fact."}],
"content": [
{"type": "output_text", "text": "Here's an interesting fact."}
],
}
],
}
@@ -124,7 +126,7 @@ async def test_azure_responses_api_status_error():
captured_request_body = json.loads(kwargs["data"])
import httpx
# Create a proper httpx Response object
response_content = json.dumps(mock_response_data).encode("utf-8")
response = httpx.Response(
@@ -149,18 +151,17 @@ async def test_azure_responses_api_status_error():
)
# Verify that 'status' field is not present in any of the input messages
print("Final request body:", json.dumps(captured_request_body, indent=4, default=str))
print(
"Final request body:", json.dumps(captured_request_body, indent=4, default=str)
)
assert "input" in captured_request_body, "Request body should contain 'input' field"
expected_input = [
{
"content": "tell me an interesting fact",
"role": "user"
},
{"content": "tell me an interesting fact", "role": "user"},
{
"id": "rs_0ab687487834d9df0068e462a1b2d88197aabbc832c9ba5316",
"summary": [],
"type": "reasoning"
"type": "reasoning",
},
{
"id": "msg_0ab687487834d9df0068e462a1df188197b74b1eef05102c18",
@@ -169,18 +170,15 @@ async def test_azure_responses_api_status_error():
"annotations": [],
"text": "very good morning",
"type": "output_text",
"logprobs": []
"logprobs": [],
}
],
"role": "assistant",
"type": "message"
"type": "message",
},
{
"role": "user",
"content": "tell me another"
}
{"role": "user", "content": "tell me another"},
]
assert captured_request_body["input"] == expected_input, (
f"Request body input should match expected format without 'status' field.\n"
f"Expected: {json.dumps(expected_input, indent=2)}\n"
@@ -193,9 +191,9 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix():
"""
Test that Azure-specific headers like 'x-request-id' and 'apim-request-id'
are properly forwarded with 'llm_provider-' prefix in response._hidden_params["headers"].
Issue: https://github.com/BerriAI/litellm/issues/16538
The fix ensures that processed headers (with llm_provider- prefix) are stored
in response._hidden_params["headers"] instead of additional_headers, making them
accessible via completion.headers in the same way as the completion API.
@@ -253,12 +251,12 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix():
# Check that the response has the expected headers structure
assert hasattr(response, "_hidden_params"), "Response should have _hidden_params"
assert "additional_headers" in response._hidden_params, (
"Response _hidden_params should contain 'additional_headers' with the LLM provider headers"
)
assert (
"additional_headers" in response._hidden_params
), "Response _hidden_params should contain 'additional_headers' with the LLM provider headers"
headers = response._hidden_params["additional_headers"]
# Verify that Azure-specific headers are present with llm_provider- prefix
assert "llm_provider-x-request-id" in headers, (
f"Response should contain 'llm_provider-x-request-id' header. "
@@ -268,12 +266,17 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix():
f"Response should contain 'llm_provider-apim-request-id' header. "
f"Headers: {list(headers.keys())}"
)
# Verify the header values match
assert headers["llm_provider-x-request-id"] == "12086715-aca3-4006-a29f-2f1e1d552043"
assert headers["llm_provider-apim-request-id"] == "25664b0d-cf4b-4e10-8d27-c7272e7efd49"
assert (
headers["llm_provider-x-request-id"] == "12086715-aca3-4006-a29f-2f1e1d552043"
)
assert (
headers["llm_provider-apim-request-id"]
== "25664b0d-cf4b-4e10-8d27-c7272e7efd49"
)
assert headers["llm_provider-x-ms-region"] == "Sweden Central"
# Also verify openai-compatible headers are included
assert "x-ratelimit-limit-tokens" in headers
assert "x-ratelimit-remaining-tokens" in headers
tests/llm_translation/test_azure_ai.py
+69 -38
View File
@@ -283,8 +283,8 @@ async def test_azure_ai_request_format():
litellm._turn_on_debug()
# Set up the test parameters
api_key = os.getenv("AZURE_API_KEY")
api_base = os.getenv("AZURE_API_BASE")
api_key = os.getenv("AZURE_AI_API_KEY")
api_base = os.getenv("AZURE_AI_API_BASE")
model = "azure_ai/gpt-4.1-mini"
messages = [
{"role": "user", "content": "hi"},
@@ -310,17 +310,17 @@ async def test_azure_gpt5_reasoning(model):
messages=[{"role": "user", "content": "What is the capital of France?"}],
reasoning_effort="minimal",
max_tokens=10,
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
)
print("response: ", response)
assert response.choices[0].message.content is not None
def test_completion_azure():
try:
from litellm import completion_cost
litellm.set_verbose = False
## Test azure call
response = completion(
@@ -331,7 +331,7 @@ def test_completion_azure():
"content": "Hello, how are you?",
}
],
api_key="os.environ/AZURE_API_KEY",
api_key="os.environ/AZURE_AI_API_KEY",
)
print(f"response: {response}")
print(f"response hidden params: {response._hidden_params}")
@@ -358,7 +358,7 @@ def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base):
response = completion(
model="azure_ai/gpt-4.1-mini",
api_base=api_base,
api_key=os.getenv("AZURE_API_KEY"),
api_key=os.getenv("AZURE_AI_API_KEY"),
messages=[{"role": "user", "content": "What is the meaning of life?"}],
)
@@ -374,13 +374,15 @@ async def test_azure_ai_model_router():
"""
Test Azure AI model router non-streaming response cost tracking.
Verifies that the flat cost of $0.14 per M input tokens is applied.
Tests the pattern: azure_ai/model_router/<deployment-name>
Where deployment-name is the Azure deployment (e.g., "azure-model-router").
The model_router prefix is stripped before sending to Azure API.
"""
from litellm.llms.azure_ai.cost_calculator import calculate_azure_model_router_flat_cost
from litellm.llms.azure_ai.cost_calculator import (
calculate_azure_model_router_flat_cost,
)
litellm._turn_on_debug()
response = await litellm.acompletion(
model="azure_ai/model_router/azure-model-router",
@@ -394,23 +396,22 @@ async def test_azure_ai_model_router():
tracked_cost = response._hidden_params["response_cost"]
assert tracked_cost > 0
print("Tracked cost: ", tracked_cost)
# Verify flat cost is included using the helper function
usage = response.usage
if usage and usage.prompt_tokens:
expected_flat_cost = calculate_azure_model_router_flat_cost(
model="model_router/azure-model-router",
prompt_tokens=usage.prompt_tokens
model="model_router/azure-model-router", prompt_tokens=usage.prompt_tokens
)
print(f"Prompt tokens: {usage.prompt_tokens}")
print(f"Expected flat cost: ${expected_flat_cost:.9f}")
print(f"Total tracked cost: ${tracked_cost:.9f}")
# Total cost should be at least the flat cost
assert tracked_cost >= expected_flat_cost, (
f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}"
)
assert (
tracked_cost >= expected_flat_cost
), f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}"
# Verify the flat cost is non-zero
assert expected_flat_cost > 0, "Flat cost should be greater than 0"
@@ -445,15 +446,20 @@ async def test_azure_ai_model_router_streaming_model_in_chunk():
# The model should NOT be azure-model-router (the request model)
# It should be the actual model from the response (e.g., gpt-4.1-nano, gpt-5-nano, etc.)
for model in chunks_with_model:
assert model != "azure-model-router", f"Chunk model should be actual model, not request model. Got: {model}"
assert (
model != "azure-model-router"
), f"Chunk model should be actual model, not request model. Got: {model}"
# The actual model should be a real model name like gpt-4.1-nano, gpt-5-nano, etc.
print(f"Verified chunk has actual model: {model}")
class AzureModelRouterStreamingCallback(litellm.integrations.custom_logger.CustomLogger):
class AzureModelRouterStreamingCallback(
litellm.integrations.custom_logger.CustomLogger
):
"""
Custom callback to capture streaming cost tracking for Azure Model Router.
"""
def __init__(self):
self.standard_logging_payload = None
self.response_cost = None
@@ -466,17 +472,21 @@ class AzureModelRouterStreamingCallback(litellm.integrations.custom_logger.Custo
self.async_success_called = True
self.standard_logging_payload = kwargs.get("standard_logging_object")
self.complete_streaming_response = kwargs.get("complete_streaming_response")
if self.standard_logging_payload:
self.response_cost = self.standard_logging_payload.get("response_cost")
print(f"standard_logging_payload model: {self.standard_logging_payload.get('model')}")
print(
f"standard_logging_payload model: {self.standard_logging_payload.get('model')}"
)
print(f"standard_logging_payload response_cost: {self.response_cost}")
if self.complete_streaming_response:
print(f"complete_streaming_response model: {self.complete_streaming_response.model}")
print(f"complete_streaming_response usage: {self.complete_streaming_response.usage}")
print(
f"complete_streaming_response model: {self.complete_streaming_response.model}"
)
print(
f"complete_streaming_response usage: {self.complete_streaming_response.usage}"
)
@pytest.mark.asyncio
@@ -504,10 +514,16 @@ async def test_azure_ai_model_router_streaming_cost_with_stream_options():
full_response = ""
chunks_with_model = []
async for chunk in response:
print(f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}")
print(
f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}"
)
if chunk.model:
chunks_with_model.append(chunk.model)
if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
if (
chunk.choices
and chunk.choices[0].delta
and chunk.choices[0].delta.content
):
full_response += chunk.choices[0].delta.content
print(f"Full streamed response: {full_response}")
@@ -515,27 +531,42 @@ async def test_azure_ai_model_router_streaming_cost_with_stream_options():
# Give async logging time to complete
import asyncio
await asyncio.sleep(1)
# Verify callback was called
assert test_callback.async_success_called is True, "async_log_success_event was not called"
assert test_callback.standard_logging_payload is not None, "standard_logging_payload is None"
assert (
test_callback.async_success_called is True
), "async_log_success_event was not called"
assert (
test_callback.standard_logging_payload is not None
), "standard_logging_payload is None"
# Check response cost
print(f"Final response_cost: {test_callback.response_cost}")
# The first chunk may have the request model (azure-model-router) because it's created
# before the API response is received. Subsequent chunks should have the actual model.
# At least some chunks should have the actual model (not azure-model-router)
actual_model_chunks = [m for m in chunks_with_model if m != "azure-model-router"]
assert len(actual_model_chunks) > 0, "No chunks had the actual model from the API response"
actual_model_chunks = [
m for m in chunks_with_model if m != "azure-model-router"
]
assert (
len(actual_model_chunks) > 0
), "No chunks had the actual model from the API response"
print(f"Chunks with actual model: {actual_model_chunks}")
# Verify response cost is tracked - this is the main goal of this test
assert test_callback.response_cost is not None, "response_cost is None with stream_options"
assert test_callback.response_cost > 0, f"response_cost should be > 0, got {test_callback.response_cost}"
print(f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}")
assert (
test_callback.response_cost is not None
), "response_cost is None with stream_options"
assert (
test_callback.response_cost > 0
), f"response_cost should be > 0, got {test_callback.response_cost}"
print(
f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}"
)
finally:
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = []
litellm.callbacks = []
tests/llm_translation/test_azure_o_series.py
+27 -13
View File
@@ -24,9 +24,9 @@ class TestAzureOpenAIO3Mini(BaseOSeriesModelsTest, BaseLLMChatTest):
litellm.in_memory_llm_clients_cache.flush_cache()
return {
"model": "azure/o3-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": "2024-12-01-preview"
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": "2024-12-01-preview",
}
def get_client(self):
@@ -187,13 +187,31 @@ async def test_azure_o1_series_response_format_extra_params():
litellm.set_verbose = True
client = AsyncAzureOpenAI(
api_key="fake-api-key",
base_url="https://openai-prod-test.openai.azure.com/openai/deployments/o1/chat/completions?api-version=2025-01-01-preview",
api_version="2025-01-01-preview"
api_key="fake-api-key",
base_url="https://openai-prod-test.openai.azure.com/openai/deployments/o1/chat/completions?api-version=2025-01-01-preview",
api_version="2025-01-01-preview",
)
tools = [{'type': 'function', 'function': {'name': 'get_current_time', 'description': 'Get the current time in a given location.', 'parameters': {'type': 'object', 'properties': {'location': {'type': 'string', 'description': 'The city name, e.g. San Francisco'}}, 'required': ['location']}}}]
response_format = {'type': 'json_object'}
tools = [
{
"type": "function",
"function": {
"name": "get_current_time",
"description": "Get the current time in a given location.",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name, e.g. San Francisco",
}
},
"required": ["location"],
},
},
}
]
response_format = {"type": "json_object"}
tool_choice = "auto"
with patch.object(
client.chat.completions.with_raw_response, "create"
@@ -208,7 +226,7 @@ async def test_azure_o1_series_response_format_extra_params():
messages=[{"role": "user", "content": "Hello! return a json object"}],
tools=tools,
response_format=response_format,
tool_choice=tool_choice
tool_choice=tool_choice,
)
except Exception as e:
print(f"Error: {e}")
@@ -220,7 +238,3 @@ async def test_azure_o1_series_response_format_extra_params():
assert request_body["tools"] == tools
assert request_body["response_format"] == response_format
assert request_body["tool_choice"] == tool_choice
tests/llm_translation/test_azure_openai.py
+39 -37
View File
@@ -208,8 +208,8 @@ class TestAzureEmbedding(BaseLLMEmbeddingTest):
def get_base_embedding_call_args(self) -> dict:
return {
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
}
def get_custom_llm_provider(self) -> litellm.LlmProviders:
@@ -618,8 +618,8 @@ def test_azure_safety_result():
response = completion(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version="2024-12-01-preview",
messages=[{"role": "user", "content": "Hello world"}],
)
@@ -671,6 +671,8 @@ def test_completion_azure_deployment_id():
)
# Add any assertions here to check the response
print(response)
def test_azure_with_content_safety_error():
"""
Verify user can access innererror from the Azure OpenAI exception
@@ -679,55 +681,55 @@ def test_azure_with_content_safety_error():
from litellm.exceptions import ContentPolicyViolationError
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from unittest.mock import MagicMock
mock_exception = Exception("The response was filtered due to the prompt triggering Azure OpenAI's content management policy")
mock_exception = Exception(
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy"
)
mock_exception.body = {
"innererror": {
"code": "ResponsibleAIPolicyViolation",
"content_filter_result": {
"hate": {
"filtered": False,
"severity": "safe"
},
"jailbreak": {
"filtered": False,
"detected": False
},
"self_harm": {
"filtered": False,
"severity": "safe"
},
"sexual": {
"filtered": False,
"severity": "safe"
},
"violence": {
"filtered": True,
"severity": "high"
}
}
"hate": {"filtered": False, "severity": "safe"},
"jailbreak": {"filtered": False, "detected": False},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": True, "severity": "high"},
},
}
}
mock_response = MagicMock()
mock_response.status_code = 400
mock_exception.response = mock_response
with pytest.raises(ContentPolicyViolationError) as exc_info:
exception_type(
model="azure/gpt-4o-new-test",
original_exception=mock_exception,
custom_llm_provider="azure"
custom_llm_provider="azure",
)
e = exc_info.value
print("got exception=", e)
assert e.provider_specific_fields is not None
print("got provider_specific_fields=", e.provider_specific_fields)
assert e.provider_specific_fields.get("innererror") is not None
assert e.provider_specific_fields["innererror"]["code"] == "ResponsibleAIPolicyViolation"
assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["filtered"] is True
assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["severity"] == "high"
assert (
e.provider_specific_fields["innererror"]["code"]
== "ResponsibleAIPolicyViolation"
)
assert (
e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][
"filtered"
]
is True
)
assert (
e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][
"severity"
]
== "high"
)
def test_azure_openai_with_prompt_cache_key():
@@ -737,9 +739,9 @@ def test_azure_openai_with_prompt_cache_key():
litellm._turn_on_debug()
response = litellm.completion(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version="2024-12-01-preview",
messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
prompt_cache_key="test_streaming_azure_openai",
)
)
tests/llm_translation/test_router_llm_translation_tests.py
+8 -8
View File
@@ -66,8 +66,8 @@ def test_router_azure_acompletion():
print("Router Test Azure - Acompletion, Acompletion with stream")
# remove api key from env to repro how proxy passes key to router
old_api_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
old_api_key = os.environ["AZURE_AI_API_KEY"]
os.environ.pop("AZURE_AI_API_KEY", None)
model_list = [
{
@@ -75,8 +75,8 @@ def test_router_azure_acompletion():
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"rpm": 1800,
},
@@ -85,8 +85,8 @@ def test_router_azure_acompletion():
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"rpm": 1800,
},
@@ -126,9 +126,9 @@ def test_router_azure_acompletion():
asyncio.run(test2())
print("\n Passed Streaming")
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
router.reset()
except Exception as e:
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
print(f"FAILED TEST")
pytest.fail(f"Got unexpected exception on router! - {e}")
tests/local_testing/example_config_yaml/azure_config.yaml
+2 -2
View File
@@ -4,12 +4,12 @@ model_list:
model: azure/gpt-4.1-mini
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
tpm: 20_000
- model_name: gpt-4-team2
litellm_params:
model: azure/gpt-4
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
api_base: https://openai-gpt-4-test-v-2.openai.azure.com/
tpm: 100_000
tests/local_testing/test_acooldowns_router.py
+1 -1
View File
@@ -31,7 +31,7 @@ def _make_model_list():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
tests/local_testing/test_assistants.py
+3 -4
View File
@@ -37,10 +37,11 @@ V0 Scope:
- Run Thread -> `/v1/threads/{thread_id}/run`
"""
def _add_azure_related_dynamic_params(data: dict) -> dict:
data["api_version"] = "2024-02-15-preview"
data["api_base"] = os.getenv("AZURE_API_BASE")
data["api_key"] = os.getenv("AZURE_API_KEY")
data["api_base"] = os.getenv("AZURE_AI_API_BASE")
data["api_key"] = os.getenv("AZURE_AI_API_KEY")
return data
@@ -236,8 +237,6 @@ async def test_aarun_thread_litellm(sync_mode, provider, is_streaming):
"""
import openai
try:
get_assistants_data = {
"custom_llm_provider": provider,
tests/local_testing/test_azure_openai.py
+6 -4
View File
@@ -40,11 +40,13 @@ async def test_aaaaazure_tenant_id_auth(respx_mock: MockRouter):
PROD Test
"""
litellm.disable_aiohttp_transport = True # since this uses respx, we need to set use_aiohttp_transport to False
litellm.disable_aiohttp_transport = (
True # since this uses respx, we need to set use_aiohttp_transport to False
)
# Clear the HTTP client cache to ensure respx mocking works
# This is critical because respx only intercepts clients created AFTER mocking is active
if hasattr(litellm, 'in_memory_llm_clients_cache'):
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
router = Router(
@@ -53,7 +55,7 @@ async def test_aaaaazure_tenant_id_auth(respx_mock: MockRouter):
"model_name": "gpt-3.5-turbo",
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"tenant_id": os.getenv("AZURE_TENANT_ID"),
"client_id": os.getenv("AZURE_CLIENT_ID"),
"client_secret": os.getenv("AZURE_CLIENT_SECRET"),
tests/local_testing/test_azure_perf.py
+4 -4
View File
@@ -9,8 +9,8 @@
# from openai import AsyncAzureOpenAI
# client = AsyncAzureOpenAI(
# api_key=os.getenv("AZURE_API_KEY"),
# azure_endpoint=os.getenv("AZURE_API_BASE"), # type: ignore
# api_key=os.getenv("AZURE_AI_API_KEY"),
# azure_endpoint=os.getenv("AZURE_AI_API_BASE"), # type: ignore
# api_version=os.getenv("AZURE_API_VERSION"),
# )
@@ -19,8 +19,8 @@
# "model_name": "azure-test",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# }
tests/local_testing/test_caching.py
+116 -37
View File
@@ -147,7 +147,12 @@ def test_caching_dynamic_args(): # test in memory cache
port=_redis_port_env,
password=_redis_password_env,
)
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
@@ -173,7 +178,12 @@ def test_caching_v2(): # test in memory cache
try:
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
@@ -200,9 +210,18 @@ def test_caching_with_ttl():
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, ttl=0, mock_response="Hello world from cache test 1"
model="gpt-3.5-turbo",
messages=messages,
caching=True,
ttl=0,
mock_response="Hello world from cache test 1",
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test 2",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test 2")
print(f"response1: {response1}")
print(f"response2: {response2}")
litellm.cache = None # disable cache
@@ -221,8 +240,18 @@ def test_caching_with_default_ttl():
try:
litellm.set_verbose = True
litellm.cache = Cache(ttl=0)
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
litellm.cache = None # disable cache
@@ -247,10 +276,16 @@ async def test_caching_with_cache_controls(sync_flag):
if sync_flag:
## TTL = 0
response1 = completion(
model="gpt-3.5-turbo", messages=messages, cache={"ttl": 0}, mock_response="Hello world"
model="gpt-3.5-turbo",
messages=messages,
cache={"ttl": 0},
mock_response="Hello world",
)
response2 = completion(
model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 10}, mock_response="Hello world"
model="gpt-3.5-turbo",
messages=messages,
cache={"s-maxage": 10},
mock_response="Hello world",
)
assert response2["id"] != response1["id"]
@@ -322,9 +357,19 @@ def test_caching_with_models_v2():
litellm.cache = Cache()
print("test2 for caching")
litellm.set_verbose = True
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response3 = completion(model="gpt-4.1-nano", messages=messages, caching=True, mock_response="Different model response")
response3 = completion(
model="gpt-4.1-nano",
messages=messages,
caching=True,
mock_response="Different model response",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
@@ -423,7 +468,10 @@ def test_embedding_caching():
text_to_embed = [embedding_large_text]
start_time = time.time()
embedding1 = embedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed,
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
end_time = time.time()
print(f"Embedding 1 response time: {end_time - start_time} seconds")
@@ -459,12 +507,18 @@ async def test_embedding_caching_individual_items_and_then_list():
"world",
]
embedding1 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed[0], caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed[0],
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
initial_prompt_tokens = embedding1.usage.prompt_tokens
await asyncio.sleep(1)
embedding2 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed[1], caching=True, mock_response="0.6,0.7,0.8,0.9,1.0"
model="text-embedding-ada-002",
input=text_to_embed[1],
caching=True,
mock_response="0.6,0.7,0.8,0.9,1.0",
)
await asyncio.sleep(1)
embedding3 = await aembedding(
@@ -480,7 +534,10 @@ async def test_embedding_caching_individual_items_and_then_list():
additional_text = "this is a new text"
text_to_embed.append(additional_text)
embedding4 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed,
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
assert embedding4.usage.prompt_tokens > embedding3.usage.prompt_tokens
@@ -490,7 +547,10 @@ async def test_embedding_caching_individual_items():
litellm.cache = Cache()
text_to_embed = "hello"
embedding1 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed,
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
await asyncio.sleep(1)
@@ -512,13 +572,13 @@ def test_embedding_caching_azure():
litellm.cache = Cache()
text_to_embed = [embedding_large_text]
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_API_BASE"]
api_key = os.environ["AZURE_AI_API_KEY"]
api_base = os.environ["AZURE_AI_API_BASE"]
api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
start_time = time.time()
print("AZURE CONFIGS")
@@ -560,8 +620,8 @@ def test_embedding_caching_azure():
pytest.fail("Error occurred: Embedding caching failed")
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_API_KEY"] = api_key
os.environ["AZURE_AI_API_BASE"] = api_base
os.environ["AZURE_AI_API_KEY"] = api_key
# test_embedding_caching_azure()
@@ -851,9 +911,18 @@ def test_redis_cache_completion():
model="gpt-3.5-turbo", messages=messages, caching=True, max_tokens=20
)
response3 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, temperature=0.5, mock_response="Different params response"
model="gpt-3.5-turbo",
messages=messages,
caching=True,
temperature=0.5,
mock_response="Different params response",
)
response4 = completion(
model="gpt-4o-mini",
messages=messages,
caching=True,
mock_response="Different model response",
)
response4 = completion(model="gpt-4o-mini", messages=messages, caching=True, mock_response="Different model response")
print("\nresponse 1", response1)
print("\nresponse 2", response2)
@@ -1127,7 +1196,11 @@ async def test_redis_cache_atext_completion():
print("test for caching, atext_completion")
response1 = await litellm.atext_completion(
model="gpt-3.5-turbo-instruct", prompt=prompt, max_tokens=40, temperature=1, mock_response="Hello world from cache test"
model="gpt-3.5-turbo-instruct",
prompt=prompt,
max_tokens=40,
temperature=1,
mock_response="Hello world from cache test",
)
await asyncio.sleep(0.5)
@@ -1458,11 +1531,17 @@ def test_cache_override():
# test embedding
response1 = embedding(
model="text-embedding-ada-002", input=["hello who are you"], caching=False, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=["hello who are you"],
caching=False,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
response2 = embedding(
model="text-embedding-ada-002", input=["hello who are you"], caching=False, mock_response="0.6,0.7,0.8,0.9,1.0"
model="text-embedding-ada-002",
input=["hello who are you"],
caching=False,
mock_response="0.6,0.7,0.8,0.9,1.0",
)
# When caching=False, responses should have different IDs
@@ -2787,7 +2866,7 @@ def test_caching_thinking_args_hit(): # test in memory cache
async def test_cache_key_in_hidden_params_acompletion():
"""
Test that cache_key is present in _hidden_params on cache hits for acompletion.
Validates fix for missing x-litellm-cache-key header on proxy cache hits.
"""
litellm.cache = Cache(
@@ -2796,10 +2875,10 @@ async def test_cache_key_in_hidden_params_acompletion():
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
unique_content = f"test cache key hidden params {uuid.uuid4()}"
messages = [{"role": "user", "content": unique_content}]
# First call - cache miss
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
@@ -2807,12 +2886,12 @@ async def test_cache_key_in_hidden_params_acompletion():
mock_response="test response",
caching=True,
)
print(f"Response 1 _hidden_params: {response1._hidden_params}")
assert response1._hidden_params.get("cache_hit") is not True
await asyncio.sleep(0.5)
# Second call - cache hit
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
@@ -2820,17 +2899,17 @@ async def test_cache_key_in_hidden_params_acompletion():
mock_response="test response",
caching=True,
)
print(f"Response 2 _hidden_params: {response2._hidden_params}")
# Verify cache hit occurred
assert response2._hidden_params.get("cache_hit") is True
# Verify cache_key is present in _hidden_params
assert "cache_key" in response2._hidden_params
assert response2._hidden_params["cache_key"] is not None
# Verify both responses have same ID (cache hit)
assert response1.id == response2.id
litellm.cache = None
tests/local_testing/test_caching_ssl.py
+2 -2
View File
@@ -59,9 +59,9 @@ def test_caching_router():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
tests/local_testing/test_class.py
+4 -4
View File
@@ -56,9 +56,9 @@
# # "model_name": "gpt-3.5-turbo", # openai model name
# # "litellm_params": { # params for litellm completion/embedding call
# # "model": "azure/gpt-4.1-mini",
# # "api_key": os.getenv("AZURE_API_KEY"),
# # "api_key": os.getenv("AZURE_AI_API_KEY"),
# # "api_version": os.getenv("AZURE_API_VERSION"),
# # "api_base": os.getenv("AZURE_API_BASE"),
# # "api_base": os.getenv("AZURE_AI_API_BASE"),
# # },
# # }
# # ]
@@ -94,9 +94,9 @@
# # "model_name": "gpt-3.5-turbo", # openai model name
# # "litellm_params": { # params for litellm completion/embedding call
# # "model": "azure/gpt-4.1-mini",
# # "api_key": os.getenv("AZURE_API_KEY"),
# # "api_key": os.getenv("AZURE_AI_API_KEY"),
# # "api_version": os.getenv("AZURE_API_VERSION"),
# # "api_base": os.getenv("AZURE_API_BASE"),
# # "api_base": os.getenv("AZURE_AI_API_BASE"),
# # },
# # }
# # ],
tests/local_testing/test_completion.py
+35 -93
View File
@@ -132,7 +132,6 @@ def test_null_role_response():
assert response.choices[0].message.role == "assistant"
def predibase_mock_post(url, data=None, json=None, headers=None, timeout=None):
mock_response = MagicMock()
mock_response.status_code = 200
@@ -286,7 +285,9 @@ def test_completion_claude_3_empty_response():
},
]
try:
response = litellm.completion(model="claude-sonnet-4-5-20250929", messages=messages)
response = litellm.completion(
model="claude-sonnet-4-5-20250929", messages=messages
)
print(response)
except litellm.InternalServerError as e:
pytest.skip(f"InternalServerError - {str(e)}")
@@ -849,8 +850,8 @@ def test_completion_mistral_azure():
litellm.set_verbose = True
response = completion(
model="mistral/Mistral-large-nmefg",
api_key=os.environ["MISTRAL_AZURE_API_KEY"],
api_base=os.environ["MISTRAL_AZURE_API_BASE"],
api_key=os.environ["MISTRAL_AZURE_AI_API_KEY"],
api_base=os.environ["MISTRAL_AZURE_AI_API_BASE"],
max_tokens=5,
messages=[
{
@@ -996,59 +997,6 @@ def test_completion_gpt4_vision():
pytest.fail(f"Error occurred: {e}")
# test_completion_gpt4_vision()
def test_completion_azure_gpt4_vision():
# azure/gpt-4, vision takes 5-seconds to respond
try:
litellm.set_verbose = True
response = completion(
model="azure/gpt-4-vision",
timeout=5,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "https://avatars.githubusercontent.com/u/29436595?v=4"
},
},
],
}
],
base_url="https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions",
api_key=os.getenv("AZURE_VISION_API_KEY"),
enhancements={"ocr": {"enabled": True}, "grounding": {"enabled": True}},
dataSources=[
{
"type": "AzureComputerVision",
"parameters": {
"endpoint": "https://gpt-4-vision-enhancement.cognitiveservices.azure.com/",
"key": os.environ["AZURE_VISION_ENHANCE_KEY"],
},
}
],
)
print(response)
except openai.APIError as e:
pass
except openai.APITimeoutError:
print("got a timeout error")
pass
except openai.RateLimitError as e:
print("got a rate liimt error", e)
pass
except openai.APIStatusError as e:
print("got an api status error", e)
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_azure_gpt4_vision()
@@ -1751,7 +1699,6 @@ def test_completion_openai_pydantic(model, api_version):
pytest.fail(f"Error occurred: {e}")
def test_completion_text_openai():
try:
# litellm.set_verbose =True
@@ -2341,9 +2288,9 @@ def test_completion_azure_extra_headers():
response = completion(
model="azure/gpt-4.1-mini",
messages=messages,
api_base=os.getenv("AZURE_API_BASE"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version="2023-07-01-preview",
api_key=os.getenv("AZURE_API_KEY"),
api_key=os.getenv("AZURE_AI_API_KEY"),
extra_headers={
"Authorization": "my-bad-key",
"Ocp-Apim-Subscription-Key": "hello-world-testing",
@@ -2379,8 +2326,8 @@ def test_completion_azure_ad_token():
litellm.set_verbose = True
old_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
old_key = os.environ["AZURE_AI_API_KEY"]
os.environ.pop("AZURE_AI_API_KEY", None)
http_client = Client()
@@ -2396,7 +2343,7 @@ def test_completion_azure_ad_token():
except Exception as e:
pass
finally:
os.environ["AZURE_API_KEY"] = old_key
os.environ["AZURE_AI_API_KEY"] = old_key
mock_client.assert_called_once()
request = mock_client.call_args[0][0]
@@ -2412,8 +2359,8 @@ def test_completion_azure_key_completion_arg():
# DO NOT REMOVE THIS TEST. No MATTER WHAT Happens!
# If you want to remove it, speak to Ishaan!
# Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this
old_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
old_key = os.environ["AZURE_AI_API_KEY"]
os.environ.pop("AZURE_AI_API_KEY", None)
try:
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
@@ -2430,9 +2377,9 @@ def test_completion_azure_key_completion_arg():
print("Hidden Params", response._hidden_params)
assert response._hidden_params["custom_llm_provider"] == "azure"
os.environ["AZURE_API_KEY"] = old_key
os.environ["AZURE_AI_API_KEY"] = old_key
except Exception as e:
os.environ["AZURE_API_KEY"] = old_key
os.environ["AZURE_AI_API_KEY"] = old_key
pytest.fail(f"Error occurred: {e}")
@@ -2443,8 +2390,8 @@ async def test_re_use_azure_async_client():
import openai
client = openai.AsyncAzureOpenAI(
azure_endpoint=os.environ["AZURE_API_BASE"],
api_key=os.environ["AZURE_API_KEY"],
azure_endpoint=os.environ["AZURE_AI_API_BASE"],
api_key=os.environ["AZURE_AI_API_KEY"],
api_version="2023-07-01-preview",
)
## Test azure call
@@ -2525,13 +2472,13 @@ def test_completion_azure2():
try:
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = False
api_base = os.environ["AZURE_API_BASE"]
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_AI_API_BASE"]
api_key = os.environ["AZURE_AI_API_KEY"]
api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
## Test azure call
response = completion(
@@ -2546,9 +2493,9 @@ def test_completion_azure2():
# Add any assertions here to check the response
print(response)
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_AI_API_BASE"] = api_base
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_KEY"] = api_key
os.environ["AZURE_AI_API_KEY"] = api_key
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -2562,13 +2509,13 @@ def test_completion_azure3():
try:
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
litellm.api_base = os.environ["AZURE_API_BASE"]
litellm.api_key = os.environ["AZURE_API_KEY"]
litellm.api_base = os.environ["AZURE_AI_API_BASE"]
litellm.api_key = os.environ["AZURE_AI_API_KEY"]
litellm.api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
## Test azure call
response = completion(
@@ -2580,9 +2527,9 @@ def test_completion_azure3():
# Add any assertions here to check the response
print(response)
os.environ["AZURE_API_BASE"] = litellm.api_base
os.environ["AZURE_AI_API_BASE"] = litellm.api_base
os.environ["AZURE_API_VERSION"] = litellm.api_version
os.environ["AZURE_API_KEY"] = litellm.api_key
os.environ["AZURE_AI_API_KEY"] = litellm.api_key
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -2594,7 +2541,7 @@ def test_completion_azure3():
# new azure test for using litellm. vars,
# use the following vars in this test and make an azure_api_call
# litellm.api_type = self.azure_api_type
# litellm.api_base = self.azure_api_base
# litellm.api_base = self.AZURE_AI_API_BASE
# litellm.api_version = self.azure_api_version
# litellm.api_key = self.api_key
def test_completion_azure_with_litellm_key():
@@ -2604,14 +2551,14 @@ def test_completion_azure_with_litellm_key():
#### set litellm vars
litellm.api_type = "azure"
litellm.api_base = os.environ["AZURE_API_BASE"]
litellm.api_base = os.environ["AZURE_AI_API_BASE"]
litellm.api_version = os.environ["AZURE_API_VERSION"]
litellm.api_key = os.environ["AZURE_API_KEY"]
litellm.api_key = os.environ["AZURE_AI_API_KEY"]
######### UNSET ENV VARs for this ################
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
######### UNSET OpenAI vars for this ##############
openai.api_type = ""
@@ -2627,9 +2574,9 @@ def test_completion_azure_with_litellm_key():
print(response)
######### RESET ENV VARs for this ################
os.environ["AZURE_API_BASE"] = litellm.api_base
os.environ["AZURE_AI_API_BASE"] = litellm.api_base
os.environ["AZURE_API_VERSION"] = litellm.api_version
os.environ["AZURE_API_KEY"] = litellm.api_key
os.environ["AZURE_AI_API_KEY"] = litellm.api_key
######### UNSET litellm vars
litellm.api_type = None
@@ -3081,7 +3028,6 @@ async def test_completion_bedrock_httpx_models(sync_mode, model):
pytest.fail(f"An error occurred - {str(e)}")
# test_completion_bedrock_titan()
@@ -3256,7 +3202,6 @@ def test_completion_anyscale_api():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="anyscale stopped serving public api endpoints")
def test_completion_anyscale_2():
try:
@@ -3871,9 +3816,6 @@ async def test_dynamic_azure_params(stream, sync_mode):
raise e
@pytest.mark.parametrize(
"model",
["gpt-4o", "azure/gpt-4.1-mini"],
tests/local_testing/test_config.py
+10 -11
View File
@@ -47,8 +47,8 @@ async def test_delete_deployment():
litellm_params = LiteLLM_Params(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
encrypted_litellm_params = litellm_params.dict(exclude_none=True)
@@ -131,8 +131,8 @@ async def test_add_existing_deployment():
litellm_params = LiteLLM_Params(
model="gpt-3.5-turbo",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params)
@@ -186,8 +186,8 @@ async def test_db_error_new_model_check():
litellm_params = LiteLLM_Params(
model="gpt-3.5-turbo",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params)
@@ -233,8 +233,8 @@ async def test_db_error_new_model_check():
litellm_params = LiteLLM_Params(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
@@ -251,8 +251,8 @@ def _create_model_list(flag_value: Literal[0, 1], master_key: str):
new_litellm_params = LiteLLM_Params(
model="azure/gpt-4.1-mini-3",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
@@ -421,4 +421,3 @@ def test_litellm_proxy_responses_api_config():
assert (
config.custom_llm_provider == LlmProviders.LITELLM_PROXY
), "custom_llm_provider should be LITELLM_PROXY"
tests/local_testing/test_configs/test_bad_config.yaml
+4 -4
View File
@@ -6,16 +6,16 @@ model_list:
- model_name: working-azure-gpt-3.5-turbo
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
- model_name: azure-gpt-3.5-turbo
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_base: os.environ/AZURE_AI_API_BASE
api_key: bad-key
- model_name: azure-embedding
litellm_params:
model: azure/text-embedding-ada-002
api_base: os.environ/AZURE_API_BASE
api_base: os.environ/AZURE_AI_API_BASE
api_key: bad-key
tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml
+1 -1
View File
@@ -3,7 +3,7 @@ model_list:
litellm_params:
model: azure/gpt-4.1-mini
api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
litellm_settings:
tests/local_testing/test_configs/test_config_no_auth.yaml
+7 -7
View File
@@ -11,7 +11,7 @@ model_list:
model_name: azure-model
- litellm_params:
api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
model: azure/gpt-4.1-mini
model_name: azure-cloudflare-model
- litellm_params:
@@ -49,8 +49,8 @@ model_list:
id: 79fc75bf-8e1b-47d5-8d24-9365a854af03
model_name: test_openai_models
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
model: azure/text-embedding-ada-002
model_info:
@@ -94,16 +94,16 @@ model_list:
mode: image_generation
model_name: dall-e-3
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-06-01-preview
model: azure/
model_info:
mode: image_generation
model_name: dall-e-2
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
model: azure/text-embedding-ada-002
model_info:
tests/local_testing/test_configs/test_custom_logger.yaml
+4 -4
View File
@@ -2,8 +2,8 @@ model_list:
- model_name: Azure OpenAI GPT-4 Canada
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: "2023-07-01-preview"
model_info:
mode: chat
@@ -12,8 +12,8 @@ model_list:
- model_name: azure-embedding-model
litellm_params:
model: azure/text-embedding-ada-002
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: "2023-07-01-preview"
model_info:
mode: embedding
tests/local_testing/test_embedding.py
+6 -6
View File
@@ -296,13 +296,13 @@ def test_openai_embedding_timeouts():
def test_openai_azure_embedding():
try:
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_API_BASE"]
api_key = os.environ["AZURE_AI_API_KEY"]
api_base = os.environ["AZURE_AI_API_BASE"]
api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
response = embedding(
model="azure/text-embedding-ada-002",
@@ -314,8 +314,8 @@ def test_openai_azure_embedding():
print(response)
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_API_KEY"] = api_key
os.environ["AZURE_AI_API_BASE"] = api_base
os.environ["AZURE_AI_API_KEY"] = api_key
except Exception as e:
pytest.fail(f"Error occurred: {e}")
tests/local_testing/test_exceptions.py
+58 -44
View File
@@ -162,8 +162,8 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
temporary_secret_key = os.environ["AWS_SECRET_ACCESS_KEY"]
os.environ["AWS_SECRET_ACCESS_KEY"] = "bad-key"
elif model == "azure/gpt-4.1-mini":
temporary_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = "bad-key"
temporary_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = "bad-key"
elif model == "claude-3-5-haiku-20241022":
temporary_key = os.environ["ANTHROPIC_API_KEY"]
os.environ["ANTHROPIC_API_KEY"] = "bad-key"
@@ -175,9 +175,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
os.environ["AI21_API_KEY"] = "bad-key"
elif "togethercomputer" in model:
temporary_key = os.environ["TOGETHERAI_API_KEY"]
os.environ["TOGETHERAI_API_KEY"] = (
"sk-test-togetherai-key-808"
)
os.environ["TOGETHERAI_API_KEY"] = "sk-test-togetherai-key-808"
elif model in litellm.openrouter_models:
temporary_key = os.environ["OPENROUTER_API_KEY"]
os.environ["OPENROUTER_API_KEY"] = "bad-key"
@@ -212,7 +210,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
if model == "gpt-3.5-turbo":
os.environ["OPENAI_API_KEY"] = temporary_key
elif model == "chatgpt-test":
os.environ["AZURE_API_KEY"] = temporary_key
os.environ["AZURE_AI_API_KEY"] = temporary_key
azure = True
elif model == "claude-3-5-haiku-20241022":
os.environ["ANTHROPIC_API_KEY"] = temporary_key
@@ -259,17 +257,17 @@ def test_completion_azure_exception():
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
old_azure_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = "good morning"
old_azure_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = "good morning"
response = completion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "hello"}],
)
os.environ["AZURE_API_KEY"] = old_azure_key
os.environ["AZURE_AI_API_KEY"] = old_azure_key
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
os.environ["AZURE_API_KEY"] = old_azure_key
os.environ["AZURE_AI_API_KEY"] = old_azure_key
print("good job got the correct error for azure when key not set")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -303,8 +301,8 @@ async def asynctest_completion_azure_exception():
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
old_azure_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = "good morning"
old_azure_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = "good morning"
response = await litellm.acompletion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "hello"}],
@@ -312,7 +310,7 @@ async def asynctest_completion_azure_exception():
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
os.environ["AZURE_API_KEY"] = old_azure_key
os.environ["AZURE_AI_API_KEY"] = old_azure_key
print("good job got the correct error for azure when key not set")
print(e)
except Exception as e:
@@ -495,6 +493,7 @@ def test_completion_bedrock_invalid_role_exception():
== "litellm.BadRequestError: Invalid Message passed in {'role': 'very-bad-role', 'content': 'hello'}"
)
@pytest.mark.skip(reason="OpenAI exception changed to a generic error")
def test_content_policy_exceptionimage_generation_openai():
try:
@@ -773,7 +772,15 @@ def test_litellm_predibase_exception():
@pytest.mark.parametrize(
"provider", ["predibase", "vertex_ai_beta", "anthropic", "databricks", "watsonx", "fireworks_ai"]
"provider",
[
"predibase",
"vertex_ai_beta",
"anthropic",
"databricks",
"watsonx",
"fireworks_ai",
],
)
def test_exception_mapping(provider):
"""
@@ -826,14 +833,14 @@ def test_fireworks_ai_exception_mapping():
2. Text-based rate limit detection (the main issue fixed)
3. Generic 400 errors that should NOT be rate limits
4. ExceptionCheckers utility function
Related to: https://github.com/BerriAI/litellm/pull/11455
Based on Fireworks AI documentation: https://docs.fireworks.ai/tools-sdks/python-client/api-reference
"""
import litellm
from litellm.llms.fireworks_ai.common_utils import FireworksAIException
from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers
# Test scenarios covering all important cases
test_scenarios = [
{
@@ -855,57 +862,63 @@ def test_fireworks_ai_exception_mapping():
"expected_exception": litellm.BadRequestError,
},
]
# Test each scenario
for scenario in test_scenarios:
mock_exception = FireworksAIException(
status_code=scenario["status_code"],
message=scenario["message"],
headers={}
status_code=scenario["status_code"], message=scenario["message"], headers={}
)
try:
response = litellm.completion(
model="fireworks_ai/llama-v3p1-70b-instruct",
messages=[{"role": "user", "content": "Hello"}],
mock_response=mock_exception,
)
pytest.fail(f"Expected {scenario['expected_exception'].__name__} to be raised")
pytest.fail(
f"Expected {scenario['expected_exception'].__name__} to be raised"
)
except scenario["expected_exception"] as e:
if scenario["expected_exception"] == litellm.RateLimitError:
assert "rate limit" in str(e).lower() or "429" in str(e)
except Exception as e:
pytest.fail(f"Expected {scenario['expected_exception'].__name__} but got {type(e).__name__}: {e}")
pytest.fail(
f"Expected {scenario['expected_exception'].__name__} but got {type(e).__name__}: {e}"
)
# Test ExceptionCheckers.is_error_str_rate_limit() method directly
# Test cases that should return True (rate limit detected)
rate_limit_strings = [
"429 rate limit exceeded",
"Rate limit exceeded, please try again later",
"Rate limit exceeded, please try again later",
"RATE LIMIT ERROR",
"Error 429: rate limit",
'{"error":{"type":"invalid_request_error","message":"rate limit exceeded, please try again later"}}',
"HTTP 429 Too Many Requests",
]
for error_str in rate_limit_strings:
assert ExceptionCheckers.is_error_str_rate_limit(error_str), f"Should detect rate limit in: {error_str}"
assert ExceptionCheckers.is_error_str_rate_limit(
error_str
), f"Should detect rate limit in: {error_str}"
# Test cases that should return False (not rate limit)
non_rate_limit_strings = [
"400 Bad Request",
"Authentication failed",
"Authentication failed",
"Invalid model specified",
"Context window exceeded",
"Internal server error",
"",
"Some other error message",
]
for error_str in non_rate_limit_strings:
assert not ExceptionCheckers.is_error_str_rate_limit(error_str), f"Should NOT detect rate limit in: {error_str}"
assert not ExceptionCheckers.is_error_str_rate_limit(
error_str
), f"Should NOT detect rate limit in: {error_str}"
# Test edge cases
assert not ExceptionCheckers.is_error_str_rate_limit(None) # type: ignore
assert not ExceptionCheckers.is_error_str_rate_limit(42) # type: ignore
@@ -1142,6 +1155,7 @@ def test_openai_gateway_timeout_error():
"""
openai_client = OpenAI()
mapped_target = openai_client.chat.completions.with_raw_response # type: ignore
def _return_exception(*args, **kwargs):
import datetime
@@ -1175,13 +1189,17 @@ def test_openai_gateway_timeout_error():
setattr(exception, k, v)
raise exception
try:
try:
with patch.object(
mapped_target,
"create",
side_effect=_return_exception,
):
litellm.completion(model="openai/gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello world"}], client=openai_client)
litellm.completion(
model="openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello world"}],
client=openai_client,
)
pytest.fail("Expected to raise Timeout")
except litellm.Timeout as e:
assert e.status_code == 504
@@ -1350,7 +1368,7 @@ def test_context_window_exceeded_error_from_litellm_proxy():
def test_bad_request_error_with_response_without_request():
"""
Test that BadRequestError handles Response objects without a request attribute.
This simulates a real scenario where a Response is created without a request
(e.g., in tests or when manually creating error responses), and we need to
ensure it doesn't raise RuntimeError when the exception is created.
@@ -1362,8 +1380,7 @@ def test_bad_request_error_with_response_without_request():
# Create a Response without a request (simulates the scenario that was failing)
response_without_request = Response(status_code=400, text="Bad Request")
# Test that extract_and_raise_litellm_exception can handle this
args = {
"response": response_without_request,
@@ -1371,17 +1388,17 @@ def test_bad_request_error_with_response_without_request():
"model": "gpt-3.5-turbo",
"custom_llm_provider": "openai",
}
# This should raise BadRequestError without RuntimeError
with pytest.raises(litellm.BadRequestError) as exc_info:
extract_and_raise_litellm_exception(**args)
# Verify the exception was created successfully
error = exc_info.value
assert error is not None
assert error.model == "gpt-3.5-turbo"
assert error.llm_provider == "openai"
# Verify the exception has a response (should be minimal error response)
assert error.response is not None
# The response should have a request (minimal error response has one)
@@ -1420,6 +1437,3 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model):
assert exc_info.value.code == "invalid_value"
assert exc_info.value.param is not None
assert exc_info.value.type == "invalid_request_error"
tests/local_testing/test_loadtest_router.py
+2 -2
View File
@@ -39,8 +39,8 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# },
tests/local_testing/test_prompt_injection_detection.py
+2 -2
View File
@@ -108,9 +108,9 @@ async def test_prompt_injection_llm_eval():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
tests/local_testing/test_router_budget_limiter.py
+3 -2
View File
@@ -75,9 +75,9 @@ async def test_provider_budgets_e2e_test():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"model_info": {"id": "azure-model-id"},
},
@@ -609,6 +609,7 @@ async def test_deployment_budgets_e2e_test_expect_to_fail():
assert "Exceeded budget for deployment" in str(exc_info.value)
@pytest.mark.flaky(retries=6, delay=2)
@pytest.mark.asyncio
async def test_tag_budgets_e2e_test_expect_to_fail():
tests/local_testing/test_router_caching.py
+2 -2
View File
@@ -268,8 +268,8 @@ async def test_acompletion_caching_on_router_caching_groups():
"model_name": "azure-gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
},
"tpm": 100000,
tests/local_testing/test_router_client_init.py
+1 -1
View File
@@ -85,7 +85,7 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab
To allow for local testing without real credentials, first must mock Azure SDK authentication functions
and environment variables.
"""
monkeypatch.delenv("AZURE_API_KEY", raising=False)
monkeypatch.delenv("AZURE_AI_API_KEY", raising=False)
litellm.enable_azure_ad_token_refresh = True
# mock the token provider function
mocked_func_generating_token = MagicMock(return_value="test_token")
tests/local_testing/test_router_cooldown_handlers.py
+2 -2
View File
@@ -45,9 +45,9 @@ async def test_cooldown_badrequest_error():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
tests/local_testing/test_router_debug_logs.py
+2 -2
View File
@@ -34,9 +34,9 @@ def test_async_fallbacks(caplog):
"model_name": "azure/gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"mock_response": "Hello world",
},
"tpm": 240000,
tests/local_testing/test_router_fallbacks.py
+43 -35
View File
@@ -70,7 +70,7 @@ def test_sync_fallbacks():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -79,9 +79,9 @@ def test_sync_fallbacks():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -92,7 +92,7 @@ def test_sync_fallbacks():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -132,7 +132,9 @@ def test_sync_fallbacks():
response = router.completion(**kwargs)
print(f"response: {response}")
time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 3 # 1 init call + 2 retries (fallback not counted as previous)
assert (
customHandler.previous_models == 3
) # 1 init call + 2 retries (fallback not counted as previous)
print("Passed ! Test router_fallbacks: test_sync_fallbacks()")
router.reset()
@@ -153,7 +155,7 @@ async def test_async_fallbacks():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -162,9 +164,9 @@ async def test_async_fallbacks():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -175,7 +177,7 @@ async def test_async_fallbacks():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -220,7 +222,9 @@ async def test_async_fallbacks():
await asyncio.sleep(
0.05
) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 3 # 1 init call + 2 retries (fallback not counted as previous)
assert (
customHandler.previous_models == 3
) # 1 init call + 2 retries (fallback not counted as previous)
router.reset()
except litellm.Timeout as e:
pass
@@ -242,7 +246,7 @@ def test_sync_fallbacks_embeddings():
"model": "azure/text-embedding-ada-002",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -292,7 +296,7 @@ async def test_async_fallbacks_embeddings():
"model": "azure/text-embedding-ada-002",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -348,7 +352,7 @@ def test_dynamic_fallbacks_sync():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -357,9 +361,9 @@ def test_dynamic_fallbacks_sync():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -370,7 +374,7 @@ def test_dynamic_fallbacks_sync():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -403,7 +407,9 @@ def test_dynamic_fallbacks_sync():
response = router.completion(**kwargs)
print(f"response: {response}")
time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models >= 3 # 1 init call, retries, 1 fallback (count varies with cooldown timing)
assert (
customHandler.previous_models >= 3
) # 1 init call, retries, 1 fallback (count varies with cooldown timing)
router.reset()
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
@@ -425,7 +431,7 @@ async def test_dynamic_fallbacks_async():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -434,9 +440,9 @@ async def test_dynamic_fallbacks_async():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -447,7 +453,7 @@ async def test_dynamic_fallbacks_async():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -489,7 +495,9 @@ async def test_dynamic_fallbacks_async():
await asyncio.sleep(
0.05
) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models >= 3 # 1 init call, retries, 1 fallback (count varies with cooldown timing)
assert (
customHandler.previous_models >= 3
) # 1 init call, retries, 1 fallback (count varies with cooldown timing)
router.reset()
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
@@ -562,7 +570,7 @@ def test_sync_fallbacks_streaming():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -571,9 +579,9 @@ def test_sync_fallbacks_streaming():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -584,7 +592,7 @@ def test_sync_fallbacks_streaming():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -643,7 +651,7 @@ async def test_async_fallbacks_max_retries_per_request():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -652,9 +660,9 @@ async def test_async_fallbacks_max_retries_per_request():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -665,7 +673,7 @@ async def test_async_fallbacks_max_retries_per_request():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -750,9 +758,9 @@ def test_ausage_based_routing_fallbacks():
def get_azure_params(deployment_name: str):
params = {
"model": f"azure/{deployment_name}",
"api_key": os.environ["AZURE_API_KEY"],
"api_key": os.environ["AZURE_AI_API_KEY"],
"api_version": os.environ["AZURE_API_VERSION"],
"api_base": os.environ["AZURE_API_BASE"],
"api_base": os.environ["AZURE_AI_API_BASE"],
}
return params
@@ -855,7 +863,7 @@ def test_custom_cooldown_times():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 24000000,
},
@@ -863,9 +871,9 @@ def test_custom_cooldown_times():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 1,
},
tests/local_testing/test_router_init.py
+14 -14
View File
@@ -41,9 +41,9 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
@@ -97,9 +97,9 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# },
# },
# ]
@@ -135,7 +135,7 @@
# "model_name": "azure-cloudflare",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": "https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1",
# },
@@ -202,9 +202,9 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "timeout": 0.000001,
# "stream_timeout": 0.000_001,
# },
@@ -255,9 +255,9 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "timeout": 200, # regular calls will not timeout, stream calls will
# "stream_timeout": 10,
# },
@@ -349,7 +349,7 @@
# "model_name": "gpt-4-vision-enhancements",
# "litellm_params": {
# "model": "azure/gpt-4-vision",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/",
# "dataSources": [
# {
@@ -616,9 +616,9 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
@@ -661,9 +661,9 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
tests/local_testing/test_router_timeout.py
+2 -2
View File
@@ -31,8 +31,8 @@ def test_router_timeouts():
"model_name": "openai-gpt-4",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "os.environ/AZURE_API_BASE",
"api_key": "os.environ/AZURE_AI_API_KEY",
"api_base": "os.environ/AZURE_AI_API_BASE",
"api_version": "os.environ/AZURE_API_VERSION",
},
"tpm": 80000,
tests/local_testing/test_router_utils.py
+20 -9
View File
@@ -35,7 +35,7 @@ def test_returned_settings():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -99,7 +99,7 @@ def test_update_kwargs_before_fallbacks_unit_test():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
@@ -136,7 +136,7 @@ async def test_update_kwargs_before_fallbacks(call_type):
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
@@ -266,6 +266,7 @@ async def test_call_router_callbacks_on_success():
)
assert increment["increment_value"] == 1
@pytest.mark.serial
@pytest.mark.asyncio
async def test_call_router_callbacks_on_failure():
@@ -486,7 +487,9 @@ def test_router_get_deployment_credentials_with_provider():
)
# Test getting credentials by model_id
credentials = router.get_deployment_credentials_with_provider(model_id="openai-deployment-1")
credentials = router.get_deployment_credentials_with_provider(
model_id="openai-deployment-1"
)
assert credentials is not None
assert credentials["api_key"] == "sk-test-123"
assert credentials["custom_llm_provider"] == "openai"
@@ -499,14 +502,16 @@ def test_router_get_deployment_credentials_with_provider():
assert credentials2["custom_llm_provider"] == "anthropic"
# Test with non-existent model
credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent")
credentials3 = router.get_deployment_credentials_with_provider(
model_id="non-existent"
)
assert credentials3 is None
def test_router_get_deployment_credentials_with_provider_wildcard():
"""
Test that get_deployment_credentials_with_provider handles wildcard patterns.
When a model like openai/gpt-4o is requested and the config has openai/*,
the method should resolve the wildcard pattern and return credentials.
"""
@@ -533,20 +538,26 @@ def test_router_get_deployment_credentials_with_provider_wildcard():
)
# Test wildcard pattern matching for OpenAI
credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-4o")
credentials = router.get_deployment_credentials_with_provider(
model_id="openai/gpt-4o"
)
assert credentials is not None
assert credentials["api_key"] == "sk-wildcard-123"
assert credentials["custom_llm_provider"] == "openai"
assert credentials["api_base"] == "https://api.openai.com/v1"
# Test wildcard pattern matching for Anthropic
credentials2 = router.get_deployment_credentials_with_provider(model_id="anthropic/claude-3-opus")
credentials2 = router.get_deployment_credentials_with_provider(
model_id="anthropic/claude-3-opus"
)
assert credentials2 is not None
assert credentials2["api_key"] == "sk-ant-wildcard-456"
assert credentials2["custom_llm_provider"] == "anthropic"
# Test with non-matching model
credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro")
credentials3 = router.get_deployment_credentials_with_provider(
model_id="vertex_ai/gemini-pro"
)
assert credentials3 is None
tests/local_testing/test_timeout.py
+4 -4
View File
@@ -111,8 +111,8 @@ def test_hanging_request_azure():
"model_name": "azure-gpt",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_base": os.environ["AZURE_API_BASE"],
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_AI_API_BASE"],
"api_key": os.environ["AZURE_AI_API_KEY"],
},
},
{
@@ -175,8 +175,8 @@ def test_hanging_request_openai():
"model_name": "azure-gpt",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_base": os.environ["AZURE_API_BASE"],
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_AI_API_BASE"],
"api_key": os.environ["AZURE_AI_API_KEY"],
},
},
{
tests/local_testing/test_tpm_rpm_routing_v2.py
+5 -11
View File
@@ -39,9 +39,7 @@ from create_mock_standard_logging_payload import create_standard_logging_payload
def test_tpm_rpm_updated():
test_cache = DualCache()
lowest_tpm_logger = LowestTPMLoggingHandler(
router_cache=test_cache
)
lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
deployment_id = "1234"
deployment = "azure/gpt-4.1-mini"
@@ -108,9 +106,7 @@ def test_get_available_deployments():
"model_info": {"id": "5678"},
},
]
lowest_tpm_logger = LowestTPMLoggingHandler(
router_cache=test_cache
)
lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
## DEPLOYMENT 1 ##
total_tokens = 50
@@ -669,9 +665,7 @@ def test_return_potential_deployments():
"""
test_cache = DualCache()
lowest_tpm_logger = LowestTPMLoggingHandler(
router_cache=test_cache
)
lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)
args: Dict = {
"healthy_deployments": [
@@ -731,8 +725,8 @@ async def test_tpm_rpm_routing_model_name_checks():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"mock_response": "Hey, how's it going?",
},
}
tests/logging_callback_tests/test_alerting.py
+10 -15
View File
@@ -641,7 +641,7 @@ async def test_outage_alerting_called(
"model_name": model,
"litellm_params": {
"model": model,
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": api_base,
"vertex_location": vertex_location,
"vertex_project": vertex_project,
@@ -749,7 +749,7 @@ async def test_region_outage_alerting_called(
"model_name": model,
"litellm_params": {
"model": model,
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": api_base,
"vertex_location": vertex_location,
"vertex_project": vertex_project,
@@ -760,7 +760,7 @@ async def test_region_outage_alerting_called(
"model_name": model,
"litellm_params": {
"model": model,
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": api_base,
"vertex_location": vertex_location,
"vertex_project": "vertex_project-2",
@@ -868,7 +868,9 @@ async def test_langfuse_trace_id():
returned_trace_id = trace_url.split("/")[-1]
assert returned_trace_id == litellm_logging_obj._get_trace_id(service_name="langfuse")
assert returned_trace_id == litellm_logging_obj._get_trace_id(
service_name="langfuse"
)
@pytest.mark.asyncio
@@ -1007,7 +1009,7 @@ async def test_soft_budget_alerts():
# Verify alert message contains correct percentage
alert_message = mock_send_alert.call_args[1]["message"]
print("GOT MESSAGE\n\n", alert_message)
expected_message = (
@@ -1077,10 +1079,10 @@ key_no_max_budget_info = CallInfo(
async def test_soft_budget_alerts_webhook(entity_info):
"""
Tests that soft budget alerts are triggered for different entity types.
Tests:
- Key with max budget
- Team
- Team
- User
- Key without max budget
"""
@@ -1097,7 +1099,7 @@ async def test_soft_budget_alerts_webhook(entity_info):
# Verify the webhook event
call_args = mock_send_alert.call_args[1]
logged_webhook_event: WebhookEvent = call_args["user_info"]
# Validate the webhook event has all expected fields
assert logged_webhook_event.spend == entity_info.spend
assert logged_webhook_event.soft_budget == entity_info.soft_budget
@@ -1106,10 +1108,3 @@ async def test_soft_budget_alerts_webhook(entity_info):
assert logged_webhook_event.user_email == entity_info.user_email
assert logged_webhook_event.key_alias == entity_info.key_alias
assert logged_webhook_event.event_group == entity_info.event_group
tests/logging_callback_tests/test_custom_callback_router.py
+27 -21
View File
@@ -267,7 +267,10 @@ class CompletionCustomHandler(
try:
print("CompletionCustomHandler.async_log_success_event, kwargs: ", kwargs)
self.states.append("async_success")
print("############### CompletionCustomHandler async success, kwargs: ", kwargs)
print(
"############### CompletionCustomHandler async success, kwargs: ",
kwargs,
)
## START TIME
assert isinstance(start_time, datetime)
## END TIME
@@ -396,9 +399,9 @@ async def test_async_chat_azure():
"model_name": "gpt-4.1-nano", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"model_info": {"base_model": "azure/gpt-4.1-mini"},
"tpm": 240000,
@@ -443,7 +446,7 @@ async def test_async_chat_azure():
"model": "azure/gpt-4o-new-test",
"api_key": "my-bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -483,9 +486,9 @@ async def test_async_embedding_azure():
"model_name": "azure-embedding-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -506,7 +509,7 @@ async def test_async_embedding_azure():
"model": "azure/text-embedding-ada-002",
"api_key": "my-bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -549,7 +552,7 @@ async def test_async_chat_azure_with_fallbacks():
"model": "azure/gpt-4.1-mini",
"api_key": "my-bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -608,9 +611,9 @@ async def test_async_completion_azure_caching():
"model_name": "gpt-4.1-nano", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -664,23 +667,23 @@ async def test_async_completion_azure_caching_streaming():
)
litellm.callbacks = [customHandler_caching]
unique_time = uuid.uuid4()
# Use Router instead of direct litellm.acompletion to get router-specific metadata
model_list = [
{
"model_name": "gpt-4.1-nano",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
},
]
router = Router(model_list=model_list)
response1 = await router.acompletion(
model="gpt-4.1-nano",
messages=[
@@ -725,12 +728,16 @@ async def test_async_embedding_azure_caching():
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
router = Router(model_list=[{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "openai/text-embedding-ada-002",
},
}])
router = Router(
model_list=[
{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "openai/text-embedding-ada-002",
},
}
]
)
litellm.callbacks = [customHandler_caching]
unique_time = time.time()
response1 = await router.aembedding(
@@ -818,4 +825,3 @@ async def test_rate_limit_error_callback():
assert "original_model_group" in mock_client.call_args.kwargs
assert mock_client.call_args.kwargs["original_model_group"] == "my-test-gpt"
tests/old_proxy_tests/tests/load_test_q.py
+2 -2
View File
@@ -26,7 +26,7 @@ config = {
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.environ["AZURE_API_KEY"],
"api_key": os.environ["AZURE_AI_API_KEY"],
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-07-01-preview",
},
@@ -34,7 +34,7 @@ config = {
]
}
print("STARTING LOAD TEST Q")
print(os.environ["AZURE_API_KEY"])
print(os.environ["AZURE_AI_API_KEY"])
response = requests.post(
url=f"{base_url}/key/generate",
tests/pass_through_unit_tests/test_unit_test_passthrough_router.py
+1 -1
View File
@@ -102,7 +102,7 @@ class TestPassthroughEndpointRouter(unittest.TestCase):
mock_get_secret.return_value = "env_azure_key"
result = self.router.get_credentials("azure", None)
self.assertEqual(result, "env_azure_key")
mock_get_secret.assert_called_once_with("AZURE_API_KEY")
mock_get_secret.assert_called_once_with("AZURE_AI_API_KEY")
def test_default_env_variable_method(self):
"""
tests/proxy_unit_tests/example_config_yaml/azure_config.yaml
+2 -2
View File
@@ -4,12 +4,12 @@ model_list:
model: azure/gpt-4.1-mini
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
tpm: 20_000
- model_name: gpt-4-team2
litellm_params:
model: azure/gpt-4
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
api_base: https://openai-gpt-4-test-v-2.openai.azure.com/
tpm: 100_000
tests/proxy_unit_tests/test_configs/test_bad_config.yaml
+4 -4
View File
@@ -6,16 +6,16 @@ model_list:
- model_name: working-azure-gpt-3.5-turbo
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
- model_name: azure-gpt-3.5-turbo
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_base: os.environ/AZURE_AI_API_BASE
api_key: bad-key
- model_name: azure-embedding
litellm_params:
model: azure/text-embedding-ada-002
api_base: os.environ/AZURE_API_BASE
api_base: os.environ/AZURE_AI_API_BASE
api_key: bad-key
tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml
+1 -1
View File
@@ -3,7 +3,7 @@ model_list:
litellm_params:
model: azure/gpt-4.1-mini
api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
litellm_settings:
tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml
+7 -7
View File
@@ -11,7 +11,7 @@ model_list:
model_name: azure-model
- litellm_params:
api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
model: azure/gpt-4.1-mini
model_name: azure-cloudflare-model
- litellm_params:
@@ -49,8 +49,8 @@ model_list:
id: 79fc75bf-8e1b-47d5-8d24-9365a854af03
model_name: test_openai_models
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
model: azure/text-embedding-ada-002
model_info:
@@ -94,16 +94,16 @@ model_list:
mode: image_generation
model_name: dall-e-3
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-06-01-preview
model: azure/
model_info:
mode: image_generation
model_name: dall-e-2
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
model: azure/text-embedding-ada-002
model_info:
tests/proxy_unit_tests/test_proxy_pass_user_config.py
+5 -4
View File
@@ -54,8 +54,9 @@ def client_no_auth():
@pytest.mark.skipif(
os.environ.get("AZURE_API_KEY") is None or os.environ.get("OPENAI_API_KEY") is None,
reason="AZURE_API_KEY or OPENAI_API_KEY not set - skipping integration test"
os.environ.get("AZURE_AI_API_KEY") is None
or os.environ.get("OPENAI_API_KEY") is None,
reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test",
)
def test_chat_completion(client_no_auth):
global headers
@@ -69,9 +70,9 @@ def test_chat_completion(client_no_auth):
model_name="user-azure-instance",
litellm_params=CompletionRequest(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_version=os.getenv("AZURE_API_VERSION"),
api_base=os.getenv("AZURE_API_BASE"),
api_base=os.getenv("AZURE_AI_API_BASE"),
timeout=10,
),
tpm=240000,
tests/proxy_unit_tests/test_proxy_server.py
+104 -81
View File
@@ -119,7 +119,7 @@ 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_API_BASE", "http://fake-azure-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")
@@ -178,7 +178,7 @@ def test_chat_completion(mock_acompletion, client_no_auth):
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.
@@ -188,33 +188,41 @@ def test_chat_completion_malformed_messages_returns_400(client_no_auth):
# 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"}]
"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)
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}"
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')}"
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)}")
@@ -342,7 +350,7 @@ def test_chat_completion_forward_llm_provider_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:
@@ -351,7 +359,7 @@ def test_chat_completion_forward_llm_provider_auth_headers(
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",
@@ -360,7 +368,7 @@ def test_chat_completion_forward_llm_provider_auth_headers(
],
"max_tokens": 10,
}
# Headers including LLM provider auth
request_headers = {
"Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped)
@@ -368,17 +376,17 @@ def test_chat_completion_forward_llm_provider_auth_headers(
"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
@@ -389,19 +397,23 @@ def test_chat_completion_forward_llm_provider_auth_headers(
# 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"✓ 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)}")
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")
@@ -2406,9 +2418,7 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch):
test_model_list_2 = [{"model_name": "model-b"}]
called_model_lists = []
async def fake_perform_health_check(
model_list, details, max_concurrency=None
):
async def fake_perform_health_check(model_list, details, max_concurrency=None):
called_model_lists.append(copy.deepcopy(model_list))
return (["healthy"], ["unhealthy"])
@@ -2452,13 +2462,14 @@ async def test_background_health_check_skip_disabled_models(monkeypatch):
test_model_list = [
{"model_name": "model-a"},
{"model_name": "model-b", "model_info": {"disable_background_health_check": True}},
{
"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
):
async def fake_perform_health_check(model_list, details, max_concurrency=None):
called_model_lists.append(copy.deepcopy(model_list))
return (["healthy"], [])
@@ -2500,15 +2511,15 @@ def test_get_timeout_from_request():
@pytest.mark.parametrize(
"ui_exists, ui_has_content",
[
(True, True), # UI path exists and 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
(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
@@ -2516,44 +2527,54 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content
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")
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))}")
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}")
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)
@@ -2562,7 +2583,9 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content
# 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.")
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)
@@ -2574,17 +2597,17 @@ 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"
- 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"]
"callbacks": ["otel", "langsmith"],
},
"environment_variables": {
"LANGFUSE_PUBLIC_KEY": "test-public-key",
@@ -2595,51 +2618,53 @@ async def test_get_config_callbacks_with_all_types(client_no_auth):
"OTEL_ENDPOINT": "http://localhost:4317",
"LANGSMITH_API_KEY": "test-langsmith-key",
},
"general_settings": {}
"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"]
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]
@@ -2654,13 +2679,13 @@ async def test_get_config_callbacks_environment_variables(client_no_auth):
for each callback type. Values are returned as-is from the config (no decryption).
"""
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"]
"callbacks": ["otel"],
},
"environment_variables": {
"LANGFUSE_PUBLIC_KEY": "test-public-key",
@@ -2670,21 +2695,21 @@ async def test_get_config_callbacks_environment_variables(client_no_auth):
"OTEL_ENDPOINT": "http://localhost:4317",
"OTEL_HEADERS": "key=value",
},
"general_settings": {}
"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
@@ -2692,7 +2717,7 @@ async def test_get_config_callbacks_environment_variables(client_no_auth):
assert langfuse_callback is not None
assert langfuse_callback["type"] == "success"
assert "variables" in langfuse_callback
# Verify langfuse env vars are present (values returned as-is, no decryption)
langfuse_vars = langfuse_callback["variables"]
assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars
@@ -2701,15 +2726,13 @@ async def test_get_config_callbacks_environment_variables(client_no_auth):
assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "test-secret-key"
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
)
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
# Verify otel env vars are present
otel_vars = otel_callback["variables"]
assert "OTEL_EXPORTER" in otel_vars
tests/test_litellm/llms/azure/test_azure_common_utils.py
+2 -2
View File
@@ -460,7 +460,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type):
"api_key": "test-api-key",
"api_version": os.getenv("AZURE_API_VERSION", "2023-05-15"),
"api_base": os.getenv(
"AZURE_API_BASE", "https://test.openai.azure.com"
"AZURE_AI_API_BASE", "https://test.openai.azure.com"
),
},
}
@@ -670,7 +670,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used_azure_text(call_ty
"api_key": "test-api-key",
"api_version": os.getenv("AZURE_API_VERSION", "2023-05-15"),
"api_base": os.getenv(
"AZURE_API_BASE", "https://test.openai.azure.com"
"AZURE_AI_API_BASE", "https://test.openai.azure.com"
),
},
}
tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py
+2 -1
View File
@@ -3,6 +3,7 @@ Tests for Azure AI Anthropic CountTokens transformation.
Verifies that the CountTokens API uses the correct authentication headers.
"""
import os
import sys
@@ -40,7 +41,7 @@ class TestAzureAIAnthropicCountTokensConfig:
assert headers["anthropic-version"] == "2023-06-01"
assert "anthropic-beta" in headers
def test_get_required_headers_includes_azure_api_key(self):
def test_get_required_headers_includes_AZURE_AI_API_KEY(self):
"""
Test that get_required_headers includes Azure api-key header.
tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
+13 -13
View File
@@ -336,7 +336,9 @@ class TestOpenAIResponsesAPIConfig:
)
assert isinstance(result, ImageGenerationPartialImageEvent)
assert result.type == ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE
assert (
result.type == ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE
)
assert result.partial_image_index == idx
assert result.b64_json == chunk["b64_json"]
@@ -689,9 +691,7 @@ class TestTransformListInputItemsRequest:
def test_openai_transform_compact_response_api_request_query_params_preserved(self):
"""Test compact URL construction preserves query params and appends path."""
# Setup
azure_style_api_base = (
"https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
)
azure_style_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
# Execute
url, data = self.openai_config.transform_compact_response_api_request(
@@ -731,12 +731,12 @@ class TestTransformListInputItemsRequest:
def test_azure_transform_list_input_items_request_minimal(self):
"""Test Azure implementation with minimal parameters"""
# Setup
azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
# Execute
url, params = self.azure_config.transform_list_input_items_request(
response_id=self.response_id,
api_base=azure_api_base,
api_base=AZURE_AI_API_BASE,
litellm_params=self.litellm_params,
headers=self.headers,
)
@@ -749,12 +749,12 @@ class TestTransformListInputItemsRequest:
def test_azure_transform_list_input_items_request_url_construction(self):
"""Test Azure implementation URL construction with response_id in path"""
# Setup
azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
# Execute
url, params = self.azure_config.transform_list_input_items_request(
response_id=self.response_id,
api_base=azure_api_base,
api_base=AZURE_AI_API_BASE,
litellm_params=self.litellm_params,
headers=self.headers,
)
@@ -768,12 +768,12 @@ class TestTransformListInputItemsRequest:
def test_azure_transform_list_input_items_request_with_all_params(self):
"""Test Azure implementation with all optional parameters"""
# Setup
azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
# Execute
url, params = self.azure_config.transform_list_input_items_request(
response_id=self.response_id,
api_base=azure_api_base,
api_base=AZURE_AI_API_BASE,
litellm_params=self.litellm_params,
headers=self.headers,
after="cursor_after_123",
@@ -1128,9 +1128,9 @@ class TestPhaseParameter:
phase = getattr(output_item, "phase", None)
expected = "commentary" if idx == 0 else "final_answer"
assert phase == expected, (
f"output[{idx}] phase={phase!r}, expected {expected!r}"
)
assert (
phase == expected
), f"output[{idx}] phase={phase!r}, expected {expected!r}"
def test_streaming_output_item_done_preserves_phase(self):
"""OutputItemDoneEvent must preserve phase on its item."""
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
+4 -8
View File
@@ -28,7 +28,7 @@ def test_get_api_key():
assert get_api_key(
custom_litellm_key_header=None,
api_key=bearer_token,
azure_api_key_header=None,
AZURE_AI_API_KEY_header=None,
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
@@ -59,7 +59,7 @@ def test_get_api_key_with_custom_litellm_key_header(
assert get_api_key(
custom_litellm_key_header=custom_litellm_key_header,
api_key=None,
azure_api_key_header=None,
AZURE_AI_API_KEY_header=None,
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
@@ -334,7 +334,6 @@ async def test_proxy_admin_expired_key_from_cache():
"litellm.proxy.auth.user_api_key_auth._delete_cache_key_object",
new_callable=AsyncMock,
) as mock_delete_cache:
mock_get_key_object.return_value = expired_token
# Set attributes on proxy_server module (these are imported inside _user_api_key_auth_builder)
@@ -372,7 +371,7 @@ async def test_proxy_admin_expired_key_from_cache():
await _user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}", # Add Bearer prefix
azure_api_key_header="",
AZURE_AI_API_KEY_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
@@ -609,7 +608,6 @@ class TestJWTOAuth2Coexistence:
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
new_callable=AsyncMock,
) as mock_jwt_auth:
litellm.proxy.proxy_server.jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
@@ -674,7 +672,6 @@ class TestJWTOAuth2Coexistence:
new_callable=AsyncMock,
return_value=mock_jwt_result,
) as mock_jwt_auth:
litellm.proxy.proxy_server.jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
@@ -726,7 +723,6 @@ class TestJWTOAuth2Coexistence:
new_callable=AsyncMock,
return_value=mock_oauth2_response,
) as mock_oauth2:
result = await user_api_key_auth(
request=mock_request,
api_key=f"Bearer {jwt_like_token}",
@@ -846,7 +842,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls():
await _user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
AZURE_AI_API_KEY_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
+223 -116
View File
@@ -34,8 +34,8 @@ def llm_router() -> Router:
"model_name": "azure-gpt-3-5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-2",
"api_key": "azure_api_key",
"api_base": "azure_api_base",
"api_key": "AZURE_AI_API_KEY",
"api_base": "AZURE_AI_API_BASE",
"api_version": "azure_api_version",
},
"model_info": {
@@ -106,16 +106,18 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router:
Asserts 'create_file' is called with the correct arguments
"""
import litellm
import litellm.proxy.proxy_server as ps
from litellm import Router
from litellm.proxy._types import LitellmUserRoles
import litellm.proxy.proxy_server as ps
from litellm.proxy.utils import ProxyLogging
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
# Mock create_file as an async function
mock_create_file = mocker.patch("litellm.files.main.create_file", new=mocker.AsyncMock())
mock_create_file = mocker.patch(
"litellm.files.main.create_file", new=mocker.AsyncMock()
)
proxy_logging_obj = ProxyLogging(
user_api_key_cache=DualCache(default_in_memory_ttl=1)
@@ -127,7 +129,14 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router:
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
class DummyManagedFiles(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
# Handle both dict and object forms of create_file_request
if isinstance(create_file_request, dict):
file_data = create_file_request.get("file")
@@ -135,12 +144,12 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router:
else:
file_data = create_file_request.file
purpose_data = create_file_request.purpose
# Call the mocked litellm.files.main.create_file to ensure asserts work
await litellm.files.main.create_file(
custom_llm_provider="azure",
model="azure/chatgpt-v-2",
api_key="azure_api_key",
api_key="AZURE_AI_API_KEY",
file=file_data[1],
purpose=purpose_data,
)
@@ -153,6 +162,7 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router:
)
# Return a dummy response object as needed by the test
from litellm.types.llms.openai import OpenAIFileObject
return OpenAIFileObject(
id="dummy-id",
object="file",
@@ -162,17 +172,21 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router:
purpose=purpose_data,
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
# Manually add the hook to the proxy_hook_mapping
@@ -214,7 +228,7 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router:
if (
kwargs.get("custom_llm_provider") == "azure"
and kwargs.get("model") == "azure/chatgpt-v-2"
and kwargs.get("api_key") == "azure_api_key"
and kwargs.get("api_key") == "AZURE_AI_API_KEY"
):
azure_call_found = True
break
@@ -245,8 +259,8 @@ def test_target_storage_invokes_storage_backend(
"""
Ensure target_storage is parsed and invokes the storage backend service.
"""
from litellm.proxy._types import LitellmUserRoles
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
@@ -304,8 +318,8 @@ def test_target_storage_with_target_models(
"""
Ensure target_storage and target_model_names are parsed and passed through.
"""
from litellm.proxy._types import LitellmUserRoles
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
@@ -611,7 +625,9 @@ def test_create_file_for_each_model(
assert openai_call_found, "OpenAI call not found with expected parameters"
def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_router: Router):
def test_create_file_with_expires_after(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
"""
Test that expires_after is properly parsed and passed through when creating a file
"""
@@ -624,18 +640,25 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_
proxy_logging_obj._add_proxy_hooks(llm_router)
class DummyManagedFiles(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
# Verify expires_after is in the request
if isinstance(create_file_request, dict):
expires_after = create_file_request.get("expires_after")
else:
expires_after = getattr(create_file_request, "expires_after", None)
# Verify expires_after was passed correctly
assert expires_after is not None, "expires_after should be in the request"
assert expires_after["anchor"] == "created_at"
assert expires_after["seconds"] == 2592000
# Return a dummy response
return OpenAIFileObject(
id="file-abc123",
@@ -646,17 +669,21 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_
purpose="fine-tune",
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles()
@@ -688,7 +715,9 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_
assert result["purpose"] == "fine-tune"
def test_create_file_with_expires_after_missing_anchor(mocker: MockerFixture, monkeypatch, llm_router: Router):
def test_create_file_with_expires_after_missing_anchor(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
"""
Test that an error is returned when expires_after[anchor] is missing
"""
@@ -717,10 +746,15 @@ def test_create_file_with_expires_after_missing_anchor(mocker: MockerFixture, mo
assert response.status_code == 400
error_detail = response.json()
assert "expires_after" in error_detail["error"]["message"].lower() or "both" in error_detail["error"]["message"].lower()
assert (
"expires_after" in error_detail["error"]["message"].lower()
or "both" in error_detail["error"]["message"].lower()
)
def test_create_file_with_expires_after_missing_seconds(mocker: MockerFixture, monkeypatch, llm_router: Router):
def test_create_file_with_expires_after_missing_seconds(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
"""
Test that an error is returned when expires_after[seconds] is missing
"""
@@ -749,10 +783,15 @@ def test_create_file_with_expires_after_missing_seconds(mocker: MockerFixture, m
assert response.status_code == 400
error_detail = response.json()
assert "expires_after" in error_detail["error"]["message"].lower() or "both" in error_detail["error"]["message"].lower()
assert (
"expires_after" in error_detail["error"]["message"].lower()
or "both" in error_detail["error"]["message"].lower()
)
def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monkeypatch, llm_router: Router):
def test_create_file_with_expires_after_valid_values(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
"""
Test that expires_after works with valid anchor and seconds values
"""
@@ -765,18 +804,25 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk
proxy_logging_obj._add_proxy_hooks(llm_router)
class DummyManagedFiles(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
# Verify expires_after is in the request
if isinstance(create_file_request, dict):
expires_after = create_file_request.get("expires_after")
else:
expires_after = getattr(create_file_request, "expires_after", None)
# Verify expires_after was passed correctly
assert expires_after is not None, "expires_after should be in the request"
assert expires_after["anchor"] == "created_at"
assert expires_after["seconds"] == 3600
return OpenAIFileObject(
id="file-abc123",
object="file",
@@ -786,17 +832,21 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk
purpose="fine-tune",
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles()
@@ -827,7 +877,9 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk
assert result["purpose"] == "fine-tune"
def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, llm_router: Router):
def test_create_file_without_expires_after(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
"""
Test that file creation works normally without expires_after
"""
@@ -840,16 +892,25 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l
proxy_logging_obj._add_proxy_hooks(llm_router)
class DummyManagedFiles(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
# Verify expires_after is None when not provided
if isinstance(create_file_request, dict):
expires_after = create_file_request.get("expires_after")
else:
expires_after = getattr(create_file_request, "expires_after", None)
# expires_after should be None when not provided
assert expires_after is None, "expires_after should be None when not provided"
assert (
expires_after is None
), "expires_after should be None when not provided"
return OpenAIFileObject(
id="file-abc123",
object="file",
@@ -859,17 +920,21 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l
purpose="fine-tune",
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles()
@@ -898,11 +963,13 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l
assert result["purpose"] == "fine-tune"
def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, llm_router: Router):
def test_managed_files_with_loadbalancing(
mocker: MockerFixture, monkeypatch, llm_router: Router
):
"""
Test that managed files work with loadbalancing when both target_model_names
and enable_loadbalancing_on_batch_endpoints are enabled.
This ensures that the priority order is correct:
- managed files should take precedence over deprecated loadbalancing
- managed files internally use llm_router.acreate_file() which provides loadbalancing
@@ -912,28 +979,34 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll
# Enable loadbalancing on batch endpoints
monkeypatch.setattr("litellm.enable_loadbalancing_on_batch_endpoints", True)
proxy_logging_obj = ProxyLogging(
user_api_key_cache=DualCache(default_in_memory_ttl=1)
)
proxy_logging_obj._add_proxy_hooks(llm_router)
# Track calls to verify loadbalancing through router
router_acreate_file_calls = []
class ManagedFilesWithLoadbalancing(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
# Verify we receive the target model names
assert len(target_model_names_list) > 0, "Should have target_model_names_list"
assert (
len(target_model_names_list) > 0
), "Should have target_model_names_list"
# Simulate what managed files does - call llm_router.acreate_file for each model
# This is where loadbalancing happens internally
for model in target_model_names_list:
router_acreate_file_calls.append({
"model": model,
"via_router": True
})
router_acreate_file_calls.append({"model": model, "via_router": True})
# Return a managed file ID (base64 encoded)
return OpenAIFileObject(
id="litellm_managed_file_abc123",
@@ -944,23 +1017,29 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll
purpose="batch",
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
proxy_logging_obj.proxy_hook_mapping["managed_files"] = ManagedFilesWithLoadbalancing()
proxy_logging_obj.proxy_hook_mapping[
"managed_files"
] = ManagedFilesWithLoadbalancing()
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
@@ -971,12 +1050,12 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
# Create batch file content
test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}'
test_file = ("batch_data.jsonl", test_file_content, "application/jsonl")
# Make request with both target_model_names AND enable_loadbalancing_on_batch_endpoints
response = client.post(
"/v1/files",
@@ -987,7 +1066,7 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll
},
headers={"Authorization": "Bearer test-key"},
)
# Verify success
assert response.status_code == 200, response.text
finally:
@@ -995,13 +1074,17 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll
result = response.json()
assert result["id"] == "litellm_managed_file_abc123"
assert result["purpose"] == "batch"
# Verify that managed files was called (via router for loadbalancing)
# This proves that managed files took precedence over deprecated loadbalancing
assert len(router_acreate_file_calls) == 2, "Should have called router for both models"
assert (
len(router_acreate_file_calls) == 2
), "Should have called router for both models"
assert router_acreate_file_calls[0]["model"] == "azure-gpt-3-5-turbo"
assert router_acreate_file_calls[1]["model"] == "gpt-3.5-turbo"
assert all(call["via_router"] for call in router_acreate_file_calls), "All calls should go through router"
assert all(
call["via_router"] for call in router_acreate_file_calls
), "All calls should go through router"
def test_create_file_with_nested_litellm_metadata(
@@ -1009,22 +1092,29 @@ def test_create_file_with_nested_litellm_metadata(
):
"""
Test that nested litellm_metadata is correctly parsed from form data in bracket notation.
Regression test for: litellm_metadata[spend_logs_metadata][owner] format should be
correctly parsed into nested dictionary structure.
"""
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.types.llms.openai import OpenAIFileObject
proxy_logging_obj = ProxyLogging(
user_api_key_cache=DualCache(default_in_memory_ttl=1)
)
proxy_logging_obj._add_proxy_hooks(llm_router)
captured_litellm_metadata = {}
class DummyManagedFiles(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
# Capture litellm_metadata for verification
if isinstance(create_file_request, dict):
captured_litellm_metadata.update(
@@ -1034,7 +1124,7 @@ def test_create_file_with_nested_litellm_metadata(
captured_litellm_metadata.update(
getattr(create_file_request, "litellm_metadata", {})
)
return OpenAIFileObject(
id="file-test-123",
object="file",
@@ -1044,28 +1134,32 @@ def test_create_file_with_nested_litellm_metadata(
purpose="fine-tune",
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles()
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
)
test_file_content = b'{"prompt": "Hello", "completion": "Hi"}'
test_file = ("test.jsonl", test_file_content, "application/jsonl")
# Test with nested litellm_metadata in bracket notation
response = client.post(
"/v1/files",
@@ -1080,12 +1174,12 @@ def test_create_file_with_nested_litellm_metadata(
},
headers={"Authorization": "Bearer test-key"},
)
# Verify success
assert response.status_code == 200
result = response.json()
assert result["id"] == "file-test-123"
# Verify nested metadata was correctly parsed
assert "spend_logs_metadata" in captured_litellm_metadata
assert captured_litellm_metadata["spend_logs_metadata"]["owner"] == "john_doe"
@@ -1099,26 +1193,33 @@ def test_create_file_with_deep_nested_litellm_metadata(
):
"""
Test that deeply nested litellm_metadata is correctly parsed from form data.
Regression test for: litellm_metadata[a][b][c] format should be correctly parsed.
"""
import litellm.proxy.proxy_server as ps
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.proxy._types import LitellmUserRoles
import litellm.proxy.proxy_server as ps
from litellm.types.llms.openai import OpenAIFileObject
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
proxy_logging_obj = ProxyLogging(
user_api_key_cache=DualCache(default_in_memory_ttl=1)
)
proxy_logging_obj._add_proxy_hooks(llm_router)
captured_litellm_metadata = {}
class DummyManagedFiles(BaseFileEndpoints):
async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
async def acreate_file(
self,
llm_router,
create_file_request,
target_model_names_list,
litellm_parent_otel_span,
user_api_key_dict,
):
if isinstance(create_file_request, dict):
captured_litellm_metadata.update(
create_file_request.get("litellm_metadata", {})
@@ -1127,7 +1228,7 @@ def test_create_file_with_deep_nested_litellm_metadata(
captured_litellm_metadata.update(
getattr(create_file_request, "litellm_metadata", {})
)
return OpenAIFileObject(
id="file-test-456",
object="file",
@@ -1137,33 +1238,37 @@ def test_create_file_with_deep_nested_litellm_metadata(
purpose="batch",
status="uploaded",
)
async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
raise NotImplementedError("Not implemented for test")
async def afile_list(self, purpose, litellm_parent_otel_span):
raise NotImplementedError("Not implemented for test")
async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_delete(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
async def afile_content(
self, file_id, litellm_parent_otel_span, llm_router, **data
):
raise NotImplementedError("Not implemented for test")
proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles()
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user"
)
try:
test_file_content = b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo"}}'
test_file = ("nested.jsonl", test_file_content, "application/jsonl")
# Test with deeply nested metadata
response = client.post(
"/v1/files",
@@ -1177,12 +1282,12 @@ def test_create_file_with_deep_nested_litellm_metadata(
},
headers={"Authorization": "Bearer test-key"},
)
# Verify success
assert response.status_code == 200, response.text
result = response.json()
assert result["id"] == "file-test-456"
# Verify deeply nested metadata was correctly parsed
assert "config" in captured_litellm_metadata
assert "database" in captured_litellm_metadata["config"]
@@ -1356,7 +1461,9 @@ def test_file_team_injects_when_caller_sends_nothing(
# ---------------------------------------------------------------------------
def _post_file_raw(monkeypatch, llm_router: Router, team_metadata: dict, form_data: dict):
def _post_file_raw(
monkeypatch, llm_router: Router, team_metadata: dict, form_data: dict
):
"""POST /v1/files and return the raw response (no status assertion)."""
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
tests/test_litellm/test_router.py
+48 -36
View File
@@ -24,9 +24,9 @@ def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
@@ -646,8 +646,15 @@ def test_arouter_responses_api_bridge():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = {"id": "resp_test", "object": "response", "status": "completed", "output": []}
mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}'
mock_response.json.return_value = {
"id": "resp_test",
"object": "response",
"status": "completed",
"output": [],
}
mock_response.text = (
'{"id": "resp_test", "object": "response", "status": "completed", "output": []}'
)
with patch.object(client, "post", return_value=mock_response) as mock_post:
try:
@@ -2147,7 +2154,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint()
)
assert credentials is not None
assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com"
assert (
credentials["aws_bedrock_runtime_endpoint"]
== "https://bedrock-runtime.us-east-1.amazonaws.com"
)
assert credentials["aws_access_key_id"] == "test-access-key"
assert credentials["aws_secret_access_key"] == "test-secret-key"
assert credentials["aws_region_name"] == "us-east-1"
@@ -2169,11 +2179,11 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name():
credential_values={
"api_key": "resolved-api-key",
"api_base": "https://resolved.openai.azure.com",
"api_version": "2024-02-01"
}
"api_version": "2024-02-01",
},
)
]
router = litellm.Router(
model_list=[
{
@@ -2197,7 +2207,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name():
assert credentials["custom_llm_provider"] == "azure"
# Ensure credential name is removed after resolution
assert "litellm_credential_name" not in credentials
# Cleanup
litellm.credential_list = []
@@ -2302,7 +2312,10 @@ async def test_aguardrail_helper():
# Mock the original function
async def mock_original_function(**kwargs):
return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")}
return {
"result": "success",
"selected_guardrail": kwargs.get("selected_guardrail"),
}
result = await router._aguardrail_helper(
model="content-filter",
@@ -2336,7 +2349,10 @@ async def test_aguardrail():
# Mock the original function
async def mock_original_function(**kwargs):
return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")}
return {
"result": "success",
"selected_guardrail": kwargs.get("selected_guardrail"),
}
result = await router.aguardrail(
guardrail_name="content-filter",
@@ -2346,6 +2362,7 @@ async def test_aguardrail():
assert result["result"] == "success"
assert result["selected_guardrail"]["id"] == "guardrail-1"
@pytest.mark.asyncio
async def test_anthropic_messages_call_type_is_cached():
"""
@@ -2417,36 +2434,33 @@ async def test_anthropic_messages_call_type_is_cached():
additional_headers=None,
),
)
cache = DualCache()
deployment_check = PromptCachingDeploymentCheck(cache=cache)
prompt_cache = PromptCachingCache(cache=cache)
# Create messages with enough tokens to pass the caching threshold
test_messages = [
{
"role": "user",
"role": "user",
"content": [
{
"type": "text",
"type": "text",
"text": "test long message here" * 1024,
"cache_control": {
"type": "ephemeral",
"ttl": "5m"
}
"cache_control": {"type": "ephemeral", "ttl": "5m"},
}
]
],
}
]
test_model_id = "test-model-id-123"
# Create a payload with anthropic_messages call type
payload = create_standard_logging_payload()
payload["call_type"] = CallTypes.anthropic_messages.value
payload["messages"] = test_messages
payload["model"] = "anthropic/claude-3-5-sonnet-20240620"
payload["model_id"] = test_model_id
# Log the success event (should cache the model_id)
await deployment_check.async_log_success_event(
kwargs={"standard_logging_object": payload},
@@ -2454,19 +2468,23 @@ async def test_anthropic_messages_call_type_is_cached():
start_time=1234567890.0,
end_time=1234567891.0,
)
# Small delay to ensure cache write completes
await asyncio.sleep(0.1)
# Verify that the model_id was actually cached
cached_result = await prompt_cache.async_get_model_id(
messages=test_messages,
tools=None,
)
# This assertion will FAIL if anthropic_messages is filtered out
assert cached_result is not None, "Model ID should be cached for anthropic_messages call type"
assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}"
assert (
cached_result is not None
), "Model ID should be cached for anthropic_messages call type"
assert (
cached_result["model_id"] == test_model_id
), f"Expected {test_model_id}, got {cached_result['model_id']}"
def test_update_kwargs_with_deployment_propagates_model_tags():
@@ -2682,9 +2700,7 @@ def test_credential_name_injected_as_tag():
)
kwargs: dict = {"metadata": {"tags": ["A.101"]}}
deployment = router.get_deployment_by_model_group_name(
model_group_name="xai-model"
)
deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model")
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
assert "Credential: xAI" in kwargs["metadata"]["tags"]
@@ -2709,9 +2725,7 @@ def test_credential_name_not_duplicated_in_tags():
)
kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}}
deployment = router.get_deployment_by_model_group_name(
model_group_name="xai-model"
)
deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model")
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1
@@ -2733,9 +2747,7 @@ def test_credential_name_not_injected_when_absent():
)
kwargs: dict = {"metadata": {"tags": ["A.101"]}}
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-model"
)
deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model")
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
assert kwargs["metadata"]["tags"] == ["A.101"]
tests/test_litellm/test_setup_wizard.py
+8 -8
View File
@@ -58,7 +58,7 @@ _ANTHROPIC = {
_AZURE = {
"id": "azure",
"name": "Azure OpenAI",
"env_key": "AZURE_API_KEY",
"env_key": "AZURE_AI_API_KEY",
"models": [],
"test_model": None,
"needs_api_base": True,
@@ -123,8 +123,8 @@ def test_build_config_master_key_quoted():
def test_build_config_does_not_mutate_env_vars():
"""_build_config must not modify the caller's env_vars dict."""
env_vars = {
"AZURE_API_KEY": "az-key",
"_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com",
"AZURE_AI_API_KEY": "az-key",
"_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com",
"_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment",
}
original_keys = set(env_vars.keys())
@@ -134,8 +134,8 @@ def test_build_config_does_not_mutate_env_vars():
def test_build_config_azure_uses_deployment_name():
env_vars = {
"AZURE_API_KEY": "az-key",
"_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com",
"AZURE_AI_API_KEY": "az-key",
"_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com",
"_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-gpt4o",
}
config = SetupWizard._build_config([_AZURE], env_vars, "sk-master")
@@ -147,7 +147,7 @@ def test_build_config_azure_uses_deployment_name():
def test_build_config_azure_no_deployment_skipped():
"""Azure without a deployment name should emit nothing (not fallback to gpt-4o)."""
env_vars = {"AZURE_API_KEY": "az-key"} # no deployment sentinel
env_vars = {"AZURE_AI_API_KEY": "az-key"} # no deployment sentinel
config = SetupWizard._build_config([_AZURE], env_vars, "sk-master")
# No azure model entry should be emitted when deployment name is absent
assert "model: azure/" not in config
@@ -157,7 +157,7 @@ def test_build_config_no_display_name_collision_openai_and_azure():
"""OpenAI gpt-4o and azure gpt-4o should get distinct model_name values."""
env_vars = {
"OPENAI_API_KEY": "sk-openai",
"AZURE_API_KEY": "az-key",
"AZURE_AI_API_KEY": "az-key",
"_LITELLM_AZURE_DEPLOYMENT_AZURE": "gpt-4o",
}
config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, "sk-master")
@@ -182,7 +182,7 @@ def test_build_config_internal_sentinel_keys_excluded():
"""_LITELLM_ prefixed sentinel keys must not appear in environment_variables."""
env_vars = {
"OPENAI_API_KEY": "sk-real",
"_LITELLM_AZURE_API_BASE_AZURE": "https://x.azure.com",
"_LITELLM_AZURE_AI_API_BASE_AZURE": "https://x.azure.com",
}
config = SetupWizard._build_config([_OPENAI], env_vars, "sk-master")
assert "_LITELLM_" not in config