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

Merge pull request #24755 from BerriAI/litellm_test_cleanup

Browse Source 
Litellm test cleanup
This commit is contained in:
Krrish Dholakia
2026-03-28 22:04:26 -07:00
committed by GitHub
parent c27114470e 85cc7bc433
commit a671083596
92 changed files with 1855 additions and 2964 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent_tests/test_a2a_agent.py
+57 -25
View File
@@ -1,38 +1,72 @@
"""
Simple A2A agent tests - non-streaming and streaming.
These tests validate the localhost URL retry logic: if an A2A agent's card
contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request
will fail with a connection error. LiteLLM detects this and automatically
retries using the original api_base URL instead.
Requires A2A_AGENT_URL environment variable to be set.
Run with:
A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s
These tests use a mocked A2A client to avoid network/env dependencies.
"""
import os
import pytest
from types import SimpleNamespace
from uuid import uuid4
import pytest
def get_a2a_agent_url():
"""Get A2A agent URL from environment, skip test if not set."""
url = os.environ.get("A2A_AGENT_URL")
return url
class MockA2AResponse:
def __init__(self, text: str):
self._payload = {
"id": str(uuid4()),
"jsonrpc": "2.0",
"result": {
"message": {
"role": "agent",
"parts": [{"kind": "text", "text": text}],
"messageId": uuid4().hex,
}
},
}
def model_dump(self, mode="json", exclude_none=True):
return self._payload
class MockA2AStreamingChunk(MockA2AResponse):
def __init__(self, text: str, state: str):
super().__init__(text=text)
self._payload["result"]["status"] = {"state": state}
class MockA2AClient:
def __init__(self):
self._litellm_agent_card = SimpleNamespace(
name="mock-agent", url="http://mock-agent.local"
)
async def send_message(self, request):
return MockA2AResponse(text="hello")
def send_message_streaming(self, request):
async def _stream():
yield MockA2AStreamingChunk(text="hel", state="in_progress")
yield MockA2AStreamingChunk(text="hello", state="completed")
return _stream()
@pytest.fixture
def mock_a2a_client(monkeypatch):
import litellm.a2a_protocol.main as a2a_main
async def _fake_create_a2a_client(base_url, timeout=60.0, extra_headers=None):
return MockA2AClient()
monkeypatch.setattr(a2a_main, "create_a2a_client", _fake_create_a2a_client)
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=5)
async def test_a2a_non_streaming():
async def test_a2a_non_streaming(mock_a2a_client):
"""Test non-streaming A2A request."""
from a2a.types import MessageSendParams, SendMessageRequest
from litellm.a2a_protocol import asend_message
api_base = get_a2a_agent_url()
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
@@ -46,7 +80,7 @@ async def test_a2a_non_streaming():
response = await asend_message(
request=request,
api_base=api_base,
api_base="http://mock",
)
assert response is not None
@@ -54,13 +88,11 @@ async def test_a2a_non_streaming():
@pytest.mark.asyncio
async def test_a2a_streaming():
async def test_a2a_streaming(mock_a2a_client):
"""Test streaming A2A request."""
from a2a.types import MessageSendParams, SendStreamingMessageRequest
from litellm.a2a_protocol import asend_message_streaming
api_base = get_a2a_agent_url()
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
@@ -75,7 +107,7 @@ async def test_a2a_streaming():
chunks = []
async for chunk in asend_message_streaming(
request=request,
api_base=api_base,
api_base="http://mock",
):
chunks.append(chunk)
print(f"\nStreaming chunk: {chunk}")
tests/audio_tests/test_audio_speech.py
+1 -1
View File
@@ -286,7 +286,7 @@ async def test_speech_litellm_vertex_async_with_voice_ssml():
def test_audio_speech_cost_calc():
from litellm.integrations.custom_logger import CustomLogger
model = "azure/azure-tts"
model = "azure/tts"
api_base = os.getenv("AZURE_TTS_API_BASE")
api_key = os.getenv("AZURE_TTS_API_KEY")
tests/batches_tests/test_fine_tuning_api.py
-72
View File
@@ -123,78 +123,6 @@ async def test_create_fine_tune_jobs_async():
pass
@pytest.mark.asyncio
async def test_azure_create_fine_tune_jobs_async():
try:
verbose_logger.setLevel(logging.DEBUG)
file_name = "azure_fine_tune.jsonl"
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
file_id = "file-5e4b20ecbd724182b9964f3cd2ab7212"
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
model="gpt-35-turbo-1106",
training_file=file_id,
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
)
print(
"response from litellm.create_fine_tuning_job=", create_fine_tuning_response
)
assert create_fine_tuning_response.id is not None
# response from Example/mocked endpoint
assert create_fine_tuning_response.model == "davinci-002"
# list fine tuning jobs
print("listing ft jobs")
ft_jobs = await litellm.alist_fine_tuning_jobs(
limit=2,
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
)
print("response from litellm.list_fine_tuning_jobs=", ft_jobs)
# cancel ft job
response = await litellm.acancel_fine_tuning_job(
fine_tuning_job_id=create_fine_tuning_response.id,
custom_llm_provider="azure",
api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
api_base="https://exampleopenaiendpoint-production.up.railway.app",
)
print("response from litellm.cancel_fine_tuning_job=", response)
assert response.status == "cancelled"
assert response.id == create_fine_tuning_response.id
except openai.RateLimitError:
pass
except Exception as e:
if "Job has already completed" in str(e):
pass
else:
pytest.fail(f"Error occurred: {e}")
pass
def test_azure_trainingtype_defaults_to_one():
"""
Azure requires trainingType in extra_body. When omitted, AzureOpenAIFineTuningAPI defaults it to 1.
"""
from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI
handler = AzureOpenAIFineTuningAPI()
create_data = {"model": "gpt-4o-mini", "training_file": "file-test"}
handler._ensure_training_type(create_data)
assert "extra_body" in create_data
assert create_data["extra_body"]["trainingType"] == 1
@pytest.mark.asyncio()
async def test_create_vertex_fine_tune_jobs_mocked():
load_vertex_ai_credentials()
tests/batches_tests/test_openai_batches_and_files.py
+27 -23
View File
@@ -75,8 +75,8 @@ def load_vertex_ai_credentials():
service_account_key_data = {}
# Update the service_account_key_data with environment variables
private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "")
private_key = os.environ.get("GCS_PRIVATE_KEY", "")
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
private_key = private_key.replace("\\n", "\n")
service_account_key_data["private_key_id"] = private_key_id
service_account_key_data["private_key"] = private_key
@@ -234,9 +234,9 @@ def cleanup_azure_ft_models():
import requests
client = AzureOpenAI(
api_key=os.getenv("AZURE_FT_API_KEY"),
azure_endpoint=os.getenv("AZURE_FT_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
api_key=os.getenv("AZURE_AI_API_KEY"),
azure_endpoint=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_AI_API_VERSION"),
)
_list_ft_jobs = client.fine_tuning.jobs.list()
@@ -577,7 +577,10 @@ async def test_vertex_list_batches(monkeypatch):
monkeypatch.setattr(
"litellm.llms.vertex_ai.batches.handler.VertexAIBatchPrediction._ensure_access_token",
lambda self, credentials, project_id, custom_llm_provider: ("mock-token", "litellm-test-project"),
lambda self, credentials, project_id, custom_llm_provider: (
"mock-token",
"litellm-test-project",
),
)
with patch(
@@ -648,7 +651,7 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error():
async def test_delete_batch_output_file():
"""
Test that deleting a batch output file works correctly.
This test verifies the fix for:
- When a batch is retrieved and has an output_file_id, the file object is properly stored
- The output file can be deleted without validation errors
@@ -656,11 +659,11 @@ async def test_delete_batch_output_file():
"""
litellm._turn_on_debug()
print("Testing delete batch output file")
file_name = "openai_batch_completions.jsonl"
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
# Create file for batch
file_obj = await litellm.acreate_file(
file=open(file_path, "rb"),
@@ -669,7 +672,7 @@ async def test_delete_batch_output_file():
)
print("Response from creating file=", file_obj)
batch_input_file_id = file_obj.id
# Create batch
create_batch_response = await litellm.acreate_batch(
completion_window="24h",
@@ -678,36 +681,37 @@ async def test_delete_batch_output_file():
custom_llm_provider="openai",
)
print("Batch created with ID=", create_batch_response.id)
# Retrieve batch to get output_file_id
retrieved_batch = await litellm.aretrieve_batch(
batch_id=create_batch_response.id,
custom_llm_provider="openai"
batch_id=create_batch_response.id, custom_llm_provider="openai"
)
print("Retrieved batch=", retrieved_batch)
# If batch has completed and has output file, test deleting it
if retrieved_batch.output_file_id:
print(f"Testing deletion of output file: {retrieved_batch.output_file_id}")
# This is the key test - deleting the output file should work
# without validation errors (file_object should not be None)
delete_output_file_response = await litellm.afile_delete(
file_id=retrieved_batch.output_file_id,
custom_llm_provider="openai"
file_id=retrieved_batch.output_file_id, custom_llm_provider="openai"
)
print("Delete output file response=", delete_output_file_response)
assert delete_output_file_response.id == retrieved_batch.output_file_id
assert delete_output_file_response.deleted is True or hasattr(delete_output_file_response, 'id')
assert delete_output_file_response.deleted is True or hasattr(
delete_output_file_response, "id"
)
print("✓ Successfully deleted batch output file")
else:
print("⚠ Batch has not completed yet or no output file available, skipping output file deletion test")
print(
"⚠ Batch has not completed yet or no output file available, skipping output file deletion test"
)
# Clean up - delete the input file
delete_input_file_response = await litellm.afile_delete(
file_id=batch_input_file_id,
custom_llm_provider="openai"
file_id=batch_input_file_id, custom_llm_provider="openai"
)
print("Delete input file response=", delete_input_file_response)
assert delete_input_file_response.id == batch_input_file_id
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/image_gen_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/litellm_utils_tests/test_health_check.py
+16 -40
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",
@@ -244,35 +246,6 @@ async def test_audio_transcription_health_check():
print(response)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model", ["azure/gpt-4o-realtime-preview", "openai/gpt-4o-realtime-preview"]
)
async def test_async_realtime_health_check(model, mocker):
"""
Test Health Check with Valid models passes
"""
mock_websocket = AsyncMock()
mock_connect = AsyncMock().__aenter__.return_value = mock_websocket
mocker.patch("websockets.connect", return_value=mock_connect)
litellm.set_verbose = True
model_params = {
"model": model,
}
if model == "azure/gpt-4o-realtime-preview":
model_params["api_base"] = os.getenv("AZURE_REALTIME_API_BASE")
model_params["api_key"] = os.getenv("AZURE_REALTIME_API_KEY")
model_params["api_version"] = os.getenv("AZURE_REALTIME_API_VERSION")
response = await litellm.ahealth_check(
model_params=model_params,
mode="realtime",
)
print(response)
assert response == {}
def test_update_litellm_params_for_health_check():
"""
Test if _update_litellm_params_for_health_check correctly:
@@ -500,7 +473,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 +632,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/litellm_utils_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
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_anthropic_completion.py
+6
View File
@@ -1350,6 +1350,9 @@ def test_anthropic_text_editor():
@pytest.mark.parametrize("spec", ["anthropic", "openai"])
@pytest.mark.skipif(
os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set"
)
def test_anthropic_mcp_server_tool_use(spec: str):
litellm._turn_on_debug()
@@ -1391,6 +1394,9 @@ def test_anthropic_mcp_server_tool_use(spec: str):
@pytest.mark.parametrize(
"model", ["openai/gpt-4.1", "anthropic/claude-sonnet-4-20250514"]
)
@pytest.mark.skipif(
os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set"
)
def test_anthropic_mcp_server_responses_api(model: str):
from litellm import responses
tests/llm_translation/test_azure_ai.py
+70 -68
View File
@@ -188,35 +188,6 @@ def test_azure_ai_services_with_api_version():
)
@pytest.mark.skip(reason="Skipping due to cohere ssl issues")
def test_completion_azure_ai_command_r():
try:
import os
litellm.set_verbose = True
os.environ["AZURE_AI_API_BASE"] = os.getenv("AZURE_COHERE_API_BASE", "")
os.environ["AZURE_AI_API_KEY"] = os.getenv("AZURE_COHERE_API_KEY", "")
response = completion(
model="azure_ai/command-r-plus",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "What is the meaning of life?"}
],
}
],
) # type: ignore
assert "azure_ai" in response.model
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_azure_deepseek_reasoning_content():
import json
@@ -283,8 +254,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 +281,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 +302,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 +329,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,18 +345,20 @@ 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",
messages=[{"role": "user", "content": "hi who is this"}],
api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/",
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
)
print("response: ", response)
@@ -394,23 +367,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 +417,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 +443,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 +485,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 +502,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_bedrock_completion.py
+15 -60
View File
@@ -217,60 +217,6 @@ def test_completion_bedrock_claude_external_client_auth():
# test_completion_bedrock_claude_external_client_auth()
@pytest.mark.skip(reason="Expired token, need to renew")
def test_completion_bedrock_claude_sts_client_auth():
print("\ncalling bedrock claude external client auth")
import os
aws_access_key_id = os.environ["AWS_TEMP_ACCESS_KEY_ID"]
aws_secret_access_key = os.environ["AWS_TEMP_SECRET_ACCESS_KEY"]
aws_region_name = os.environ["AWS_REGION_NAME"]
aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"]
try:
import boto3
litellm.set_verbose = True
response = completion(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
messages=messages,
max_tokens=10,
temperature=0.1,
aws_region_name=aws_region_name,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
)
response = embedding(
model="cohere.embed-multilingual-v3",
input=["hello world"],
aws_region_name="us-east-1",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
)
response = completion(
model="gpt-3.5-turbo",
messages=messages,
aws_region_name="us-east-1",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
)
# Add any assertions here to check the response
print(response)
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.fixture()
def bedrock_session_token_creds():
print("\ncalling oidc auto to get aws_session_token credentials")
@@ -3413,7 +3359,8 @@ def test_bedrock_openai_imported_model():
print(f"URL: {url}")
assert "bedrock-runtime.us-east-1.amazonaws.com" in url
assert (
"arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy" in url
"arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy"
in url
)
assert "/invoke" in url
@@ -3850,10 +3797,12 @@ def test_bedrock_openai_error_handling():
assert exc_info.value.status_code == 422
print("✓ Error handling works correctly")
# ============================================================================
# Nova Grounding (web_search_options) Unit Tests (Mocked)
# ============================================================================
def test_bedrock_nova_grounding_web_search_options_non_streaming():
"""
Unit test for Nova grounding using web_search_options parameter (non-streaming).
@@ -3907,7 +3856,9 @@ def test_bedrock_nova_grounding_web_search_options_non_streaming():
break
assert system_tool_found, "systemTool with nova_grounding should be present"
print(f"✓ web_search_options correctly transformed to systemTool (non-streaming)")
print(
f"✓ web_search_options correctly transformed to systemTool (non-streaming)"
)
def test_bedrock_nova_grounding_with_function_tools():
@@ -3987,7 +3938,9 @@ def test_bedrock_nova_grounding_with_function_tools():
assert tool["systemTool"]["name"] == "nova_grounding"
system_tool_found = True
assert function_tool_found, "Function tool (get_stock_price) should be present"
assert (
function_tool_found
), "Function tool (get_stock_price) should be present"
assert system_tool_found, "systemTool (nova_grounding) should be present"
print(f"✓ Both function tools and web_search_options correctly combined")
@@ -4092,10 +4045,12 @@ def test_bedrock_nova_grounding_request_transformation():
mock_post.return_value = MagicMock(
status_code=200,
json=lambda: {
"output": {"message": {"role": "assistant", "content": [{"text": "Test"}]}},
"output": {
"message": {"role": "assistant", "content": [{"text": "Test"}]}
},
"stopReason": "end_turn",
"usage": {"inputTokens": 10, "outputTokens": 5}
}
"usage": {"inputTokens": 10, "outputTokens": 5},
},
)
try:
tests/llm_translation/test_clarifai_completion.py
-109
View File
@@ -1,109 +0,0 @@
import sys, os
import traceback
from dotenv import load_dotenv
import asyncio, logging
load_dotenv()
import os, io
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
import litellm
from litellm import (
embedding,
completion,
acompletion,
acreate,
completion_cost,
Timeout,
ModelResponse,
)
from litellm import RateLimitError
# litellm.num_retries = 3
litellm.cache = None
litellm.success_callback = []
user_message = "Write a short poem about the sky"
messages = [{"content": user_message, "role": "user"}]
@pytest.fixture(autouse=True)
def reset_callbacks():
print("\npytest fixture - resetting callbacks")
litellm.success_callback = []
litellm._async_success_callback = []
litellm.failure_callback = []
litellm.callbacks = []
@pytest.mark.skip(reason="Account rate limited.")
def test_completion_clarifai_claude_2_1():
print("calling clarifai claude completion")
import os
clarifai_pat = os.environ["CLARIFAI_API_KEY"]
try:
response = completion(
model="clarifai/anthropic.completion.claude-2_1",
num_retries=3,
messages=messages,
max_tokens=10,
temperature=0.1,
)
print(response)
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occured: {e}")
@pytest.mark.skip(reason="Account rate limited")
def test_completion_clarifai_mistral_large():
try:
litellm.set_verbose = True
response: ModelResponse = completion(
model="clarifai/mistralai.completion.mistral-small",
messages=messages,
num_retries=3,
max_tokens=10,
temperature=0.78,
)
# Add any assertions here to check the response
assert len(response.choices) > 0
assert len(response.choices[0].message.content) > 0
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Account rate limited")
@pytest.mark.asyncio
def test_async_completion_clarifai():
import asyncio
litellm.set_verbose = True
async def test_get_response():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(
model="clarifai/openai.chat-completion.GPT-4",
messages=messages,
num_retries=3,
timeout=10,
api_key=os.getenv("CLARIFAI_API_KEY"),
)
print(f"response: {response}")
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
asyncio.run(test_get_response())
tests/llm_translation/test_rerank.py
-42
View File
@@ -148,48 +148,6 @@ async def test_basic_rerank_together_ai(sync_mode):
raise e
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.skip(reason="Skipping test due to Cohere RBAC issues")
async def test_basic_rerank_azure_ai(sync_mode):
import os
litellm.set_verbose = True
if sync_mode is True:
response = litellm.rerank(
model="azure_ai/Cohere-rerank-v3-multilingual-ko",
query="hello",
documents=["hello", "world"],
top_n=3,
api_key=os.getenv("AZURE_AI_COHERE_API_KEY"),
api_base=os.getenv("AZURE_AI_COHERE_API_BASE"),
)
print("re rank response: ", response)
assert response.id is not None
assert response.results is not None
assert_response_shape(response, custom_llm_provider="together_ai")
else:
response = await litellm.arerank(
model="azure_ai/Cohere-rerank-v3-multilingual-ko",
query="hello",
documents=["hello", "world"],
top_n=3,
api_key=os.getenv("AZURE_AI_COHERE_API_KEY"),
api_base=os.getenv("AZURE_AI_COHERE_API_BASE"),
)
print("async re rank response: ", response)
assert response.id is not None
assert response.results is not None
assert_response_shape(response, custom_llm_provider="together_ai")
@pytest.mark.asyncio()
@pytest.mark.parametrize("version", ["v1", "v2"])
async def test_rerank_custom_api_base(version):
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/llm_translation/test_snowflake.py
+58 -118
View File
@@ -46,45 +46,51 @@ def mock_snowflake_streaming_response_chunks() -> List[str]:
Mock streaming response chunks for Snowflake.
"""
return [
json.dumps({
"id": "chatcmpl-snowflake-stream-123",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-7b",
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "The"},
"finish_reason": None,
}
],
}),
json.dumps({
"id": "chatcmpl-snowflake-stream-123",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-7b",
"choices": [
{
"index": 0,
"delta": {"content": " sky"},
"finish_reason": None,
}
],
}),
json.dumps({
"id": "chatcmpl-snowflake-stream-123",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-7b",
"choices": [
{
"index": 0,
"delta": {"content": " is blue"},
"finish_reason": "stop",
}
],
}),
json.dumps(
{
"id": "chatcmpl-snowflake-stream-123",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-7b",
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "The"},
"finish_reason": None,
}
],
}
),
json.dumps(
{
"id": "chatcmpl-snowflake-stream-123",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-7b",
"choices": [
{
"index": 0,
"delta": {"content": " sky"},
"finish_reason": None,
}
],
}
),
json.dumps(
{
"id": "chatcmpl-snowflake-stream-123",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "mistral-7b",
"choices": [
{
"index": 0,
"delta": {"content": " is blue"},
"finish_reason": "stop",
}
],
}
),
]
@@ -120,6 +126,7 @@ def test_chat_completion_snowflake(sync_mode):
async_handler = AsyncHTTPHandler()
with patch.object(AsyncHTTPHandler, "post", return_value=mock_response):
import asyncio
response = asyncio.run(
acompletion(
model="snowflake/mistral-7b",
@@ -148,16 +155,16 @@ def test_chat_completion_snowflake_stream(sync_mode):
if sync_mode:
sync_handler = HTTPHandler()
mock_chunks = mock_snowflake_streaming_response_chunks()
def mock_iter_lines():
for chunk in mock_chunks:
for line in [f"data: {chunk}", "data: [DONE]"]:
yield line
mock_response = MagicMock()
mock_response.iter_lines.side_effect = mock_iter_lines
mock_response.status_code = 200
with patch.object(HTTPHandler, "post", return_value=mock_response):
response = completion(
model="snowflake/mistral-7b",
@@ -167,28 +174,28 @@ def test_chat_completion_snowflake_stream(sync_mode):
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
client=sync_handler,
)
chunks_received = []
for chunk in response:
chunks_received.append(chunk)
assert len(chunks_received) > 0
else:
async_handler = AsyncHTTPHandler()
mock_chunks = mock_snowflake_streaming_response_chunks()
async def mock_iter_lines():
for chunk in mock_chunks:
for line in [f"data: {chunk}", "data: [DONE]"]:
yield line
mock_response = MagicMock()
mock_response.iter_lines.side_effect = mock_iter_lines
mock_response.status_code = 200
with patch.object(AsyncHTTPHandler, "post", return_value=mock_response):
import asyncio
async def test_async_stream():
response = await acompletion(
model="snowflake/mistral-7b",
@@ -198,78 +205,11 @@ def test_chat_completion_snowflake_stream(sync_mode):
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
client=async_handler,
)
chunks_received = []
async for chunk in response:
chunks_received.append(chunk)
assert len(chunks_received) > 0
asyncio.run(test_async_stream())
@pytest.mark.skip(reason="Requires Snowflake credentials - run manually when needed")
def test_snowflake_tool_calling_responses_api():
"""
Test Snowflake tool calling with Responses API.
Requires SNOWFLAKE_JWT and SNOWFLAKE_ACCOUNT_ID environment variables.
"""
import litellm
# Skip if credentials not available
if not os.getenv("SNOWFLAKE_JWT") or not os.getenv("SNOWFLAKE_ACCOUNT_ID"):
pytest.skip("Snowflake credentials not available")
litellm.drop_params = False # We now support tools!
tools = [
{
"type": "function",
"name": "get_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
}
},
"required": ["location"],
},
}
]
try:
# Test with tool_choice to force tool use
response = responses(
model="snowflake/claude-3-5-sonnet",
input="What's the weather in Paris?",
tools=tools,
tool_choice={"type": "function", "function": {"name": "get_weather"}},
max_output_tokens=200,
)
assert response is not None
assert hasattr(response, "output")
assert len(response.output) > 0
# Verify tool call was made
tool_call_found = False
for item in response.output:
if hasattr(item, "type") and item.type == "function_call":
tool_call_found = True
assert item.name == "get_weather"
assert hasattr(item, "arguments")
print(f"✅ Tool call detected: {item.name}({item.arguments})")
break
assert tool_call_found, "Expected tool call but none was found"
except APIConnectionError as e:
if "JWT token is invalid" in str(e):
pytest.skip("Invalid Snowflake JWT token")
elif "Application failed to respond" in str(e) or "502" in str(e):
pytest.skip(f"Snowflake API unavailable: {e}")
else:
raise
tests/load_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
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_alangfuse.py
-263
View File
@@ -599,233 +599,6 @@ def test_langfuse_logging_function_calling():
# test_langfuse_logging_function_calling()
@pytest.mark.skip(reason="Need to address this on main")
def test_aaalangfuse_existing_trace_id():
"""
When existing trace id is passed, don't set trace params -> prevents overwriting the trace
Pass 1 logging object with a trace
Pass 2nd logging object with the trace id
Assert no changes to the trace
"""
# Test - if the logs were sent to the correct team on langfuse
import datetime
import litellm
from litellm.integrations.langfuse.langfuse import LangFuseLogger
langfuse_Logger = LangFuseLogger(
langfuse_public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
langfuse_secret=os.getenv("LANGFUSE_PROJECT2_SECRET"),
)
litellm.success_callback = ["langfuse"]
# langfuse_args = {'kwargs': { 'start_time': 'end_time': datetime.datetime(2024, 5, 1, 7, 31, 29, 903685), 'user_id': None, 'print_verbose': <function print_verbose at 0x109d1f420>, 'level': 'DEFAULT', 'status_message': None}
response_obj = litellm.ModelResponse(
id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY",
choices=[
litellm.Choices(
finish_reason="stop",
index=0,
message=litellm.Message(
content="I'm sorry, I am an AI assistant and do not have real-time information. I recommend checking a reliable weather website or app for the most up-to-date weather information in Boston.",
role="assistant",
),
)
],
created=1714573888,
model="gpt-3.5-turbo-0125",
object="chat.completion",
system_fingerprint="fp_3b956da36b",
usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51),
)
### NEW TRACE ###
message = [{"role": "user", "content": "what's the weather in boston"}]
langfuse_args = {
"response_obj": response_obj,
"kwargs": {
"model": "gpt-3.5-turbo",
"litellm_params": {
"acompletion": False,
"api_key": None,
"force_timeout": 600,
"logger_fn": None,
"verbose": False,
"custom_llm_provider": "openai",
"api_base": "https://api.openai.com/v1/",
"litellm_call_id": None,
"model_alias_map": {},
"completion_call_id": None,
"metadata": None,
"model_info": None,
"proxy_server_request": None,
"preset_cache_key": None,
"no-log": False,
"stream_response": {},
},
"messages": message,
"optional_params": {"temperature": 0.1, "extra_body": {}},
"start_time": "2024-05-01 07:31:27.986164",
"stream": False,
"user": None,
"call_type": "completion",
"litellm_call_id": None,
"completion_start_time": "2024-05-01 07:31:29.903685",
"temperature": 0.1,
"extra_body": {},
"input": [{"role": "user", "content": "what's the weather in boston"}],
"api_key": "my-api-key",
"additional_args": {
"complete_input_dict": {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "what's the weather in boston"}
],
"temperature": 0.1,
"extra_body": {},
}
},
"log_event_type": "successful_api_call",
"end_time": "2024-05-01 07:31:29.903685",
"cache_hit": None,
"response_cost": 6.25e-05,
},
"start_time": datetime.datetime(2024, 5, 1, 7, 31, 27, 986164),
"end_time": datetime.datetime(2024, 5, 1, 7, 31, 29, 903685),
"user_id": None,
"print_verbose": litellm.print_verbose,
"level": "DEFAULT",
"status_message": None,
}
langfuse_response_object = langfuse_Logger.log_event(**langfuse_args)
import langfuse
langfuse_client = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
)
trace_id = langfuse_response_object["trace_id"]
assert trace_id is not None
langfuse_client.flush()
time.sleep(2)
print(langfuse_client.get_trace(id=trace_id))
initial_langfuse_trace = langfuse_client.get_trace(id=trace_id)
### EXISTING TRACE ###
new_metadata = {"existing_trace_id": trace_id}
new_messages = [{"role": "user", "content": "What do you know?"}]
new_response_obj = litellm.ModelResponse(
id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY",
choices=[
litellm.Choices(
finish_reason="stop",
index=0,
message=litellm.Message(
content="What do I know?",
role="assistant",
),
)
],
created=1714573888,
model="gpt-3.5-turbo-0125",
object="chat.completion",
system_fingerprint="fp_3b956da36b",
usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51),
)
langfuse_args = {
"response_obj": new_response_obj,
"kwargs": {
"model": "gpt-3.5-turbo",
"litellm_params": {
"acompletion": False,
"api_key": None,
"force_timeout": 600,
"logger_fn": None,
"verbose": False,
"custom_llm_provider": "openai",
"api_base": "https://api.openai.com/v1/",
"litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e",
"model_alias_map": {},
"completion_call_id": None,
"metadata": new_metadata,
"model_info": None,
"proxy_server_request": None,
"preset_cache_key": None,
"no-log": False,
"stream_response": {},
},
"messages": new_messages,
"optional_params": {"temperature": 0.1, "extra_body": {}},
"start_time": "2024-05-01 07:31:27.986164",
"stream": False,
"user": None,
"call_type": "completion",
"litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e",
"completion_start_time": "2024-05-01 07:31:29.903685",
"temperature": 0.1,
"extra_body": {},
"input": [{"role": "user", "content": "what's the weather in boston"}],
"api_key": "my-api-key",
"additional_args": {
"complete_input_dict": {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "what's the weather in boston"}
],
"temperature": 0.1,
"extra_body": {},
}
},
"log_event_type": "successful_api_call",
"end_time": "2024-05-01 07:31:29.903685",
"cache_hit": None,
"response_cost": 6.25e-05,
},
"start_time": datetime.datetime(2024, 5, 1, 7, 31, 27, 986164),
"end_time": datetime.datetime(2024, 5, 1, 7, 31, 29, 903685),
"user_id": None,
"print_verbose": litellm.print_verbose,
"level": "DEFAULT",
"status_message": None,
}
langfuse_response_object = langfuse_Logger.log_event(**langfuse_args)
new_trace_id = langfuse_response_object["trace_id"]
assert new_trace_id == trace_id
langfuse_client.flush()
time.sleep(2)
print(langfuse_client.get_trace(id=trace_id))
new_langfuse_trace = langfuse_client.get_trace(id=trace_id)
initial_langfuse_trace_dict = dict(initial_langfuse_trace)
initial_langfuse_trace_dict.pop("updatedAt")
initial_langfuse_trace_dict.pop("timestamp")
new_langfuse_trace_dict = dict(new_langfuse_trace)
new_langfuse_trace_dict.pop("updatedAt")
new_langfuse_trace_dict.pop("timestamp")
assert initial_langfuse_trace_dict == new_langfuse_trace_dict
@pytest.mark.skipif(
condition=not os.environ.get("OPENAI_API_KEY", False),
reason="Authentication missing for openai",
@@ -928,42 +701,6 @@ async def test_make_request():
)
@pytest.mark.skip(
reason="local only test, use this to verify if dynamic langfuse logging works as expected"
)
def test_aaalangfuse_dynamic_logging():
"""
pass in langfuse credentials via completion call
assert call is logged.
Covers the team-logging scenario.
"""
from litellm._uuid import uuid
import langfuse
trace_id = str(uuid.uuid4())
_ = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey"}],
mock_response="Hey! how's it going?",
langfuse_public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
langfuse_secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
metadata={"trace_id": trace_id},
success_callback=["langfuse"],
)
time.sleep(3)
langfuse_client = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
)
langfuse_client.get_trace(id=trace_id)
import datetime
generation_params = {
tests/local_testing/test_arize_ai.py
+2 -2
View File
@@ -48,8 +48,8 @@ async def test_async_dynamic_arize_config():
messages=[{"role": "user", "content": "hi test from arize dynamic config"}],
temperature=0.1,
user="OTEL_USER",
arize_api_key=os.getenv("ARIZE_SPACE_2_API_KEY"),
arize_space_key=os.getenv("ARIZE_SPACE_2_KEY"),
arize_api_key=os.getenv("ARIZE_SPACE_API_KEY"),
arize_space_key=os.getenv("ARIZE_SPACE_KEY"),
)
await asyncio.sleep(2)
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 -138
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
@@ -177,34 +176,6 @@ def predibase_mock_post(url, data=None, json=None, headers=None, timeout=None):
return mock_response
# @pytest.mark.skip(reason="local-only test")
@pytest.mark.asyncio
async def test_completion_predibase():
try:
litellm.set_verbose = True
# with patch("requests.post", side_effect=predibase_mock_post):
response = await litellm.acompletion(
model="predibase/llama-3-8b-instruct",
tenant_id="c4768f95",
api_key=os.getenv("PREDIBASE_API_KEY"),
messages=[{"role": "user", "content": "who are u?"}],
max_tokens=10,
timeout=5,
)
print(response)
except litellm.Timeout as e:
print("got a timeout error from predibase")
pass
except litellm.ServiceUnavailableError as e:
pass
except litellm.InternalServerError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_predibase()
@@ -286,7 +257,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 +822,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 +969,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 +1671,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 +2260,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 +2298,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 +2315,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 +2331,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 +2349,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 +2362,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 +2444,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 +2465,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 +2481,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 +2499,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 +2513,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 +2523,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 +2546,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 +3000,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 +3174,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:
@@ -3293,23 +3210,6 @@ def test_mistral_anyscale_stream():
print(chunk["choices"][0]["delta"].get("content", ""), end="")
# test_completion_anyscale_2()
# def test_completion_with_fallbacks_multiple_keys():
# print(f"backup key 1: {os.getenv('BACKUP_OPENAI_API_KEY_1')}")
# print(f"backup key 2: {os.getenv('BACKUP_OPENAI_API_KEY_2')}")
# backup_keys = [{"api_key": os.getenv("BACKUP_OPENAI_API_KEY_1")}, {"api_key": os.getenv("BACKUP_OPENAI_API_KEY_2")}]
# try:
# api_key = "bad-key"
# response = completion(
# model="gpt-3.5-turbo", messages=messages, force_timeout=120, fallbacks=backup_keys, api_key=api_key
# )
# # Add any assertions here to check the response
# print(response)
# except Exception as e:
# error_str = traceback.format_exc()
# pytest.fail(f"Error occurred: {error_str}")
# test_completion_with_fallbacks_multiple_keys()
def test_petals():
try:
@@ -3871,9 +3771,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
+79 -67
View File
@@ -108,7 +108,11 @@ def test_openai_embedding_3():
"model, api_base, api_key",
[
# ("azure/text-embedding-ada-002", None, None),
("together_ai/BAAI/bge-base-en-v1.5", None, None), # Updated to current Together AI embedding model
(
"together_ai/BAAI/bge-base-en-v1.5",
None,
None,
), # Updated to current Together AI embedding model
],
)
@pytest.mark.parametrize("sync_mode", [True, False])
@@ -193,9 +197,9 @@ def _azure_ai_image_mock_response(*args, **kwargs):
"model, api_base, api_key",
[
(
"azure_ai/Cohere-embed-v3-multilingual-jzu",
"https://Cohere-embed-v3-multilingual-jzu.eastus2.models.ai.azure.com",
os.getenv("AZURE_AI_COHERE_API_KEY_2"),
"azure_ai/Cohere-embed-v3-multilingual-2",
os.getenv("AZURE_AI_API_BASE"),
os.getenv("AZURE_AI_API_KEY"),
)
],
)
@@ -292,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",
@@ -310,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}")
@@ -353,11 +357,11 @@ def test_openai_azure_embedding_optional_arg():
)
mock_client.assert_called_once_with(
model="test",
input=["test"],
extra_body={"azure_ad_token": "test"},
timeout=600,
extra_headers={"X-Stainless-Raw-Response": "true"}
model="test",
input=["test"],
extra_body={"azure_ad_token": "test"},
timeout=600,
extra_headers={"X-Stainless-Raw-Response": "true"},
)
# Verify azure_ad_token is passed in extra_body, not as a direct parameter
assert "azure_ad_token" not in mock_client.call_args.kwargs
@@ -545,7 +549,7 @@ def test_bedrock_embedding_cohere():
"good morning from litellm, attempting to embed data",
"lets test a second string for good measure",
],
aws_region_name="os.environ/AWS_REGION_NAME_2",
aws_region_name="os.environ/AWS_REGION_NAME",
)
assert isinstance(
response["data"][0]["embedding"], list
@@ -811,7 +815,7 @@ def test_watsonx_embeddings(monkeypatch):
monkeypatch.setenv("WATSONX_PROJECT_ID", "mock-project-id")
client = HTTPHandler()
# Track the actual request made
captured_request = {}
@@ -820,7 +824,7 @@ def test_watsonx_embeddings(monkeypatch):
captured_request["url"] = url
captured_request["headers"] = kwargs.get("headers", {})
captured_request["data"] = kwargs.get("data")
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
@@ -843,10 +847,12 @@ def test_watsonx_embeddings(monkeypatch):
print(f"response: {response}")
assert isinstance(response.usage, litellm.Usage)
# Verify the request was made correctly
assert "Authorization" in captured_request["headers"]
assert captured_request["headers"]["Authorization"] == "Bearer mock-watsonx-token"
assert (
captured_request["headers"]["Authorization"] == "Bearer mock-watsonx-token"
)
assert "us-south.ml.cloud.ibm.com" in captured_request["url"]
except litellm.RateLimitError as e:
pass
@@ -1256,9 +1262,7 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input):
# Call the function we want to test
try:
litellm.embedding(
model="jina_ai/jina-embeddings-v4", input=input_data
)
litellm.embedding(model="jina_ai/jina-embeddings-v4", input=input_data)
except Exception as e:
pytest.fail(
f"litellm.embedding call failed with an unexpected exception: {e}"
@@ -1285,105 +1289,113 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input):
def test_encoding_format_none_not_omitted_from_openai_sdk():
"""
Test that encoding_format=None is explicitly sent to OpenAI SDK.
This test verifies that when encoding_format is not provided by the user,
liteLLM explicitly sets it to None rather than omitting it. This prevents
the OpenAI SDK from adding its default value of 'base64'.
Without this fix:
- OpenAI SDK adds encoding_format='base64' as default when parameter is missing
- This causes issues with providers that don't support encoding_format (like Gemini)
With this fix:
- encoding_format=None is explicitly passed
- OpenAI SDK respects the explicit None and doesn't add defaults
"""
with patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") as mock_get_client:
with patch(
"litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client"
) as mock_get_client:
# Create a mock client instance
mock_client_instance = MagicMock()
mock_get_client.return_value = mock_client_instance
# Mock the embeddings.with_raw_response.create method
mock_response = MagicMock()
mock_response.parse.return_value = MagicMock(
model_dump=lambda: {
'data': [{'embedding': [0.1, 0.2, 0.3], 'index': 0}],
'model': 'text-embedding-ada-002',
'object': 'list',
'usage': {'prompt_tokens': 1, 'total_tokens': 1}
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
"model": "text-embedding-ada-002",
"object": "list",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
}
)
mock_response.headers = {}
mock_client_instance.embeddings.with_raw_response.create.return_value = mock_response
mock_client_instance.embeddings.with_raw_response.create.return_value = (
mock_response
)
# Call the embedding function without encoding_format
response = embedding(
model="text-embedding-ada-002",
input="Hello world",
)
# Get the call arguments to verify what was sent to OpenAI SDK
call_args = mock_client_instance.embeddings.with_raw_response.create.call_args
assert call_args is not None, "OpenAI SDK embeddings.create should have been called"
assert (
call_args is not None
), "OpenAI SDK embeddings.create should have been called"
call_kwargs = call_args[1] # Get kwargs
# The key assertion: encoding_format should be in the request with value None
# This prevents OpenAI SDK from adding its default 'base64' value
assert 'encoding_format' in call_kwargs, (
assert "encoding_format" in call_kwargs, (
"encoding_format should be explicitly passed to OpenAI SDK "
"(even if None) to prevent SDK from adding default value"
)
assert call_kwargs['encoding_format'] is None, (
"encoding_format should be None when not provided by user"
)
assert (
call_kwargs["encoding_format"] is None
), "encoding_format should be None when not provided by user"
print("✅ PASS: encoding_format=None is correctly passed to OpenAI SDK")
def test_encoding_format_explicit_value_preserved():
"""
Test that explicitly provided encoding_format values are preserved.
When user provides encoding_format='float' or 'base64', it should be
When user provides encoding_format='float' or 'base64', it should be
sent as-is to the OpenAI SDK.
"""
with patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") as mock_get_client:
with patch(
"litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client"
) as mock_get_client:
# Create a mock client instance
mock_client_instance = MagicMock()
mock_get_client.return_value = mock_client_instance
# Mock the embeddings.with_raw_response.create method
mock_response = MagicMock()
mock_response.parse.return_value = MagicMock(
model_dump=lambda: {
'data': [{'embedding': [0.1, 0.2, 0.3], 'index': 0}],
'model': 'text-embedding-ada-002',
'object': 'list',
'usage': {'prompt_tokens': 1, 'total_tokens': 1}
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
"model": "text-embedding-ada-002",
"object": "list",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
}
)
mock_response.headers = {}
mock_client_instance.embeddings.with_raw_response.create.return_value = mock_response
mock_client_instance.embeddings.with_raw_response.create.return_value = (
mock_response
)
# Test with explicit encoding_format='float'
response = embedding(
model="text-embedding-ada-002",
input="Hello world",
encoding_format="float"
model="text-embedding-ada-002", input="Hello world", encoding_format="float"
)
# Verify the encoding_format was passed correctly
call_args = mock_client_instance.embeddings.with_raw_response.create.call_args
call_kwargs = call_args[1]
assert 'encoding_format' in call_kwargs, (
"encoding_format should be in the request"
)
assert call_kwargs['encoding_format'] == 'float', (
"encoding_format should be 'float' when explicitly provided"
)
assert (
"encoding_format" in call_kwargs
), "encoding_format should be in the request"
assert (
call_kwargs["encoding_format"] == "float"
), "encoding_format should be 'float' when explicitly provided"
print("✅ PASS: encoding_format='float' is correctly preserved")
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_gcs_bucket.py
+43 -8
View File
@@ -22,6 +22,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket import (
)
from litellm.types.utils import StandardCallbackDynamicParams
from unittest.mock import patch
verbose_logger.setLevel(logging.DEBUG)
@@ -52,8 +53,8 @@ def load_vertex_ai_credentials():
service_account_key_data = {}
# Update the service_account_key_data with environment variables
private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "")
private_key = os.environ.get("GCS_PRIVATE_KEY", "")
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
private_key = private_key.replace("\\n", "\n")
service_account_key_data["private_key_id"] = private_key_id
service_account_key_data["private_key"] = private_key
@@ -686,6 +687,7 @@ async def test_basic_gcs_logger_with_folder_in_bucket_name():
if old_bucket_name is not None:
os.environ["GCS_BUCKET_NAME"] = old_bucket_name
@pytest.mark.skip(reason="This test is flaky on ci/cd")
def test_create_file_e2e():
"""
@@ -696,6 +698,7 @@ def test_create_file_e2e():
test_file = ("test.wav", test_file_content, "audio/wav")
from litellm import create_file
response = create_file(
file=test_file,
purpose="user_data",
@@ -704,6 +707,7 @@ def test_create_file_e2e():
print("response", response)
assert response is not None
@pytest.mark.skip(reason="This test is flaky on ci/cd")
def test_create_file_e2e_jsonl():
"""
@@ -714,14 +718,41 @@ def test_create_file_e2e_jsonl():
client = HTTPHandler()
example_jsonl = [{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}},{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}]
example_jsonl = [
{
"custom_id": "request-1",
"method": "POST",
"url": "/v1/chat/completions",
"body": {
"model": "gemini-1.5-flash-001",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello world!"},
],
"max_tokens": 10,
},
},
{
"custom_id": "request-2",
"method": "POST",
"url": "/v1/chat/completions",
"body": {
"model": "gemini-1.5-flash-001",
"messages": [
{"role": "system", "content": "You are an unhelpful assistant."},
{"role": "user", "content": "Hello world!"},
],
"max_tokens": 10,
},
},
]
# Create and write to the file
file_path = "example.jsonl"
with open(file_path, "w") as f:
for item in example_jsonl:
f.write(json.dumps(item) + "\n")
# Verify file content
with open(file_path, "r") as f:
content = f.read()
@@ -729,10 +760,11 @@ def test_create_file_e2e_jsonl():
assert len(content) > 0, "File is empty"
from litellm import create_file
with patch.object(client, "post") as mock_create_file:
try:
try:
response = create_file(
file=open(file_path, "rb"),
file=open(file_path, "rb"),
purpose="user_data",
custom_llm_provider="vertex_ai",
client=client,
@@ -744,4 +776,7 @@ def test_create_file_e2e_jsonl():
print(f"kwargs: {mock_create_file.call_args.kwargs}")
assert mock_create_file.call_args.kwargs["data"] is not None and len(mock_create_file.call_args.kwargs["data"]) > 0
assert (
mock_create_file.call_args.kwargs["data"] is not None
and len(mock_create_file.call_args.kwargs["data"]) > 0
)
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.py
+50 -92
View File
@@ -126,7 +126,9 @@ async def test_router_provider_wildcard_routing():
print("response 3 = ", response3)
response4 = await router.acompletion(
model=os.environ.get("CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
model=os.environ.get(
"CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"
),
messages=[{"role": "user", "content": "hello"}],
)
@@ -356,51 +358,6 @@ async def test_router_retries(sync_mode):
print(response.choices[0].message)
@pytest.mark.parametrize(
"mistral_api_base",
[
"os.environ/AZURE_MISTRAL_API_BASE",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/v1/",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/v1",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com",
],
)
@pytest.mark.skip(
reason="Router no longer creates clients, this is delegated to the provider integration."
)
def test_router_azure_ai_studio_init(mistral_api_base):
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "azure/mistral-large-latest",
"api_key": "os.environ/AZURE_MISTRAL_API_KEY",
"api_base": mistral_api_base,
},
"model_info": {"id": 1234},
}
]
)
# model_client = router._get_client(
# deployment={"model_info": {"id": 1234}}, client_type="sync_client", kwargs={}
# )
# url = getattr(model_client, "_base_url")
# uri_reference = str(getattr(url, "_uri_reference"))
# print(f"uri_reference: {uri_reference}")
# assert "/v1/" in uri_reference
# assert uri_reference.count("v1") == 1
response = router.completion(
model="azure/mistral-large-latest",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
assert response is not None
def test_exception_raising():
# this tests if the router raises an exception when invalid params are set
# in this test both deployments have bad keys - Keep this test. It validates if the router raises the most recent exception
@@ -409,8 +366,8 @@ def test_exception_raising():
try:
print("testing if router raises an exception")
old_api_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = ""
old_api_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = ""
model_list = [
{
"model_name": "gpt-3.5-turbo", # openai model name
@@ -418,7 +375,7 @@ def test_exception_raising():
"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,
@@ -446,16 +403,16 @@ def test_exception_raising():
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hello this request will fail"}],
)
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
pytest.fail(f"Should have raised an Auth Error")
except openai.AuthenticationError:
print(
"Test Passed: Caught an OPENAI AUTH Error, Good job. This is what we needed!"
)
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("Got unexpected exception on router!", e)
@@ -530,7 +487,7 @@ def test_call_one_endpoint():
# this test makes a completion calls azure/gpt-4.1-mini, it should work
try:
print("Testing calling a specific deployment")
old_api_key = os.environ["AZURE_API_KEY"]
old_api_key = os.environ["AZURE_AI_API_KEY"]
model_list = [
{
@@ -539,7 +496,7 @@ def test_call_one_endpoint():
"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_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -548,8 +505,8 @@ def test_call_one_endpoint():
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "azure/text-embedding-ada-002",
"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"],
},
"tpm": 100000,
"rpm": 10000,
@@ -562,7 +519,7 @@ def test_call_one_endpoint():
set_verbose=True,
num_retries=1,
) # type: ignore
old_api_base = os.environ.pop("AZURE_API_BASE", None)
old_api_base = os.environ.pop("AZURE_AI_API_BASE", None)
async def call_azure_completion():
response = await router.acompletion(
@@ -584,8 +541,8 @@ def test_call_one_endpoint():
asyncio.run(call_azure_completion())
asyncio.run(call_azure_embedding())
os.environ["AZURE_API_BASE"] = old_api_base
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_BASE"] = old_api_base
os.environ["AZURE_AI_API_KEY"] = old_api_key
except Exception as e:
print(f"FAILED TEST")
pytest.fail(f"Got unexpected exception on router! - {e}")
@@ -594,7 +551,6 @@ def test_call_one_endpoint():
# test_call_one_endpoint()
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_async_router_context_window_fallback(sync_mode):
@@ -708,9 +664,9 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info()
"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"),
"base_model": "azure/gpt-35-turbo",
"mock_response": "Hello world 1!",
},
@@ -762,9 +718,9 @@ def test_router_context_window_check_pre_call_check():
"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"),
"base_model": "azure/gpt-35-turbo",
"mock_response": "Hello world 1!",
},
@@ -816,9 +772,9 @@ def test_router_context_window_check_pre_call_check_out_group():
"model_name": "gpt-3.5-turbo-small", # 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"),
"base_model": "azure/gpt-35-turbo",
},
},
@@ -896,9 +852,9 @@ def test_router_region_pre_call_check(allowed_model_region):
"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"),
"base_model": "azure/gpt-35-turbo",
"region_name": allowed_model_region,
},
@@ -1173,8 +1129,8 @@ def test_azure_embedding_on_router():
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "azure/text-embedding-ada-002",
"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"],
},
"tpm": 100000,
"rpm": 10000,
@@ -1381,8 +1337,8 @@ def test_reading_keys_os_environ():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"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",
"timeout": "os.environ/AZURE_TIMEOUT",
"stream_timeout": "os.environ/AZURE_STREAM_TIMEOUT",
@@ -1394,11 +1350,11 @@ def test_reading_keys_os_environ():
router = Router(model_list=model_list)
for model in router.model_list:
assert (
model["litellm_params"]["api_key"] == os.environ["AZURE_API_KEY"]
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_API_KEY']}"
model["litellm_params"]["api_key"] == os.environ["AZURE_AI_API_KEY"]
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_AI_API_KEY']}"
assert (
model["litellm_params"]["api_base"] == os.environ["AZURE_API_BASE"]
), f"{model['litellm_params']['api_base']} vs {os.environ['AZURE_API_BASE']}"
model["litellm_params"]["api_base"] == os.environ["AZURE_AI_API_BASE"]
), f"{model['litellm_params']['api_base']} vs {os.environ['AZURE_AI_API_BASE']}"
assert (
model["litellm_params"]["api_version"]
== os.environ["AZURE_API_VERSION"]
@@ -1415,8 +1371,8 @@ def test_reading_keys_os_environ():
print("passed testing of reading keys from os.environ")
model_id = model["model_info"]["id"]
async_client: openai.AsyncAzureOpenAI = router.cache.get_cache(f"{model_id}_async_client") # type: ignore
assert async_client.api_key == os.environ["AZURE_API_KEY"]
assert async_client.base_url == os.environ["AZURE_API_BASE"]
assert async_client.api_key == os.environ["AZURE_AI_API_KEY"]
assert async_client.base_url == os.environ["AZURE_AI_API_BASE"]
assert async_client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{async_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1428,8 +1384,8 @@ def test_reading_keys_os_environ():
print("\n Testing async streaming client")
stream_async_client: openai.AsyncAzureOpenAI = router.cache.get_cache(f"{model_id}_stream_async_client") # type: ignore
assert stream_async_client.api_key == os.environ["AZURE_API_KEY"]
assert stream_async_client.base_url == os.environ["AZURE_API_BASE"]
assert stream_async_client.api_key == os.environ["AZURE_AI_API_KEY"]
assert stream_async_client.base_url == os.environ["AZURE_AI_API_BASE"]
assert stream_async_client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{stream_async_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1440,8 +1396,8 @@ def test_reading_keys_os_environ():
print("\n Testing sync client")
client: openai.AzureOpenAI = router.cache.get_cache(f"{model_id}_client") # type: ignore
assert client.api_key == os.environ["AZURE_API_KEY"]
assert client.base_url == os.environ["AZURE_API_BASE"]
assert client.api_key == os.environ["AZURE_AI_API_KEY"]
assert client.base_url == os.environ["AZURE_AI_API_BASE"]
assert client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1452,8 +1408,8 @@ def test_reading_keys_os_environ():
print("\n Testing sync stream client")
stream_client: openai.AzureOpenAI = router.cache.get_cache(f"{model_id}_stream_client") # type: ignore
assert stream_client.api_key == os.environ["AZURE_API_KEY"]
assert stream_client.base_url == os.environ["AZURE_API_BASE"]
assert stream_client.api_key == os.environ["AZURE_AI_API_KEY"]
assert stream_client.base_url == os.environ["AZURE_AI_API_BASE"]
assert stream_client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{stream_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1503,7 +1459,7 @@ def test_reading_openai_keys_os_environ():
for model in router.model_list:
assert (
model["litellm_params"]["api_key"] == os.environ["OPENAI_API_KEY"]
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_API_KEY']}"
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_AI_API_KEY']}"
assert float(model["litellm_params"]["timeout"]) == float(
os.environ["AZURE_TIMEOUT"]
), f"{model['litellm_params']['timeout']} vs {os.environ['AZURE_TIMEOUT']}"
@@ -1574,7 +1530,9 @@ def test_router_anthropic_key_dynamic():
{
"model_name": "anthropic-claude",
"litellm_params": {
"model": os.environ.get("CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
"model": os.environ.get(
"CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"
),
"api_key": anthropic_api_key,
},
}
@@ -2273,8 +2231,8 @@ async def test_router_batch_endpoints(provider):
"model_name": "my-custom-name",
"litellm_params": {
"model": "azure/gpt-4o-mini",
"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"),
},
},
]
@@ -2452,8 +2410,8 @@ def test_is_team_specific_model():
# "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"),
# "tpm": 100000,
# "rpm": 100000,
# },
@@ -2462,8 +2420,8 @@ def test_is_team_specific_model():
# "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"),
# "tpm": 500,
# "rpm": 500,
# },
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
+4 -4
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")
@@ -174,9 +174,9 @@ async def test_audio_speech_router():
{
"model_name": "tts",
"litellm_params": {
"model": "azure/azure-tts",
"api_base": os.getenv("AZURE_SWEDEN_API_BASE"),
"api_key": os.getenv("AZURE_SWEDEN_API_KEY"),
"model": "azure/tts",
"api_base": os.getenv("AZURE_TTS_API_BASE"),
"api_key": os.getenv("AZURE_TTS_API_KEY"),
},
},
]
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
-704
View File
@@ -1,704 +0,0 @@
# # this tests if the router is initialized correctly
# import asyncio
# import os
# import sys
# import time
# import traceback
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# from collections import defaultdict
# from concurrent.futures import ThreadPoolExecutor
# from dotenv import load_dotenv
# import litellm
# from litellm import Router
# load_dotenv()
# # every time we load the router we should have 4 clients:
# # Async
# # Sync
# # Async + Stream
# # Sync + Stream
# def test_init_clients():
# litellm.set_verbose = True
# import logging
# from litellm._logging import verbose_router_logger
# verbose_router_logger.setLevel(logging.DEBUG)
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "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"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# },
# ]
# router = Router(model_list=model_list, set_verbose=True)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# assert router.cache.get_cache(f"{model_id}_client") is not None
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# # check if timeout for stream/non stream clients is set correctly
# async_client = router.cache.get_cache(f"{model_id}_async_client")
# stream_async_client = router.cache.get_cache(
# f"{model_id}_stream_async_client"
# )
# assert async_client.timeout == 0.01
# assert stream_async_client.timeout == 0.000_001
# print(vars(async_client))
# print()
# print(async_client._base_url)
# assert (
# async_client._base_url
# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/"
# )
# assert (
# stream_async_client._base_url
# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/"
# )
# print("PASSED !")
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_init_clients()
# def test_init_clients_basic():
# litellm.set_verbose = True
# try:
# print("Test basic client init")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "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"),
# },
# },
# ]
# router = Router(model_list=model_list)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# assert router.cache.get_cache(f"{model_id}_client") is not None
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# print("PASSED !")
# # see if we can init clients without timeout or max retries set
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_init_clients_basic()
# def test_init_clients_basic_azure_cloudflare():
# # init azure + cloudflare
# # init OpenAI gpt-3.5
# # init OpenAI text-embedding
# # init OpenAI comptaible - Mistral/mistral-medium
# # init OpenAI compatible - xinference/bge
# litellm.set_verbose = True
# try:
# print("Test basic client init")
# model_list = [
# {
# "model_name": "azure-cloudflare",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_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",
# },
# },
# {
# "model_name": "gpt-openai",
# "litellm_params": {
# "model": "gpt-3.5-turbo",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# {
# "model_name": "text-embedding-ada-002",
# "litellm_params": {
# "model": "text-embedding-ada-002",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# {
# "model_name": "mistral",
# "litellm_params": {
# "model": "mistral/mistral-tiny",
# "api_key": os.getenv("MISTRAL_API_KEY"),
# },
# },
# {
# "model_name": "bge-base-en",
# "litellm_params": {
# "model": "xinference/bge-base-en",
# "api_base": "http://127.0.0.1:9997/v1",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# ]
# router = Router(model_list=model_list)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# assert router.cache.get_cache(f"{model_id}_client") is not None
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# print("PASSED !")
# # see if we can init clients without timeout or max retries set
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_init_clients_basic_azure_cloudflare()
# def test_timeouts_router():
# """
# Test the timeouts of the router with multiple clients. This HASas to raise a timeout error
# """
# import openai
# litellm.set_verbose = True
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "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"),
# "timeout": 0.000001,
# "stream_timeout": 0.000_001,
# },
# },
# ]
# router = Router(model_list=model_list, num_retries=0)
# print("PASSED !")
# async def test():
# try:
# await router.acompletion(
# model="gpt-3.5-turbo",
# messages=[
# {"role": "user", "content": "hello, write a 20 pg essay"}
# ],
# )
# except Exception as e:
# raise e
# asyncio.run(test())
# except openai.APITimeoutError as e:
# print(
# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
# )
# print(type(e))
# pass
# except Exception as e:
# pytest.fail(
# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
# )
# # test_timeouts_router()
# def test_stream_timeouts_router():
# """
# Test the stream timeouts router. See if it selected the correct client with stream timeout
# """
# import openai
# litellm.set_verbose = True
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "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"),
# "timeout": 200, # regular calls will not timeout, stream calls will
# "stream_timeout": 10,
# },
# },
# ]
# router = Router(model_list=model_list)
# print("PASSED !")
# data = {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "hello, write a 20 pg essay"}],
# "stream": True,
# }
# selected_client = router._get_client(
# deployment=router.model_list[0],
# kwargs=data,
# client_type=None,
# )
# print("Select client timeout", selected_client.timeout)
# assert selected_client.timeout == 10
# # make actual call
# response = router.completion(**data)
# for chunk in response:
# print(f"chunk: {chunk}")
# except openai.APITimeoutError as e:
# print(
# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
# )
# print(type(e))
# pass
# except Exception as e:
# pytest.fail(
# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
# )
# # test_stream_timeouts_router()
# def test_xinference_embedding():
# # [Test Init Xinference] this tests if we init xinference on the router correctly
# # [Test Exception Mapping] tests that xinference is an openai comptiable provider
# print("Testing init xinference")
# print(
# "this tests if we create an OpenAI client for Xinference, with the correct API BASE"
# )
# model_list = [
# {
# "model_name": "xinference",
# "litellm_params": {
# "model": "xinference/bge-base-en",
# "api_base": "os.environ/XINFERENCE_API_BASE",
# },
# }
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# assert (
# router.model_list[0]["litellm_params"]["api_base"] == "http://0.0.0.0:9997"
# ) # set in env
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "xinference"},
# )
# assert openai_client._base_url == "http://0.0.0.0:9997"
# assert "xinference" in litellm.openai_compatible_providers
# print("passed")
# # test_xinference_embedding()
# def test_router_init_gpt_4_vision_enhancements():
# try:
# # tests base_url set when any base_url with /openai/deployments passed to router
# print("Testing Azure GPT_Vision enhancements")
# model_list = [
# {
# "model_name": "gpt-4-vision-enhancements",
# "litellm_params": {
# "model": "azure/gpt-4-vision",
# "api_key": os.getenv("AZURE_API_KEY"),
# "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/",
# "dataSources": [
# {
# "type": "AzureComputerVision",
# "parameters": {
# "endpoint": "os.environ/AZURE_VISION_ENHANCE_ENDPOINT",
# "key": "os.environ/AZURE_VISION_ENHANCE_KEY",
# },
# }
# ],
# },
# }
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# assert (
# router.model_list[0]["litellm_params"]["base_url"]
# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/"
# ) # set in env
# assert (
# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][
# "endpoint"
# ]
# == os.environ["AZURE_VISION_ENHANCE_ENDPOINT"]
# )
# assert (
# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][
# "key"
# ]
# == os.environ["AZURE_VISION_ENHANCE_KEY"]
# )
# azure_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"stream": True, "model": "gpt-4-vision-enhancements"},
# client_type="async",
# )
# assert (
# azure_client._base_url
# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/"
# )
# print("passed")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# @pytest.mark.parametrize("sync_mode", [True, False])
# @pytest.mark.asyncio
# async def test_openai_with_organization(sync_mode):
# try:
# print("Testing OpenAI with organization")
# model_list = [
# {
# "model_name": "openai-bad-org",
# "litellm_params": {
# "model": "gpt-3.5-turbo",
# "organization": "org-ikDc4ex8NB",
# },
# },
# {
# "model_name": "openai-good-org",
# "litellm_params": {"model": "gpt-3.5-turbo"},
# },
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# if sync_mode:
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
# )
# print(vars(openai_client))
# assert openai_client.organization == "org-ikDc4ex8NB"
# # bad org raises error
# try:
# response = router.completion(
# model="openai-bad-org",
# messages=[{"role": "user", "content": "this is a test"}],
# )
# pytest.fail(
# "Request should have failed - This organization does not exist"
# )
# except Exception as e:
# print("Got exception: " + str(e))
# assert "header should match organization for API key" in str(
# e
# ) or "No such organization" in str(e)
# # good org works
# response = router.completion(
# model="openai-good-org",
# messages=[{"role": "user", "content": "this is a test"}],
# max_tokens=5,
# )
# else:
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
# client_type="async",
# )
# print(vars(openai_client))
# assert openai_client.organization == "org-ikDc4ex8NB"
# # bad org raises error
# try:
# response = await router.acompletion(
# model="openai-bad-org",
# messages=[{"role": "user", "content": "this is a test"}],
# )
# pytest.fail(
# "Request should have failed - This organization does not exist"
# )
# except Exception as e:
# print("Got exception: " + str(e))
# assert "header should match organization for API key" in str(
# e
# ) or "No such organization" in str(e)
# # good org works
# response = await router.acompletion(
# model="openai-good-org",
# messages=[{"role": "user", "content": "this is a test"}],
# max_tokens=5,
# )
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# def test_init_clients_azure_command_r_plus():
# # This tests that the router uses the OpenAI client for Azure/Command-R+
# # For azure/command-r-plus we need to use openai.OpenAI because of how the Azure provider requires requests being sent
# litellm.set_verbose = True
# import logging
# from litellm._logging import verbose_router_logger
# verbose_router_logger.setLevel(logging.DEBUG)
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/command-r-plus",
# "api_key": os.getenv("AZURE_COHERE_API_KEY"),
# "api_base": os.getenv("AZURE_COHERE_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# },
# ]
# router = Router(model_list=model_list, set_verbose=True)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# async_client = router.cache.get_cache(f"{model_id}_async_client")
# stream_async_client = router.cache.get_cache(
# f"{model_id}_stream_async_client"
# )
# # Assert the Async Clients used are OpenAI clients and not Azure
# # For using Azure/Command-R-Plus and Azure/Mistral the clients NEED to be OpenAI clients used
# # this is weirdness introduced on Azure's side
# assert "openai.AsyncOpenAI" in str(async_client)
# assert "openai.AsyncOpenAI" in str(stream_async_client)
# print("PASSED !")
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# @pytest.mark.asyncio
# async def test_aaaaatext_completion_with_organization():
# try:
# print("Testing Text OpenAI with organization")
# model_list = [
# {
# "model_name": "openai-bad-org",
# "litellm_params": {
# "model": "text-completion-openai/gpt-3.5-turbo-instruct",
# "api_key": os.getenv("OPENAI_API_KEY", None),
# "organization": "org-ikDc4ex8NB",
# },
# },
# {
# "model_name": "openai-good-org",
# "litellm_params": {
# "model": "text-completion-openai/gpt-3.5-turbo-instruct",
# "api_key": os.getenv("OPENAI_API_KEY", None),
# "organization": os.getenv("OPENAI_ORGANIZATION", None),
# },
# },
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
# )
# print(vars(openai_client))
# assert openai_client.organization == "org-ikDc4ex8NB"
# # bad org raises error
# try:
# response = await router.atext_completion(
# model="openai-bad-org",
# prompt="this is a test",
# )
# pytest.fail("Request should have failed - This organization does not exist")
# except Exception as e:
# print("Got exception: " + str(e))
# assert "header should match organization for API key" in str(
# e
# ) or "No such organization" in str(e)
# # good org works
# response = await router.atext_completion(
# model="openai-good-org",
# prompt="this is a test",
# max_tokens=5,
# )
# print("working response: ", response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# def test_init_clients_async_mode():
# litellm.set_verbose = True
# import logging
# from litellm._logging import verbose_router_logger
# from litellm.types.router import RouterGeneralSettings
# verbose_router_logger.setLevel(logging.DEBUG)
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "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"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# },
# ]
# router = Router(
# model_list=model_list,
# set_verbose=True,
# router_general_settings=RouterGeneralSettings(async_only_mode=True),
# )
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# # sync clients not initialized in async_only_mode=True
# assert router.cache.get_cache(f"{model_id}_client") is None
# assert router.cache.get_cache(f"{model_id}_stream_client") is None
# # only async clients initialized in async_only_mode=True
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# @pytest.mark.parametrize(
# "environment,expected_models",
# [
# ("development", ["gpt-3.5-turbo"]),
# ("production", ["gpt-4", "gpt-3.5-turbo", "gpt-4o"]),
# ],
# )
# def test_init_router_with_supported_environments(environment, expected_models):
# """
# Tests that the correct models are setup on router when LITELLM_ENVIRONMENT is set
# """
# os.environ["LITELLM_ENVIRONMENT"] = environment
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "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"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# "model_info": {"supported_environments": ["development", "production"]},
# },
# {
# "model_name": "gpt-4",
# "litellm_params": {
# "model": "openai/gpt-4",
# "api_key": os.getenv("OPENAI_API_KEY"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# "model_info": {"supported_environments": ["production"]},
# },
# {
# "model_name": "gpt-4o",
# "litellm_params": {
# "model": "openai/gpt-4o",
# "api_key": os.getenv("OPENAI_API_KEY"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# "model_info": {"supported_environments": ["production"]},
# },
# ]
# router = Router(model_list=model_list, set_verbose=True)
# _model_list = router.get_model_names()
# print("model_list: ", _model_list)
# print("expected_models: ", expected_models)
# assert set(_model_list) == set(expected_models)
# os.environ.pop("LITELLM_ENVIRONMENT")
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_streaming.py
+9 -79
View File
@@ -552,7 +552,6 @@ async def test_completion_predibase_streaming(sync_mode):
pytest.fail(f"Error occurred: {e}")
def test_completion_azure_function_calling_stream():
try:
litellm.set_verbose = False
@@ -1655,80 +1654,9 @@ def test_sagemaker_weird_response():
# test_sagemaker_weird_response()
@pytest.mark.skip(reason="Move to being a mock endpoint")
@pytest.mark.asyncio
async def test_sagemaker_streaming_async():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
litellm.set_verbose = True
response = await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
messages=messages,
temperature=0.2,
max_tokens=80,
aws_region_name=os.getenv("AWS_REGION_NAME_2"),
aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
stream=True,
)
# Add any assertions here to check the response
print(response)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
idx = 0
async for chunk in response:
# print
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
complete_response += chunk
if finished:
break
idx += 1
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"completion_response: {complete_response}")
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# asyncio.run(test_sagemaker_streaming_async())
@pytest.mark.skip(reason="costly sagemaker deployment. Move to mock implementation")
def test_completion_sagemaker_stream():
try:
response = completion(
model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
messages=messages,
temperature=0.2,
max_tokens=80,
aws_region_name=os.getenv("AWS_REGION_NAME_2"),
aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
stream=True,
)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
for idx, chunk in enumerate(response):
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
if finished:
break
complete_response += chunk
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Account deleted by IBM.")
@pytest.mark.asyncio
async def test_completion_watsonx_stream():
@@ -2725,8 +2653,8 @@ def test_azure_streaming_and_function_calling():
tool_choice="auto",
messages=messages,
stream=True,
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="2024-02-15-preview",
)
# Add any assertions here to check the response
@@ -2796,8 +2724,8 @@ async def test_azure_astreaming_and_function_calling():
tool_choice="auto",
messages=messages,
stream=True,
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="2024-02-15-preview",
caching=True,
)
@@ -2827,8 +2755,8 @@ async def test_azure_astreaming_and_function_calling():
tool_choice="auto",
messages=messages,
stream=True,
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="2024-02-15-preview",
caching=True,
)
@@ -3109,7 +3037,9 @@ def test_unit_test_custom_stream_wrapper_repeating_chunk(
print(f"expected_chunk_fail: {expected_chunk_fail}")
if (loop_amount > litellm.REPEATED_STREAMING_CHUNK_LIMIT) and expected_chunk_fail:
with pytest.raises((litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError)):
with pytest.raises(
(litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError)
):
for chunk in response:
continue
else:
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/local_testing/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/logging_callback_tests/test_alerting.py
+10 -49
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",
@@ -788,40 +788,6 @@ async def test_region_outage_alerting_called(
mock_send_alert.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.skip(reason="test only needs to run locally ")
async def test_alerting():
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "bad_key",
},
}
],
debug_level="DEBUG",
set_verbose=True,
alerting_config=AlertingConfig(
alerting_threshold=10, # threshold for slow / hanging llm responses (in seconds). Defaults to 300 seconds
webhook_url=os.getenv(
"SLACK_WEBHOOK_URL"
), # webhook you want to send alerts to
),
)
try:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
except Exception:
pass
finally:
await asyncio.sleep(3)
@pytest.mark.asyncio
async def test_langfuse_trace_id():
"""
@@ -868,7 +834,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 +975,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 +1045,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 +1065,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 +1074,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_amazing_s3_logs.py
+28 -94
View File
@@ -74,15 +74,13 @@ async def test_basic_s3_logging(sync_mode, streaming):
s3.delete_object(Bucket="load-testing-oct", Key=key)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"streaming", [(True)]
)
@pytest.mark.parametrize("streaming", [True])
@pytest.mark.flaky(retries=3, delay=1)
async def test_basic_s3_v2_logging(streaming):
from blockbuster import BlockBuster
from litellm.integrations.s3_v2 import S3Logger
s3_v2_logger = S3Logger(s3_flush_interval=1)
litellm.callbacks = [s3_v2_logger]
blockbuster = BlockBuster()
@@ -120,7 +118,7 @@ async def test_basic_s3_v2_logging(streaming):
print(f"all_s3_keys: {all_s3_keys}")
#assert that atlest one key has response.id in it
# assert that atlest one key has response.id in it
assert any(response_id in key for key in all_s3_keys)
s3 = boto3.client("s3")
# delete all objects
@@ -134,22 +132,22 @@ async def test_basic_s3_v2_logging_failure():
"""Test that S3 v2 logger makes httpx PUT request when logging failures"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.integrations.s3_v2 import S3Logger
# Create S3 logger with short flush interval
s3_v2_logger = S3Logger(s3_flush_interval=1)
# Mock the httpx client to capture the PUT request
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.raise_for_status = MagicMock()
s3_v2_logger.async_httpx_client = AsyncMock()
s3_v2_logger.async_httpx_client.put.return_value = mock_response
# Track the upload method calls
original_upload = s3_v2_logger.async_upload_data_to_s3
upload_called = False
async def mock_upload(batch_logging_element):
nonlocal upload_called
upload_called = True
@@ -157,12 +155,12 @@ async def test_basic_s3_v2_logging_failure():
url = f"https://test-bucket.s3.us-west-2.amazonaws.com/{batch_logging_element.s3_object_key}"
headers = {"Content-Type": "application/json"}
data = '{"model": "gpt-4o-mini"}'
# Make the actual httpx call we want to test
await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data)
s3_v2_logger.async_upload_data_to_s3 = mock_upload
# Configure S3 callback params
litellm.callbacks = [s3_v2_logger]
litellm.s3_callback_params = {
@@ -172,7 +170,7 @@ async def test_basic_s3_v2_logging_failure():
"s3_region_name": "us-west-2",
}
litellm.set_verbose = True
# Trigger a failure by using invalid API key
try:
response = await litellm.acompletion(
@@ -182,33 +180,33 @@ async def test_basic_s3_v2_logging_failure():
)
except Exception as e:
print(f"Expected error: {e}")
# Wait for logger to process the failure
await asyncio.sleep(5)
# Verify that our mock upload was called
assert upload_called, "S3 upload method was not called"
print("✓ S3 upload method was called")
# Verify that httpx PUT was called
s3_v2_logger.async_httpx_client.put.assert_called()
# Get the call arguments to verify the S3 URL
call_args = s3_v2_logger.async_httpx_client.put.call_args
assert call_args is not None
url = call_args[1]['url'] if 'url' in call_args[1] else call_args[0][0]
url = call_args[1]["url"] if "url" in call_args[1] else call_args[0][0]
# Verify the URL contains expected S3 endpoint
assert "test-bucket.s3.us-west-2.amazonaws.com" in url
print(f"✓ S3 PUT request made to: {url}")
# Verify headers include expected content type
headers = call_args[1]['headers']
assert headers['Content-Type'] == 'application/json'
headers = call_args[1]["headers"]
assert headers["Content-Type"] == "application/json"
print("✓ S3 request headers are correct")
# Verify JSON data was included
data = call_args[1]['data']
data = call_args[1]["data"]
assert data is not None
assert '"model": "gpt-4o-mini"' in data
print("✓ S3 request data contains expected log payload")
@@ -411,83 +409,19 @@ async def make_async_calls():
return total_time
@pytest.mark.skip(reason="flaky test on ci/cd")
def test_s3_logging_r2():
# all s3 requests need to be in one test function
# since we are modifying stdout, and pytests runs tests in parallel
# on circle ci - we only test litellm.acompletion()
try:
# redirect stdout to log_file
# litellm.cache = litellm.Cache(
# type="s3", s3_bucket_name="litellm-r2-bucket", s3_region_name="us-west-2"
# )
litellm.set_verbose = True
from litellm._logging import verbose_logger
import logging
verbose_logger.setLevel(level=logging.DEBUG)
litellm.success_callback = ["s3"]
litellm.s3_callback_params = {
"s3_bucket_name": "litellm-r2-bucket",
"s3_aws_secret_access_key": "os.environ/R2_S3_ACCESS_KEY",
"s3_aws_access_key_id": "os.environ/R2_S3_ACCESS_ID",
"s3_endpoint_url": "os.environ/R2_S3_URL",
"s3_region_name": "os.environ/R2_S3_REGION_NAME",
}
print("Testing async s3 logging")
expected_keys = []
import time
curr_time = str(time.time())
async def _test():
return await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": f"This is a test {curr_time}"}],
max_tokens=10,
temperature=0.7,
user="ishaan-2",
)
response = asyncio.run(_test())
print(f"response: {response}")
expected_keys.append(response.id)
import boto3
s3 = boto3.client(
"s3",
endpoint_url=os.getenv("R2_S3_URL"),
region_name=os.getenv("R2_S3_REGION_NAME"),
aws_access_key_id=os.getenv("R2_S3_ACCESS_ID"),
aws_secret_access_key=os.getenv("R2_S3_ACCESS_KEY"),
)
bucket_name = "litellm-r2-bucket"
# List objects in the bucket
response = s3.list_objects(Bucket=bucket_name)
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
finally:
# post, close log file and verify
# Reset stdout to the original value
print("Passed! Testing async s3 logging")
from litellm.integrations.s3_v2 import S3Logger
class TestS3Logger(S3Logger):
def __init__(self, *args, **kwargs):
self.recorded_requests = {}
self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
super().__init__(*args, **kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.recorded_requests[response_obj["id"]] = start_time
print("recorded request", self.recorded_requests)
self.logged_standard_logging_payload = kwargs["standard_logging_object"]
return await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
return await super().async_log_success_event(
kwargs, response_obj, start_time, end_time
)
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/ocr_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
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_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py
+2 -2
View File
@@ -30,7 +30,7 @@ class TestAzureAnthropicStructuredOutput(BaseAnthropicMessagesStructuredOutputTe
return "azure_ai/claude-opus-4-5"
def get_api_base(self) -> Optional[str]:
return "https://krish-mh44t553-eastus2.services.ai.azure.com/"
return "https://krris-mnb3t0vd-swedencentral.services.ai.azure.com"
def get_api_key(self) -> Optional[str]:
return os.environ.get("AZURE_ANTHROPIC_API_KEY")
return os.environ.get("AZURE_ANTHROPIC_API_KEY")
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/proxy_unit_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/store_model_in_db_tests/test_adding_passthrough_model.py
+6 -6
View File
@@ -1,12 +1,12 @@
"""
Test adding a pass through assemblyai model + api key + api base to the db
wait 20 seconds
make request
wait 20 seconds
make request
Cases to cover
1. user points api base to <proxy-base>/assemblyai
Cases to cover
1. user points api base to <proxy-base>/assemblyai
2. user points api base to <proxy-base>/asssemblyai/us
3. user points api base to <proxy-base>/assemblyai/eu
3. user points api base to <proxy-base>/assemblyai/eu
4. Bad API Key / credential - 401
"""
@@ -21,7 +21,7 @@ TEST_MASTER_KEY = "sk-1234"
PROXY_BASE_URL = "http://0.0.0.0:4000"
US_BASE_URL = f"{PROXY_BASE_URL}/assemblyai"
EU_BASE_URL = f"{PROXY_BASE_URL}/eu.assemblyai"
ASSEMBLYAI_API_KEY_ENV_VAR = "TEST_SPECIAL_ASSEMBLYAI_API_KEY"
ASSEMBLYAI_API_KEY_ENV_VAR = "ASSEMBLYAI_API_KEY"
def _delete_all_assemblyai_models_from_db():
tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json
+20
View File
@@ -0,0 +1,20 @@
{
"id": "ftjob-azure-create-123",
"object": "fine_tuning.job",
"created_at": 1735689600,
"model": "davinci-002",
"status": "cancelled",
"fine_tuned_model": null,
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
"hyperparameters": {
"n_epochs": 3,
"batch_size": null,
"learning_rate_multiplier": null
},
"organization_id": "",
"result_files": [],
"validation_file": null,
"trained_tokens": null,
"estimated_finish": null,
"error": null
}
tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json
+18
View File
@@ -0,0 +1,18 @@
{
"id": "ftjob-azure-create-123",
"object": "fine_tuning.job",
"created_at": 1735689600,
"model": "davinci-002",
"status": "canceled",
"fine_tuned_model": null,
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
"hyperparameters": {
"n_epochs": 3
},
"organization_id": null,
"result_files": null,
"validation_file": null,
"trained_tokens": null,
"estimated_finish": null,
"error": null
}
tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json
+3
View File
@@ -0,0 +1,3 @@
{
"fine_tuning_job_id": "ftjob-azure-create-123"
}
tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json
+20
View File
@@ -0,0 +1,20 @@
{
"id": "ftjob-azure-create-123",
"object": "fine_tuning.job",
"created_at": 1735689600,
"model": "davinci-002",
"status": "running",
"fine_tuned_model": null,
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
"hyperparameters": {
"n_epochs": 3,
"batch_size": null,
"learning_rate_multiplier": null
},
"organization_id": "",
"result_files": [],
"validation_file": null,
"trained_tokens": null,
"estimated_finish": null,
"error": null
}
tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json
+18
View File
@@ -0,0 +1,18 @@
{
"id": "ftjob-azure-create-123",
"object": "fine_tuning.job",
"created_at": 1735689600,
"model": "davinci-002",
"status": "running",
"fine_tuned_model": null,
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
"hyperparameters": {
"n_epochs": 3
},
"organization_id": null,
"result_files": null,
"validation_file": null,
"trained_tokens": null,
"estimated_finish": null,
"error": null
}
tests/test_litellm/expected_fine_tuning_api/azure_create_request.json
+8
View File
@@ -0,0 +1,8 @@
{
"model": "gpt-35-turbo-1106",
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
"hyperparameters": {},
"extra_body": {
"trainingType": 1
}
}
tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json
+20
View File
@@ -0,0 +1,20 @@
{
"object": "list",
"data": [
{
"id": "ftjob-azure-create-123",
"object": "fine_tuning.job",
"created_at": 1735689600,
"model": "davinci-002",
"status": "running"
},
{
"id": "ftjob-azure-prev-000",
"object": "fine_tuning.job",
"created_at": 1735603200,
"model": "davinci-002",
"status": "succeeded"
}
],
"has_more": false
}
tests/test_litellm/expected_fine_tuning_api/azure_list_request.json
+4
View File
@@ -0,0 +1,4 @@
{
"after": "ftjob-azure-prev-000",
"limit": 2
}
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/test_azure_fine_tuning_api.py
+150
View File
@@ -0,0 +1,150 @@
import json
from pathlib import Path
from unittest.mock import AsyncMock, patch
import pytest
from openai import AsyncAzureOpenAI
import litellm
from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI
def _expected_dir() -> Path:
return Path(__file__).resolve().parent.parent.parent / "expected_fine_tuning_api"
def _load_json(file_name: str) -> dict:
path = _expected_dir() / file_name
assert path.exists(), f"Expected fixture file not found: {path}"
with open(path) as f:
return json.load(f)
class _MockSDKResponse:
def __init__(self, payload: dict):
self._payload = payload
def model_dump(self) -> dict:
return self._payload
def _mock_azure_client(
create_payload: dict | None = None,
list_payload: dict | None = None,
cancel_payload: dict | None = None,
):
client = AsyncAzureOpenAI(
api_key="test-key",
api_version="2024-10-21",
azure_endpoint="https://exampleopenaiendpoint-production.up.railway.app",
)
client.fine_tuning.jobs.create = AsyncMock(
return_value=(
_MockSDKResponse(create_payload) if create_payload is not None else None
)
) # type: ignore[method-assign]
client.fine_tuning.jobs.list = AsyncMock(
return_value=list_payload
) # type: ignore[method-assign]
client.fine_tuning.jobs.cancel = AsyncMock(
return_value=(
_MockSDKResponse(cancel_payload) if cancel_payload is not None else None
)
) # type: ignore[method-assign]
return client
@pytest.mark.asyncio
async def test_azure_acreate_fine_tuning_job_request_and_output_match_expected_json():
expected_request = _load_json("azure_create_request.json")
raw_response = _load_json("azure_create_raw_response.json")
expected_output = _load_json("azure_create_expected_output.json")
mock_client = _mock_azure_client(create_payload=raw_response)
with patch.object(
AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client
):
response = await litellm.acreate_fine_tuning_job(
model="gpt-35-turbo-1106",
training_file="file-5e4b20ecbd724182b9964f3cd2ab7212",
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
api_key="test-key",
api_version="2024-10-21",
)
request_kwargs = mock_client.fine_tuning.jobs.create.call_args.kwargs
assert request_kwargs == expected_request
response_dict = response.model_dump(exclude={"_hidden_params"})
for key, expected_value in expected_output.items():
assert key in response_dict, f"Missing key in response: {key}"
assert response_dict[key] == expected_value
assert response.id is not None
assert response.model == "davinci-002"
@pytest.mark.asyncio
async def test_azure_alist_fine_tuning_jobs_request_matches_expected_json():
expected_request = _load_json("azure_list_request.json")
raw_list_response = _load_json("azure_list_raw_response.json")
mock_client = _mock_azure_client(list_payload=raw_list_response)
with patch.object(
AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client
):
response = await litellm.alist_fine_tuning_jobs(
after=expected_request["after"],
limit=expected_request["limit"],
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
api_key="test-key",
api_version="2024-10-21",
)
request_kwargs = mock_client.fine_tuning.jobs.list.call_args.kwargs
assert request_kwargs == expected_request
assert response == raw_list_response
@pytest.mark.asyncio
async def test_azure_acancel_fine_tuning_job_request_and_output_match_expected_json():
expected_request = _load_json("azure_cancel_request.json")
raw_response = _load_json("azure_cancel_raw_response.json")
expected_output = _load_json("azure_cancel_expected_output.json")
mock_client = _mock_azure_client(cancel_payload=raw_response)
with patch.object(
AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client
):
response = await litellm.acancel_fine_tuning_job(
fine_tuning_job_id=expected_request["fine_tuning_job_id"],
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
api_key="test-key",
api_version="2024-10-21",
)
request_kwargs = mock_client.fine_tuning.jobs.cancel.call_args.kwargs
assert request_kwargs == expected_request
response_dict = response.model_dump(exclude={"_hidden_params"})
for key, expected_value in expected_output.items():
assert key in response_dict, f"Missing key in response: {key}"
assert response_dict[key] == expected_value
assert response.status == "cancelled"
def test_azure_trainingtype_defaults_to_one():
handler = AzureOpenAIFineTuningAPI()
create_data = {"model": "gpt-4o-mini", "training_file": "file-test"}
handler._ensure_training_type(create_data)
assert "extra_body" in create_data
assert create_data["extra_body"]["trainingType"] == 1
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
tests/test_team_logging.py
-134
View File
@@ -59,137 +59,3 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=
if status != 200:
raise Exception(f"Request did not return a 200 status code: {status}")
@pytest.mark.skip(reason="flaky test - covered by simpler unit testing.")
@pytest.mark.asyncio
@pytest.mark.flaky(retries=12, delay=2)
async def test_aaateam_logging():
"""
-> Team 1 logs to project 1
-> Create Key
-> Make chat/completions call
-> Fetch logs from langfuse
"""
try:
async with aiohttp.ClientSession() as session:
key = await generate_key(
session, models=["fake-openai-endpoint"], team_id="team-1"
) # team-1 logs to project 1
from litellm._uuid import uuid
_trace_id = f"trace-{uuid.uuid4()}"
_request_metadata = {
"trace_id": _trace_id,
}
await chat_completion(
session,
key["key"],
model="fake-openai-endpoint",
request_metadata=_request_metadata,
)
# Test - if the logs were sent to the correct team on langfuse
import langfuse
print(f"langfuse_public_key: {os.getenv('LANGFUSE_PROJECT1_PUBLIC')}")
print(f"langfuse_secret_key: {os.getenv('LANGFUSE_HOST')}")
langfuse_client = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"),
secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"),
host="https://us.cloud.langfuse.com",
)
await asyncio.sleep(30)
print(f"searching for trace_id={_trace_id} on langfuse")
generations = langfuse_client.get_generations(trace_id=_trace_id).data
print(generations)
assert len(generations) == 1
except Exception as e:
pytest.fail(f"Unexpected error: {str(e)}")
@pytest.mark.skip(reason="todo fix langfuse credential error")
@pytest.mark.asyncio
async def test_team_2logging():
"""
-> Team 1 logs to project 2
-> Create Key
-> Make chat/completions call
-> Fetch logs from langfuse
"""
langfuse_public_key = os.getenv("LANGFUSE_PROJECT2_PUBLIC")
print(f"langfuse_public_key: {langfuse_public_key}")
langfuse_secret_key = os.getenv("LANGFUSE_PROJECT2_SECRET")
print(f"langfuse_secret_key: {langfuse_secret_key}")
langfuse_host = "https://us.cloud.langfuse.com"
try:
assert langfuse_public_key is not None
assert langfuse_secret_key is not None
except Exception as e:
# skip test if langfuse credentials are not set
return
try:
async with aiohttp.ClientSession() as session:
key = await generate_key(
session, models=["fake-openai-endpoint"], team_id="team-2"
) # team-1 logs to project 1
from litellm._uuid import uuid
_trace_id = f"trace-{uuid.uuid4()}"
_request_metadata = {
"trace_id": _trace_id,
}
await chat_completion(
session,
key["key"],
model="fake-openai-endpoint",
request_metadata=_request_metadata,
)
# Test - if the logs were sent to the correct team on langfuse
import langfuse
langfuse_client = langfuse.Langfuse(
public_key=langfuse_public_key,
secret_key=langfuse_secret_key,
host=langfuse_host,
)
await asyncio.sleep(30)
print(f"searching for trace_id={_trace_id} on langfuse")
generations = langfuse_client.get_generations(trace_id=_trace_id).data
print("Team 2 generations", generations)
# team-2 should have 1 generation with this trace id
assert len(generations) == 1
# team-1 should have 0 generations with this trace id
langfuse_client_1 = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"),
secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"),
host="https://us.cloud.langfuse.com",
)
generations_team_1 = langfuse_client_1.get_generations(
trace_id=_trace_id
).data
print("Team 1 generations", generations_team_1)
assert len(generations_team_1) == 0
except Exception as e:
pytest.fail("Team 2 logging failed: " + str(e))
tests/unified_google_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/vector_store_tests/test_azure_ai_vector_store.py
+3 -3
View File
@@ -17,10 +17,10 @@ async def test_basic_search_vector_store(sync_mode):
"vector_store_id": "my-vector-index",
"custom_llm_provider": "azure_ai",
"azure_search_service_name": "azure-kb-search",
"litellm_embedding_model": "azure/text-embedding-3-large",
"litellm_embedding_model": "azure_ai/text-embedding-3-large",
"litellm_embedding_config": {
"api_base": os.getenv("AZURE_AI_SEARCH_EMBEDDING_API_BASE"),
"api_key": os.getenv("AZURE_AI_SEARCH_EMBEDDING_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
},
"api_key": os.getenv("AZURE_SEARCH_API_KEY"),
}