mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
test: move offline router, retry and latency tests from local_testing to tests/unit (#45551)
* test: move offline router, latency-routing, retry and exception tests from local_testing into tests/unit Moves 29 offline nodes into the tests/unit files that mirror the code they exercise. Hosted-provider calls are mocked at the HTTP boundary with respx, and sleeps that only built a start/end gap are replaced by explicit timestamps. Legacy files left empty are deleted. * ci: record the live prompt caching cases left in tests/local_testing as unrun Moving the router prompt caching test to tests/unit leaves only live Anthropic and Vertex cases in tests/local_testing/test_anthropic_prompt_caching.py, which every CircleCI -k already deselects * test: freeze the clock in the moved lowest-latency routing tests The latency logger buckets usage by the current minute, so a run that crosses a minute boundary could grow the cache and fail the memory check. The moved tests now use fixed timestamps and a frozen clock in the logger module * ci: describe the mixed prompt caching file accurately in the coverage allowlist
This commit is contained in:
yuneng-jiang
GitHub
parent
6d4fa56ac4
commit
3eeca75ca0
20 changed files with 1324 additions and 1272 deletions
10
.github/ci-coverage-allowlist.yml
vendored
10
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -34,6 +34,16 @@ test_paths:
|
|||
- tests/local_testing/test_caching.py
|
||||
- tests/local_testing/test_disk_cache_unit_tests.py
|
||||
- tests/local_testing/test_gcs_cache_unit_tests.py
|
||||
- reason: >-
|
||||
Mixed file that no job has ever run in full. Every CircleCI job that globs
|
||||
tests/local_testing deselects it by name ("caching") or keeps only another keyword, so only
|
||||
its router test ran, and that test now lives in tests/unit/router_utils/pre_call_checks.
|
||||
Five of the seven left need ANTHROPIC_API_KEY or Vertex credentials. The other two,
|
||||
test_litellm_anthropic_prompt_caching_tools and test_litellm_anthropic_prompt_caching_system,
|
||||
are offline mocks that already fail on main against a stale expected request body; they need
|
||||
that fixed before they can move to tests/unit
|
||||
paths:
|
||||
- tests/local_testing/test_anthropic_prompt_caching.py
|
||||
- reason: >-
|
||||
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
|
||||
from a pull request; it needs a live gateway and provider credentials no PR job holds
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -137,6 +137,7 @@ jobs:
|
|||
include:
|
||||
- shard: core-utils
|
||||
test-path: |-
|
||||
tests/unit/batch_completion
|
||||
tests/unit/decisions
|
||||
tests/unit/litellm_core_utils
|
||||
python-version: "3.12"
|
||||
|
|
|
|||
|
|
@ -592,74 +592,3 @@ async def test_litellm_anthropic_prompt_caching_system():
|
|||
mock_post.assert_called_once_with(
|
||||
expected_url, json=expected_json, headers=expected_headers, timeout=600.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
# )
|
||||
async def test_router_with_prompt_caching(anthropic_messages):
|
||||
"""
|
||||
if prompt caching supported model called with prompt caching valid prompt,
|
||||
then 2nd call should go to the same model.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.prompt_caching_cache import PromptCachingCache
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "claude-model",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5-20250929",
|
||||
"api_key": os.environ.get("ANTHROPIC_API_KEY"),
|
||||
"mock_response": "The sky is blue.",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-model",
|
||||
"litellm_params": {
|
||||
"model": "anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"mock_response": "The sky is green.",
|
||||
},
|
||||
},
|
||||
],
|
||||
optional_pre_call_checks=["prompt_caching"],
|
||||
)
|
||||
|
||||
response = await router.acompletion(
|
||||
messages=anthropic_messages,
|
||||
model="claude-model",
|
||||
mock_response="The sky is blue.",
|
||||
)
|
||||
print("response=", response)
|
||||
|
||||
initial_model_id = response._hidden_params["model_id"]
|
||||
|
||||
await asyncio.sleep(1)
|
||||
cache = PromptCachingCache(
|
||||
cache=router.cache,
|
||||
)
|
||||
|
||||
cached_model_id = cache.get_model_id(messages=anthropic_messages, tools=None)
|
||||
|
||||
assert cached_model_id is not None
|
||||
prompt_caching_cache_key = PromptCachingCache.get_prompt_caching_cache_key(
|
||||
messages=anthropic_messages, tools=None
|
||||
)
|
||||
print(f"prompt_caching_cache_key: {prompt_caching_cache_key}")
|
||||
assert cached_model_id["model_id"] == initial_model_id
|
||||
|
||||
new_messages = anthropic_messages + [
|
||||
{"role": "user", "content": "What is the weather in SF?"}
|
||||
]
|
||||
|
||||
for _ in range(20):
|
||||
response = await router.acompletion(
|
||||
messages=new_messages,
|
||||
model="claude-model",
|
||||
mock_response="The sky is blue.",
|
||||
)
|
||||
print("response=", response)
|
||||
|
||||
assert response._hidden_params["model_id"] == initial_model_id
|
||||
|
|
|
|||
|
|
@ -7,9 +7,7 @@ load_dotenv()
|
|||
import pytest
|
||||
import openai
|
||||
import litellm
|
||||
from litellm import completion_with_retries, completion, acompletion_with_retries
|
||||
from litellm import responses_with_retries, aresponses_with_retries
|
||||
from litellm.responses.main import responses, aresponses
|
||||
from litellm import completion
|
||||
from litellm import (
|
||||
AuthenticationError,
|
||||
BadRequestError,
|
||||
|
|
@ -48,43 +46,6 @@ def test_completion_with_0_num_retries():
|
|||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
async def test_completion_with_retry_policy(sync_mode):
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from litellm.types.router import RetryPolicy
|
||||
|
||||
retry_number = 1
|
||||
retry_policy = RetryPolicy(
|
||||
BadRequestErrorRetries=10,
|
||||
ContentPolicyViolationErrorRetries=retry_number, # run 3 retries for ContentPolicyViolationErrors
|
||||
AuthenticationErrorRetries=0, # run 0 retries for AuthenticationErrorRetries
|
||||
)
|
||||
|
||||
target_function = "completion_with_retries"
|
||||
|
||||
with patch.object(litellm, target_function) as mock_completion_with_retries:
|
||||
data = {
|
||||
"model": "azure/gpt-3.5-turbo",
|
||||
"messages": [{"gm": "vibe", "role": "user"}],
|
||||
"retry_policy": retry_policy,
|
||||
"mock_response": "Exception: content_filter_policy",
|
||||
}
|
||||
try:
|
||||
if sync_mode:
|
||||
completion(**data)
|
||||
else:
|
||||
await completion(**data)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
mock_completion_with_retries.assert_called_once()
|
||||
assert (
|
||||
mock_completion_with_retries.call_args.kwargs["num_retries"] == retry_number
|
||||
)
|
||||
assert retry_policy.ContentPolicyViolationErrorRetries == retry_number
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
async def test_completion_with_retry_policy_no_error(sync_mode):
|
||||
|
|
@ -112,130 +73,3 @@ async def test_completion_with_retry_policy_no_error(sync_mode):
|
|||
await completion(**data)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_with_retries(sync_mode):
|
||||
"""
|
||||
If completion_with_retries is called with num_retries=3, and max_retries=0, then litellm.completion should receive num_retries , max_retries=0
|
||||
"""
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
|
||||
if sync_mode:
|
||||
target_function = "completion"
|
||||
else:
|
||||
target_function = "acompletion"
|
||||
|
||||
with patch.object(litellm, target_function) as mock_completion:
|
||||
if sync_mode:
|
||||
completion_with_retries(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"gm": "vibe", "role": "user"}],
|
||||
num_retries=3,
|
||||
original_function=mock_completion,
|
||||
)
|
||||
else:
|
||||
await acompletion_with_retries(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"gm": "vibe", "role": "user"}],
|
||||
num_retries=3,
|
||||
original_function=mock_completion,
|
||||
)
|
||||
mock_completion.assert_called_once()
|
||||
assert mock_completion.call_args.kwargs["num_retries"] == 0
|
||||
assert mock_completion.call_args.kwargs["max_retries"] == 0
|
||||
|
||||
|
||||
# ==================== Responses API Retry Tests ====================
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_with_retries(sync_mode):
|
||||
"""
|
||||
Test that responses() and aresponses() properly handle num_retries parameter.
|
||||
If responses_with_retries is called with num_retries=3, and max_retries=0,
|
||||
then litellm.responses should receive num_retries=0, max_retries=0
|
||||
"""
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
|
||||
if sync_mode:
|
||||
target_function = "responses"
|
||||
retry_function = responses_with_retries
|
||||
else:
|
||||
target_function = "aresponses"
|
||||
retry_function = aresponses_with_retries
|
||||
|
||||
# Mock the responses/aresponses function
|
||||
with patch(
|
||||
"litellm.responses.main.responses"
|
||||
if sync_mode
|
||||
else "litellm.responses.main.aresponses"
|
||||
) as mock_responses:
|
||||
if sync_mode:
|
||||
mock_responses.return_value = MagicMock()
|
||||
retry_function(
|
||||
model="gpt-4o",
|
||||
input="Hello, what's the weather?",
|
||||
num_retries=3,
|
||||
original_function=mock_responses,
|
||||
)
|
||||
else:
|
||||
mock_responses.return_value = AsyncMock()
|
||||
await retry_function(
|
||||
model="gpt-4o",
|
||||
input="Hello, what's the weather?",
|
||||
num_retries=3,
|
||||
original_function=mock_responses,
|
||||
)
|
||||
|
||||
mock_responses.assert_called_once()
|
||||
assert mock_responses.call_args.kwargs["num_retries"] == 0
|
||||
assert mock_responses.call_args.kwargs["max_retries"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
async def test_responses_retry_on_auth_error(sync_mode):
|
||||
"""
|
||||
Test that responses API actually retries when encountering authentication errors.
|
||||
This validates that the @client decorator properly handles responses/aresponses retries.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
num_retries = 2
|
||||
|
||||
# Mock the responses/aresponses to raise an authentication error
|
||||
if sync_mode:
|
||||
with patch.object(litellm, "responses_with_retries") as mock_retry:
|
||||
mock_retry.return_value = None
|
||||
try:
|
||||
responses(
|
||||
model="gpt-4o",
|
||||
input="Test input",
|
||||
num_retries=num_retries,
|
||||
api_key="sk-invalid-key-12345",
|
||||
)
|
||||
except Exception:
|
||||
pass # Expected to fail with invalid key
|
||||
|
||||
# Check if retry function was called (means @client decorator triggered retry)
|
||||
if mock_retry.called:
|
||||
assert mock_retry.call_args.kwargs.get("num_retries") == num_retries
|
||||
else:
|
||||
with patch.object(litellm, "aresponses_with_retries") as mock_retry:
|
||||
mock_retry.return_value = None
|
||||
try:
|
||||
await aresponses(
|
||||
model="gpt-4o",
|
||||
input="Test input",
|
||||
num_retries=num_retries,
|
||||
api_key="sk-invalid-key-12345",
|
||||
)
|
||||
except Exception:
|
||||
pass # Expected to fail with invalid key
|
||||
|
||||
# Check if retry function was called (means @client decorator triggered retry)
|
||||
if mock_retry.called:
|
||||
assert mock_retry.call_args.kwargs.get("num_retries") == num_retries
|
||||
|
|
|
|||
|
|
@ -41,16 +41,12 @@ exception_models = [
|
|||
]
|
||||
|
||||
|
||||
|
||||
|
||||
# Test 1: Context Window Errors
|
||||
|
||||
|
||||
models = ["command-nightly"]
|
||||
|
||||
|
||||
|
||||
|
||||
# for model in litellm.models_by_provider["bedrock"]:
|
||||
# test_context_window(model=model)
|
||||
# test_context_window(model="chat-bison")
|
||||
|
|
@ -185,17 +181,6 @@ def test_completion_azure_exception():
|
|||
# test_completion_azure_exception()
|
||||
|
||||
|
||||
def test_azure_embedding_exceptions():
|
||||
# CRUCIAL Test - Ensures our exceptions are readable and not overly complicated. some users have complained exceptions will randomly have another exception raised in our exception mapping
|
||||
with pytest.raises(Exception, match="Mock error") as exc_info:
|
||||
litellm.embedding(
|
||||
model="azure/text-embedding-ada-002",
|
||||
input="hello",
|
||||
mock_response="error",
|
||||
)
|
||||
assert str(exc_info.value) == "Mock error"
|
||||
|
||||
|
||||
async def asynctest_completion_azure_exception():
|
||||
try:
|
||||
import openai
|
||||
|
|
@ -315,8 +300,6 @@ def test_completion_openai_exception():
|
|||
# test_completion_openai_exception()
|
||||
|
||||
|
||||
|
||||
|
||||
def test_completion_mistral_exception():
|
||||
# test if mistral/mistral-tiny raises openai.AuthenticationError
|
||||
try:
|
||||
|
|
@ -343,30 +326,6 @@ def test_completion_mistral_exception():
|
|||
# test_completion_mistral_exception()
|
||||
|
||||
|
||||
def test_completion_bedrock_invalid_role_exception():
|
||||
"""
|
||||
Test if litellm raises a BadRequestError for an invalid role on Bedrock
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
completion(
|
||||
model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
messages=[{"role": "very-bad-role", "content": "hello"}],
|
||||
)
|
||||
|
||||
# This is important - We we previously returning a poorly formatted error string. Which was
|
||||
# litellm.BadRequestError: litellm.BadRequestError: Invalid Message passed in {'role': 'very-bad-role', 'content': 'hello'}
|
||||
assert (
|
||||
str(exc_info.value)
|
||||
== "litellm.BadRequestError: Invalid Message passed in {'role': 'very-bad-role', 'content': 'hello'}"
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
# test_content_policy_exceptionimage_generation_openai()
|
||||
|
||||
|
||||
def test_content_policy_violation_error_streaming():
|
||||
"""
|
||||
Production Test.
|
||||
|
|
@ -434,8 +393,6 @@ def test_content_policy_violation_error_streaming():
|
|||
asyncio.run(test_get_error())
|
||||
|
||||
|
||||
|
||||
|
||||
# test_completion_perplexity_exception_on_openai_client()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -31,7 +31,6 @@ _MIGRATED_FILES = (
|
|||
"tests/llm_translation/test_triton.py",
|
||||
"tests/local_testing/test_router.py",
|
||||
"tests/local_testing/test_router_custom_routing.py",
|
||||
"tests/local_testing/test_router_fallback_handlers.py",
|
||||
"tests/local_testing/test_router_fallbacks.py",
|
||||
"tests/local_testing/test_secret_detect_hook.py",
|
||||
"tests/local_testing/test_lowest_latency_routing.py",
|
||||
|
|
|
|||
|
|
@ -3,19 +3,15 @@
|
|||
|
||||
import asyncio
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
import copy
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
|
||||
|
|
@ -25,127 +21,6 @@ from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
|||
### UNIT TESTS FOR LATENCY ROUTING ###
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_latency_memory_leak(sync_mode):
|
||||
"""
|
||||
Test to make sure there's no memory leak caused by lowest latency routing
|
||||
|
||||
- make 10 calls -> check memory
|
||||
- make 11th call -> no change in memory
|
||||
"""
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "1234"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
time.sleep(5)
|
||||
end_time = time.time()
|
||||
for _ in range(10):
|
||||
if sync_mode:
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
latency_key = f"{model_group}_map"
|
||||
cache_value = copy.deepcopy(
|
||||
test_cache.get_cache(key=latency_key)
|
||||
) # MAKE SURE NO MEMORY LEAK IN CACHING OBJECT
|
||||
|
||||
if sync_mode:
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
new_cache_value = test_cache.get_cache(key=latency_key)
|
||||
# Assert that the size of the cache doesn't grow unreasonably
|
||||
assert get_size(new_cache_value) <= get_size(
|
||||
cache_value
|
||||
), f"Memory leak detected in function call! new_cache size={get_size(new_cache_value)}, old cache size={get_size(cache_value)}"
|
||||
|
||||
|
||||
def get_size(obj, seen=None):
|
||||
# From https://goshippo.com/blog/measure-real-size-any-python-object/
|
||||
# Recursively finds size of objects
|
||||
size = sys.getsizeof(obj)
|
||||
if seen is None:
|
||||
seen = set()
|
||||
obj_id = id(obj)
|
||||
if obj_id in seen:
|
||||
return 0
|
||||
seen.add(obj_id)
|
||||
if isinstance(obj, dict):
|
||||
size += sum([get_size(v, seen) for v in obj.values()])
|
||||
size += sum([get_size(k, seen) for k in obj.keys()])
|
||||
elif hasattr(obj, "__dict__"):
|
||||
size += get_size(obj.__dict__, seen)
|
||||
elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes, bytearray)):
|
||||
size += sum([get_size(i, seen) for i in obj])
|
||||
return size
|
||||
|
||||
|
||||
def test_latency_updated():
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "1234"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
time.sleep(5)
|
||||
end_time = time.time()
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
latency_key = f"{model_group}_map"
|
||||
assert (
|
||||
end_time - start_time
|
||||
== test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0]
|
||||
)
|
||||
|
||||
|
||||
# test_tpm_rpm_updated()
|
||||
|
||||
|
||||
def test_latency_updated_custom_ttl():
|
||||
"""
|
||||
Invalidate the cached request.
|
||||
|
|
@ -186,79 +61,6 @@ def test_latency_updated_custom_ttl():
|
|||
assert test_cache.get_cache(key=latency_key) is None
|
||||
|
||||
|
||||
def test_get_available_deployments():
|
||||
test_cache = DualCache()
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "azure/gpt-4.1-mini"},
|
||||
"model_info": {"id": "1234"},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "azure/gpt-4.1-mini"},
|
||||
"model_info": {"id": "5678"},
|
||||
},
|
||||
]
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
|
||||
model_group = "gpt-3.5-turbo"
|
||||
## DEPLOYMENT 1 ##
|
||||
deployment_id = "1234"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
time.sleep(3)
|
||||
end_time = time.time()
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
## DEPLOYMENT 2 ##
|
||||
deployment_id = "5678"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 20}}
|
||||
time.sleep(2)
|
||||
end_time = time.time()
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
## CHECK WHAT'S SELECTED ##
|
||||
print(
|
||||
lowest_latency_logger.get_available_deployments(
|
||||
model_group=model_group, healthy_deployments=model_list
|
||||
)
|
||||
)
|
||||
assert (
|
||||
lowest_latency_logger.get_available_deployments(
|
||||
model_group=model_group, healthy_deployments=model_list
|
||||
)["model_info"]["id"]
|
||||
== "5678"
|
||||
)
|
||||
|
||||
|
||||
async def _deploy(lowest_latency_logger, deployment_id, tokens_used, duration):
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
|
|
@ -417,91 +219,6 @@ def test_get_available_endpoints_tpm_rpm_check(ans_rpm):
|
|||
assert (d_ans and d_ans["model_info"]["id"]) == ans
|
||||
|
||||
|
||||
def test_router_get_available_deployments():
|
||||
"""
|
||||
Test if routers 'get_available_deployments' returns the fastest deployment
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-turbo",
|
||||
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
|
||||
"api_base": "https://openai-france-1234.openai.azure.com",
|
||||
"rpm": 1440,
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-35-turbo",
|
||||
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
|
||||
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
|
||||
"rpm": 6,
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
set_verbose=False,
|
||||
num_retries=3,
|
||||
) # type: ignore
|
||||
|
||||
## DEPLOYMENT 1 ##
|
||||
deployment_id = 1
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
time.sleep(3)
|
||||
end_time = time.time()
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
## DEPLOYMENT 2 ##
|
||||
deployment_id = 2
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 20}}
|
||||
time.sleep(2)
|
||||
end_time = time.time()
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
## CHECK WHAT'S SELECTED ##
|
||||
# print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
assert (
|
||||
router.get_available_deployment(model="azure-model")["model_info"]["id"] == "2"
|
||||
)
|
||||
|
||||
|
||||
# test_router_get_available_deployments()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_completion_streaming():
|
||||
messages = [
|
||||
|
|
@ -719,486 +436,3 @@ async def test_lowest_latency_routing_first_pick():
|
|||
|
||||
# assert that len(deployments) >1
|
||||
assert len(deployments) > 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("buffer", [0, 1])
|
||||
@pytest.mark.asyncio
|
||||
async def test_lowest_latency_routing_buffer(buffer):
|
||||
"""
|
||||
Allow shuffling calls within a certain latency buffer
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-turbo",
|
||||
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
|
||||
"api_base": "https://openai-france-1234.openai.azure.com",
|
||||
"rpm": 1440,
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-35-turbo",
|
||||
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
|
||||
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
|
||||
"rpm": 6,
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
set_verbose=False,
|
||||
num_retries=3,
|
||||
routing_strategy_args={"lowest_latency_buffer": buffer},
|
||||
) # type: ignore
|
||||
|
||||
## DEPLOYMENT 1 ##
|
||||
deployment_id = 1
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
time.sleep(3)
|
||||
end_time = time.time()
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
## DEPLOYMENT 2 ##
|
||||
deployment_id = 2
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
}
|
||||
}
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 20}}
|
||||
time.sleep(2)
|
||||
end_time = time.time()
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
## CHECK WHAT'S SELECTED ##
|
||||
# print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
|
||||
selected_deployments = {}
|
||||
for _ in range(50):
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
selected_deployments[
|
||||
router.get_available_deployment(model="azure-model")["model_info"]["id"]
|
||||
] = 1
|
||||
|
||||
if buffer == 0:
|
||||
assert len(selected_deployments.keys()) == 1
|
||||
else:
|
||||
assert len(selected_deployments.keys()) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_lowest_latency_routing_time_to_first_token(sync_mode):
|
||||
"""
|
||||
If a deployment has
|
||||
- a fast time to first token
|
||||
- slow latency/output token
|
||||
|
||||
test if:
|
||||
- for streaming, the deployment with fastest time to first token is picked
|
||||
- for non-streaming, fastest overall deployment is picked
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-turbo",
|
||||
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
|
||||
"api_base": "https://openai-france-1234.openai.azure.com",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-35-turbo",
|
||||
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
|
||||
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
set_verbose=False,
|
||||
num_retries=3,
|
||||
) # type: ignore
|
||||
## DEPLOYMENT 1 ##
|
||||
deployment_id = 1
|
||||
start_time = datetime.now()
|
||||
one_second_later = start_time + timedelta(seconds=1)
|
||||
|
||||
# Compute 3 seconds after the current time
|
||||
three_seconds_later = start_time + timedelta(seconds=3)
|
||||
four_seconds_later = start_time + timedelta(seconds=4)
|
||||
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
"stream": True,
|
||||
"completion_start_time": one_second_later,
|
||||
}
|
||||
|
||||
response_obj = litellm.ModelResponse(
|
||||
usage=litellm.Usage(completion_tokens=50, total_tokens=50)
|
||||
)
|
||||
end_time = four_seconds_later
|
||||
|
||||
if sync_mode:
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await router.lowestlatency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
## DEPLOYMENT 2 ##
|
||||
deployment_id = 2
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
"stream": True,
|
||||
"completion_start_time": three_seconds_later,
|
||||
}
|
||||
response_obj = litellm.ModelResponse(
|
||||
usage=litellm.Usage(completion_tokens=50, total_tokens=50)
|
||||
)
|
||||
end_time = three_seconds_later
|
||||
if sync_mode:
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await router.lowestlatency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
"""
|
||||
TESTING
|
||||
|
||||
- expect deployment 1 to be picked for streaming
|
||||
- expect deployment 2 to be picked for non-streaming
|
||||
"""
|
||||
# print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
|
||||
selected_deployments = {}
|
||||
for _ in range(3):
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
## for non-streaming
|
||||
selected_deployments[
|
||||
router.get_available_deployment(model="azure-model")["model_info"]["id"]
|
||||
] = 1
|
||||
|
||||
assert len(selected_deployments.keys()) == 1
|
||||
assert "2" in list(selected_deployments.keys())
|
||||
|
||||
selected_deployments = {}
|
||||
for _ in range(50):
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
## for non-streaming
|
||||
selected_deployments[
|
||||
router.get_available_deployment(
|
||||
model="azure-model", request_kwargs={"stream": True}
|
||||
)["model_info"]["id"]
|
||||
] = 1
|
||||
|
||||
assert len(selected_deployments.keys()) == 1
|
||||
assert "1" in list(selected_deployments.keys())
|
||||
|
||||
|
||||
def test_latency_list_trimming_discards_oldest_entry():
|
||||
"""
|
||||
When the latency list reaches max_latency_list_size, the oldest entry is
|
||||
discarded to make room for new entries. The newest entry is appended at
|
||||
the end of the list.
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
# With 1 completion token, the logged latency value equals the raw
|
||||
# response time, so we can use distinct, identifiable values.
|
||||
latencies_to_add = []
|
||||
for i in range(max_size + 1): # One more than max to trigger trimming
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
expected_latency = float(i + 1) # 1.0, 2.0, 3.0, 4.0
|
||||
end_time = start_time + expected_latency
|
||||
latencies_to_add.append(expected_latency)
|
||||
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = test_cache.get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert (
|
||||
len(latency_list) == max_size
|
||||
), f"Expected {max_size} entries, got {len(latency_list)}"
|
||||
|
||||
newest_latency = latencies_to_add[-1] # 4.0
|
||||
oldest_latency = latencies_to_add[0] # 1.0
|
||||
tolerance = 0.1
|
||||
|
||||
# Newest entry is at the end of the list.
|
||||
assert (
|
||||
abs(latency_list[-1] - newest_latency) < tolerance
|
||||
), f"Newest latency {newest_latency} should be at end, got {latency_list[-1]}"
|
||||
|
||||
# Oldest entry is no longer in the list.
|
||||
for latency in latency_list:
|
||||
assert (
|
||||
abs(latency - oldest_latency) > tolerance
|
||||
), f"Oldest latency {oldest_latency} should have been discarded, found {latency}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_latency_list_trimming_discards_oldest_entry_async():
|
||||
"""
|
||||
Async counterpart: the oldest entry is discarded when the latency list is
|
||||
trimmed.
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
latencies_to_add = []
|
||||
for i in range(max_size + 1):
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
expected_latency = float(i + 1)
|
||||
end_time = start_time + expected_latency
|
||||
latencies_to_add.append(expected_latency)
|
||||
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = await test_cache.async_get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size
|
||||
|
||||
newest_latency = latencies_to_add[-1]
|
||||
oldest_latency = latencies_to_add[0]
|
||||
tolerance = 0.1
|
||||
|
||||
assert (
|
||||
abs(latency_list[-1] - newest_latency) < tolerance
|
||||
), f"Newest latency {newest_latency} should be at end of list"
|
||||
|
||||
for latency in latency_list:
|
||||
assert (
|
||||
abs(latency - oldest_latency) > tolerance
|
||||
), f"Oldest latency {oldest_latency} should have been discarded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_penalty_discards_oldest_entry():
|
||||
"""
|
||||
Timeout penalties (1000.0) are appended to the latency list and, when the
|
||||
list is full, the oldest entry is discarded.
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
# Fill the list with max_size normal latency entries first.
|
||||
for i in range(max_size):
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
end_time = start_time + float(i + 1)
|
||||
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
# Trigger a timeout failure: this appends 1000.0 and should discard the
|
||||
# oldest normal entry (1.0).
|
||||
timeout_kwargs = {
|
||||
**kwargs,
|
||||
"exception": litellm.Timeout(
|
||||
message="Request timed out", model="test-model", llm_provider="test"
|
||||
),
|
||||
}
|
||||
|
||||
await lowest_latency_logger.async_log_failure_event(
|
||||
kwargs=timeout_kwargs,
|
||||
response_obj=None,
|
||||
start_time=time.time(),
|
||||
end_time=time.time() + 30,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = await test_cache.async_get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size
|
||||
|
||||
# Timeout penalty is the newest entry.
|
||||
assert (
|
||||
latency_list[-1] == 1000.0
|
||||
), f"Timeout penalty should be at end of list, got {latency_list[-1]}"
|
||||
|
||||
# Oldest normal entry (1.0) has been discarded.
|
||||
tolerance = 0.1
|
||||
for latency in latency_list[:-1]:
|
||||
assert (
|
||||
abs(latency - 1.0) > tolerance
|
||||
), f"Oldest latency 1.0 should have been discarded, found {latency}"
|
||||
|
||||
|
||||
def test_list_order_preserved_after_multiple_trims():
|
||||
"""
|
||||
After many trims, the list still holds the most recent `max_size` entries
|
||||
in insertion order (oldest at index 0, newest at index -1).
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
# Add 10 entries (7 more than max) to trigger multiple trims.
|
||||
all_latencies = []
|
||||
for i in range(10):
|
||||
start_time = time.time()
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
expected_latency = float(i + 1)
|
||||
end_time = start_time + expected_latency
|
||||
all_latencies.append(expected_latency)
|
||||
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = test_cache.get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size
|
||||
|
||||
# After inserting 1..10 with max_size=3, the list should be [8, 9, 10].
|
||||
expected_remaining = all_latencies[-max_size:]
|
||||
tolerance = 0.1
|
||||
|
||||
for i, expected in enumerate(expected_remaining):
|
||||
assert (
|
||||
abs(latency_list[i] - expected) < tolerance
|
||||
), f"At index {i}, expected ~{expected}, got {latency_list[i]}"
|
||||
|
|
|
|||
|
|
@ -75,31 +75,6 @@ class CustomRoutingStrategy(CustomRoutingStrategyBase):
|
|||
pass
|
||||
|
||||
|
||||
def test_reset_custom_routing_strategy():
|
||||
"""
|
||||
Setting a custom routing strategy installs instance-level overrides for
|
||||
get_available_deployment / async_get_available_deployment. Re-initializing the
|
||||
routing strategy must clear them so the class implementations are used again.
|
||||
"""
|
||||
router = _create_router()
|
||||
router.set_custom_routing_strategy(CustomRoutingStrategy(router))
|
||||
|
||||
assert "get_available_deployment" in router.__dict__
|
||||
assert "async_get_available_deployment" in router.__dict__
|
||||
|
||||
router._reset_custom_routing_strategy()
|
||||
|
||||
assert "get_available_deployment" not in router.__dict__
|
||||
assert "async_get_available_deployment" not in router.__dict__
|
||||
assert (
|
||||
router.async_get_available_deployment.__func__
|
||||
is Router.async_get_available_deployment
|
||||
)
|
||||
|
||||
# idempotent: resetting again when nothing is overridden must not raise
|
||||
router._reset_custom_routing_strategy()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_routing():
|
||||
litellm.set_verbose = True
|
||||
|
|
|
|||
|
|
@ -1,106 +0,0 @@
|
|||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import logging
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
||||
# this tests debug logs from litellm router and litellm proxy server
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger, verbose_router_logger
|
||||
from litellm.llms.custom_httpx.async_client_cleanup import close_litellm_async_clients
|
||||
|
||||
|
||||
# this tests debug logs from litellm router and litellm proxy server
|
||||
def test_async_fallbacks(caplog):
|
||||
# THIS IS A PROD TEST - DO NOT DELETE THIS. Used for testing if litellm proxy verbose logs are human readable
|
||||
litellm.set_verbose = False
|
||||
litellm.success_callback = []
|
||||
litellm.failure_callback = []
|
||||
verbose_router_logger.setLevel(level=logging.INFO)
|
||||
verbose_logger.setLevel(logging.CRITICAL + 1)
|
||||
verbose_proxy_logger.setLevel(logging.CRITICAL + 1)
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure/gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-4.1-mini",
|
||||
"api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
"mock_response": "Hello world",
|
||||
},
|
||||
"tpm": 240000,
|
||||
"rpm": 1800,
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "bad-key",
|
||||
},
|
||||
"tpm": 1000000,
|
||||
"rpm": 9000,
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
fallbacks=[{"gpt-3.5-turbo": ["azure/gpt-3.5-turbo"]}],
|
||||
num_retries=1,
|
||||
)
|
||||
|
||||
user_message = "Hello, how are you?"
|
||||
messages = [{"content": user_message, "role": "user"}]
|
||||
|
||||
async def _make_request():
|
||||
try:
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo", messages=messages, max_tokens=1
|
||||
)
|
||||
router.reset()
|
||||
except litellm.Timeout:
|
||||
pass
|
||||
except Exception as e:
|
||||
pytest.fail(f"An exception occurred: {e}")
|
||||
finally:
|
||||
router.reset()
|
||||
# Close cached aiohttp/httpx clients before the event loop ends
|
||||
# to prevent "Unclosed client session" / "Unclosed connector" warnings.
|
||||
await close_litellm_async_clients()
|
||||
|
||||
asyncio.run(_make_request())
|
||||
captured_logs = [rec.message for rec in caplog.records]
|
||||
|
||||
# on circle ci the captured logs get async cleanup noise from the gc (leaked
|
||||
# task warnings, plus aiohttp "Unclosed client session"/"Unclosed connector"
|
||||
# warnings from cached clients other router tests evicted) - filter it out
|
||||
captured_logs = [
|
||||
log
|
||||
for log in captured_logs
|
||||
if "Task exception was never retrieved" not in log
|
||||
and "Task was destroyed but it is pending" not in log
|
||||
and "get_available_deployment" not in log
|
||||
and "Selected deployment for model" not in log
|
||||
and "in the Langfuse queue" not in log
|
||||
and "Unclosed client session" not in log
|
||||
and "Unclosed connector" not in log
|
||||
]
|
||||
|
||||
print("\n Captured caplog records - ", captured_logs)
|
||||
|
||||
# Define the expected log messages
|
||||
# - error request, falling back notice, success notice
|
||||
expected_logs = [
|
||||
"Falling back to model_group = azure/gpt-3.5-turbo",
|
||||
"litellm.acompletion(model=azure/gpt-4.1-mini)\x1b[32m 200 OK\x1b[0m",
|
||||
"Successful fallback b/w models.",
|
||||
]
|
||||
|
||||
# Assert that the captured logs match the expected log messages
|
||||
assert captured_logs[-3:] == expected_logs
|
||||
|
|
@ -1,87 +0,0 @@
|
|||
import os
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
run_async_fallback,
|
||||
)
|
||||
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
||||
|
||||
# Helper function to create a Router instance
|
||||
|
||||
def create_test_router_2():
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4",
|
||||
"api_key": "very-fake-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "fake-openai-endpoint-2",
|
||||
"litellm_params": {
|
||||
"model": "openai/fake-openai-endpoint-2",
|
||||
"api_key": "working-key-since-this-is-fake-endpoint",
|
||||
"api_base": FAKE_OPENAI_API_BASE,
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("function_name", ["_acompletion", "_atext_completion"])
|
||||
async def test_multiple_fallbacks(function_name):
|
||||
"""
|
||||
Tests that if multiple fallbacks passed:
|
||||
- fallback 1 = bad configured deployment / failing endpoint
|
||||
- fallback 2 = working deployment / working endpoint
|
||||
|
||||
Assert that:
|
||||
- a success response is received from the working endpoint (fallback 2)
|
||||
"""
|
||||
router_2 = create_test_router_2()
|
||||
original_function = getattr(router_2, function_name)
|
||||
|
||||
fallback_model_group = ["gpt-4", "fake-openai-endpoint-2"]
|
||||
original_model_group = "gpt-3.5-turbo"
|
||||
original_exception = Exception("Simulated error")
|
||||
|
||||
request_kwargs: dict[str, Any] = {
|
||||
"metadata": {"previous_models": ["gpt-3.5-turbo"]}
|
||||
}
|
||||
|
||||
if function_name == "_aembedding":
|
||||
request_kwargs["input"] = "hello this is a test for run_async_fallback"
|
||||
elif function_name == "_atext_completion":
|
||||
request_kwargs["prompt"] = "hello this is a test for run_async_fallback"
|
||||
elif function_name == "_acompletion":
|
||||
request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}]
|
||||
|
||||
result = await run_async_fallback(
|
||||
litellm_router=router_2,
|
||||
original_function=original_function,
|
||||
num_retries=1,
|
||||
fallback_model_group=fallback_model_group,
|
||||
original_model_group=original_model_group,
|
||||
original_exception=original_exception,
|
||||
max_fallbacks=5,
|
||||
fallback_depth=0,
|
||||
**request_kwargs,
|
||||
)
|
||||
|
||||
print(result)
|
||||
|
||||
print(result._hidden_params)
|
||||
|
||||
assert result._hidden_params["api_base"] == FAKE_OPENAI_API_BASE
|
||||
0
tests/unit/batch_completion/__init__.py
Normal file
0
tests/unit/batch_completion/__init__.py
Normal file
|
|
@ -1,27 +1,37 @@
|
|||
"""https://github.com/BerriAI/litellm/pull/3397/commits/a7ec1772b1457594d3af48cdcb0a382279b841c7#diff-44852387ceb00aade916d6b314dfd5d180499e54f35209ae9c07179febe08b4b."""
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
"""Test batch_completion's return_exceptions."""
|
||||
import litellm
|
||||
|
||||
msg1 = [{"role": "user", "content": "hi 1"}]
|
||||
msg2 = [{"role": "user", "content": "hi 2"}]
|
||||
|
||||
|
||||
def test_batch_completion_return_exceptions_true():
|
||||
def test_batch_completion_return_exceptions_true(respx_mock: respx.MockRouter):
|
||||
"""Test batch_completion's return_exceptions.
|
||||
|
||||
With an invalid API key, we expect an error to be returned rather than raised.
|
||||
The error type may be AuthenticationError (from API) or InternalServerError
|
||||
(from connection issues), depending on network conditions.
|
||||
"""
|
||||
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
401,
|
||||
json={
|
||||
"error": {
|
||||
"message": "Incorrect API key provided: sk_xxx.",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
res = litellm.batch_completion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[msg1, msg2],
|
||||
api_key="sk_xxx", # deliberately set invalid key
|
||||
api_key="sk_xxx",
|
||||
)
|
||||
|
||||
# batch_completion should return exceptions rather than raise them
|
||||
# Accept either AuthenticationError (API rejected key) or InternalServerError (network issues)
|
||||
assert isinstance(
|
||||
res[0],
|
||||
(
|
||||
|
|
@ -7025,3 +7025,20 @@ def test_has_tool_with_name_anthropic_shape_without_type_field():
|
|||
def test_has_tool_with_name_not_a_list():
|
||||
assert not has_tool_with_name(None, "my_tool")
|
||||
assert not has_tool_with_name("not a list", "my_tool")
|
||||
|
||||
|
||||
def test_completion_bedrock_invalid_role_exception(monkeypatch):
|
||||
"""
|
||||
Test if litellm raises a BadRequestError for an invalid role on Bedrock
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "set_verbose", True)
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
litellm.completion(
|
||||
model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
messages=[{"role": "very-bad-role", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
str(exc_info.value)
|
||||
== "litellm.BadRequestError: Invalid Message passed in {'role': 'very-bad-role', 'content': 'hello'}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,9 @@
|
|||
# the ModelResponse branch, so the raw timedelta used to leak into the
|
||||
# latency list and break the Redis cache sync). Issue #33169.
|
||||
|
||||
import copy
|
||||
import json
|
||||
import sys
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
|
@ -476,3 +478,712 @@ async def test_runtime_routing_strategy_args_update_is_a_noop_without_a_selector
|
|||
|
||||
assert router.routing_strategy_args == {"ttl": 5}
|
||||
assert await _pick_streaming(router) == FAST_TTFT_ID
|
||||
|
||||
|
||||
FROZEN_NOW = datetime(2026, 1, 15, 12, 30, 15)
|
||||
FROZEN_EPOCH = FROZEN_NOW.timestamp()
|
||||
|
||||
|
||||
class _FrozenDatetime(datetime):
|
||||
@classmethod
|
||||
def now(cls, tz=None):
|
||||
return FROZEN_NOW if tz is None else FROZEN_NOW.astimezone(tz)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def frozen_latency_clock(monkeypatch):
|
||||
monkeypatch.setattr("litellm.router_strategy.lowest_latency.datetime", _FrozenDatetime)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
async def test_latency_memory_leak(sync_mode):
|
||||
"""
|
||||
Test to make sure there's no memory leak caused by lowest latency routing
|
||||
|
||||
- make 10 calls -> check memory
|
||||
- make 11th call -> no change in memory
|
||||
"""
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "1234"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
end_time = start_time + 5
|
||||
for _ in range(10):
|
||||
if sync_mode:
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
latency_key = f"{model_group}_map"
|
||||
cache_value = copy.deepcopy(test_cache.get_cache(key=latency_key))
|
||||
|
||||
if sync_mode:
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
new_cache_value = test_cache.get_cache(key=latency_key)
|
||||
assert get_size(new_cache_value) <= get_size(cache_value), (
|
||||
f"Memory leak detected in function call! new_cache size={get_size(new_cache_value)}, old cache size={get_size(cache_value)}"
|
||||
)
|
||||
|
||||
|
||||
def get_size(obj, seen=None):
|
||||
size = sys.getsizeof(obj)
|
||||
if seen is None:
|
||||
seen = set()
|
||||
obj_id = id(obj)
|
||||
if obj_id in seen:
|
||||
return 0
|
||||
seen.add(obj_id)
|
||||
if isinstance(obj, dict):
|
||||
size += sum([get_size(v, seen) for v in obj.values()])
|
||||
size += sum([get_size(k, seen) for k in obj.keys()])
|
||||
elif hasattr(obj, "__dict__"):
|
||||
size += get_size(obj.__dict__, seen)
|
||||
elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes, bytearray)):
|
||||
size += sum([get_size(i, seen) for i in obj])
|
||||
return size
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
def test_latency_updated():
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "1234"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
end_time = start_time + 5
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
latency_key = f"{model_group}_map"
|
||||
assert end_time - start_time == test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
def test_get_available_deployments():
|
||||
test_cache = DualCache()
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "azure/gpt-4.1-mini"},
|
||||
"model_info": {"id": "1234"},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "azure/gpt-4.1-mini"},
|
||||
"model_info": {"id": "5678"},
|
||||
},
|
||||
]
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "1234"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
end_time = start_time + 3
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
deployment_id = "5678"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 20}}
|
||||
end_time = start_time + 2
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
print(lowest_latency_logger.get_available_deployments(model_group=model_group, healthy_deployments=model_list))
|
||||
assert (
|
||||
lowest_latency_logger.get_available_deployments(model_group=model_group, healthy_deployments=model_list)[
|
||||
"model_info"
|
||||
]["id"]
|
||||
== "5678"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
def test_router_get_available_deployments():
|
||||
"""
|
||||
Test if routers 'get_available_deployments' returns the fastest deployment
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-turbo",
|
||||
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
|
||||
"api_base": "https://openai-france-1234.openai.azure.com",
|
||||
"rpm": 1440,
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-35-turbo",
|
||||
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
|
||||
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
|
||||
"rpm": 6,
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
set_verbose=False,
|
||||
num_retries=3,
|
||||
)
|
||||
|
||||
deployment_id = 1
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
end_time = start_time + 3
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
deployment_id = 2
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 20}}
|
||||
end_time = start_time + 2
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
assert router.get_available_deployment(model="azure-model")["model_info"]["id"] == "2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("buffer", [0, 1])
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
async def test_lowest_latency_routing_buffer(buffer):
|
||||
"""
|
||||
Allow shuffling calls within a certain latency buffer
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-turbo",
|
||||
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
|
||||
"api_base": "https://openai-france-1234.openai.azure.com",
|
||||
"rpm": 1440,
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-35-turbo",
|
||||
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
|
||||
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
|
||||
"rpm": 6,
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
set_verbose=False,
|
||||
num_retries=3,
|
||||
routing_strategy_args={"lowest_latency_buffer": buffer},
|
||||
)
|
||||
|
||||
deployment_id = 1
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 50}}
|
||||
end_time = start_time + 3
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
deployment_id = 2
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
}
|
||||
}
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 20}}
|
||||
end_time = start_time + 2
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
selected_deployments = {}
|
||||
for _ in range(50):
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
selected_deployments[router.get_available_deployment(model="azure-model")["model_info"]["id"]] = 1
|
||||
|
||||
if buffer == 0:
|
||||
assert len(selected_deployments.keys()) == 1
|
||||
else:
|
||||
assert len(selected_deployments.keys()) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
async def test_lowest_latency_routing_time_to_first_token(sync_mode):
|
||||
"""
|
||||
If a deployment has
|
||||
- a fast time to first token
|
||||
- slow latency/output token
|
||||
|
||||
test if:
|
||||
- for streaming, the deployment with fastest time to first token is picked
|
||||
- for non-streaming, fastest overall deployment is picked
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-turbo",
|
||||
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
|
||||
"api_base": "https://openai-france-1234.openai.azure.com",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-35-turbo",
|
||||
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
|
||||
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
set_verbose=False,
|
||||
num_retries=3,
|
||||
)
|
||||
deployment_id = 1
|
||||
start_time = FROZEN_NOW
|
||||
one_second_later = start_time + timedelta(seconds=1)
|
||||
|
||||
three_seconds_later = start_time + timedelta(seconds=3)
|
||||
four_seconds_later = start_time + timedelta(seconds=4)
|
||||
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 1},
|
||||
},
|
||||
"stream": True,
|
||||
"completion_start_time": one_second_later,
|
||||
}
|
||||
|
||||
response_obj = litellm.ModelResponse(usage=litellm.Usage(completion_tokens=50, total_tokens=50))
|
||||
end_time = four_seconds_later
|
||||
|
||||
if sync_mode:
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await router.lowestlatency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
deployment_id = 2
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "azure-model",
|
||||
},
|
||||
"model_info": {"id": 2},
|
||||
},
|
||||
"stream": True,
|
||||
"completion_start_time": three_seconds_later,
|
||||
}
|
||||
response_obj = litellm.ModelResponse(usage=litellm.Usage(completion_tokens=50, total_tokens=50))
|
||||
end_time = three_seconds_later
|
||||
if sync_mode:
|
||||
router.lowestlatency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
else:
|
||||
await router.lowestlatency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
"""
|
||||
TESTING
|
||||
|
||||
- expect deployment 1 to be picked for streaming
|
||||
- expect deployment 2 to be picked for non-streaming
|
||||
"""
|
||||
selected_deployments = {}
|
||||
for _ in range(3):
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
selected_deployments[router.get_available_deployment(model="azure-model")["model_info"]["id"]] = 1
|
||||
|
||||
assert len(selected_deployments.keys()) == 1
|
||||
assert "2" in list(selected_deployments.keys())
|
||||
|
||||
selected_deployments = {}
|
||||
for _ in range(50):
|
||||
print(router.get_available_deployment(model="azure-model"))
|
||||
selected_deployments[
|
||||
router.get_available_deployment(model="azure-model", request_kwargs={"stream": True})["model_info"]["id"]
|
||||
] = 1
|
||||
|
||||
assert len(selected_deployments.keys()) == 1
|
||||
assert "1" in list(selected_deployments.keys())
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
def test_latency_list_trimming_discards_oldest_entry():
|
||||
"""
|
||||
When the latency list reaches max_latency_list_size, the oldest entry is
|
||||
discarded to make room for new entries. The newest entry is appended at
|
||||
the end of the list.
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
latencies_to_add = []
|
||||
for i in range(max_size + 1):
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
expected_latency = float(i + 1)
|
||||
end_time = start_time + expected_latency
|
||||
latencies_to_add.append(expected_latency)
|
||||
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = test_cache.get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size, f"Expected {max_size} entries, got {len(latency_list)}"
|
||||
|
||||
newest_latency = latencies_to_add[-1]
|
||||
oldest_latency = latencies_to_add[0]
|
||||
tolerance = 0.1
|
||||
|
||||
assert abs(latency_list[-1] - newest_latency) < tolerance, (
|
||||
f"Newest latency {newest_latency} should be at end, got {latency_list[-1]}"
|
||||
)
|
||||
|
||||
for latency in latency_list:
|
||||
assert abs(latency - oldest_latency) > tolerance, (
|
||||
f"Oldest latency {oldest_latency} should have been discarded, found {latency}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
async def test_latency_list_trimming_discards_oldest_entry_async():
|
||||
"""
|
||||
Async counterpart: the oldest entry is discarded when the latency list is
|
||||
trimmed.
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
latencies_to_add = []
|
||||
for i in range(max_size + 1):
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
expected_latency = float(i + 1)
|
||||
end_time = start_time + expected_latency
|
||||
latencies_to_add.append(expected_latency)
|
||||
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = await test_cache.async_get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size
|
||||
|
||||
newest_latency = latencies_to_add[-1]
|
||||
oldest_latency = latencies_to_add[0]
|
||||
tolerance = 0.1
|
||||
|
||||
assert abs(latency_list[-1] - newest_latency) < tolerance, (
|
||||
f"Newest latency {newest_latency} should be at end of list"
|
||||
)
|
||||
|
||||
for latency in latency_list:
|
||||
assert abs(latency - oldest_latency) > tolerance, f"Oldest latency {oldest_latency} should have been discarded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
async def test_timeout_penalty_discards_oldest_entry():
|
||||
"""
|
||||
Timeout penalties (1000.0) are appended to the latency list and, when the
|
||||
list is full, the oldest entry is discarded.
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
for i in range(max_size):
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
end_time = start_time + float(i + 1)
|
||||
|
||||
await lowest_latency_logger.async_log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
timeout_kwargs = {
|
||||
**kwargs,
|
||||
"exception": litellm.Timeout(message="Request timed out", model="test-model", llm_provider="test"),
|
||||
}
|
||||
|
||||
await lowest_latency_logger.async_log_failure_event(
|
||||
kwargs=timeout_kwargs,
|
||||
response_obj=None,
|
||||
start_time=FROZEN_EPOCH,
|
||||
end_time=FROZEN_EPOCH + 30,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = await test_cache.async_get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size
|
||||
|
||||
assert latency_list[-1] == 1000.0, f"Timeout penalty should be at end of list, got {latency_list[-1]}"
|
||||
|
||||
tolerance = 0.1
|
||||
for latency in latency_list[:-1]:
|
||||
assert abs(latency - 1.0) > tolerance, f"Oldest latency 1.0 should have been discarded, found {latency}"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("frozen_latency_clock")
|
||||
def test_list_order_preserved_after_multiple_trims():
|
||||
"""
|
||||
After many trims, the list still holds the most recent `max_size` entries
|
||||
in insertion order (oldest at index 0, newest at index -1).
|
||||
"""
|
||||
max_size = 3
|
||||
test_cache = DualCache()
|
||||
lowest_latency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
|
||||
)
|
||||
|
||||
model_group = "gpt-3.5-turbo"
|
||||
deployment_id = "test-deployment"
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": model_group,
|
||||
"deployment": "azure/gpt-4.1-mini",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
all_latencies = []
|
||||
for i in range(10):
|
||||
start_time = FROZEN_EPOCH
|
||||
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
|
||||
expected_latency = float(i + 1)
|
||||
end_time = start_time + expected_latency
|
||||
all_latencies.append(expected_latency)
|
||||
|
||||
lowest_latency_logger.log_success_event(
|
||||
response_obj=response_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
latency_key = f"{model_group}_map"
|
||||
cached_data = test_cache.get_cache(key=latency_key)
|
||||
latency_list = cached_data[deployment_id]["latency"]
|
||||
|
||||
assert len(latency_list) == max_size
|
||||
|
||||
expected_remaining = all_latencies[-max_size:]
|
||||
tolerance = 0.1
|
||||
|
||||
for i, expected in enumerate(expected_remaining):
|
||||
assert abs(latency_list[i] - expected) < tolerance, f"At index {i}, expected ~{expected}, got {latency_list[i]}"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import copy
|
||||
import os
|
||||
import functools
|
||||
import uuid
|
||||
from typing import Final, cast
|
||||
|
|
@ -972,3 +973,105 @@ async def test_claude_code_style_session_stays_on_one_deployment_across_turns(lo
|
|||
history = [*history, {"role": "user", "content": [_text(text)]}, {"role": "assistant", "content": "ok"}]
|
||||
|
||||
assert served == [served[0]] * len(user_turns)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anthropic_messages():
|
||||
return [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Here is the full text of a complex legal agreement" * 500,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What are the key terms and conditions in this agreement?",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What are the key terms and conditions in this agreement?",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_with_prompt_caching(anthropic_messages):
|
||||
"""
|
||||
if prompt caching supported model called with prompt caching valid prompt,
|
||||
then 2nd call should go to the same model.
|
||||
"""
|
||||
from litellm.router import Router
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "claude-model",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-4-5-20250929",
|
||||
"api_key": os.environ.get("ANTHROPIC_API_KEY"),
|
||||
"mock_response": "The sky is blue.",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-model",
|
||||
"litellm_params": {
|
||||
"model": "anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"mock_response": "The sky is green.",
|
||||
},
|
||||
},
|
||||
],
|
||||
optional_pre_call_checks=["prompt_caching"],
|
||||
)
|
||||
|
||||
response = await router.acompletion(
|
||||
messages=anthropic_messages,
|
||||
model="claude-model",
|
||||
mock_response="The sky is blue.",
|
||||
)
|
||||
print("response=", response)
|
||||
|
||||
initial_model_id = response._hidden_params["model_id"]
|
||||
|
||||
cache = PromptCachingCache(
|
||||
cache=router.cache,
|
||||
)
|
||||
|
||||
cached_model_id = await _eventually(lambda: cache.get_model_id(messages=anthropic_messages, tools=None))
|
||||
|
||||
assert cached_model_id is not None
|
||||
prompt_caching_cache_key = PromptCachingCache.get_prompt_caching_cache_key(messages=anthropic_messages, tools=None)
|
||||
print(f"prompt_caching_cache_key: {prompt_caching_cache_key}")
|
||||
assert cached_model_id["model_id"] == initial_model_id
|
||||
|
||||
new_messages = anthropic_messages + [{"role": "user", "content": "What is the weather in SF?"}]
|
||||
|
||||
for _ in range(20):
|
||||
response = await router.acompletion(
|
||||
messages=new_messages,
|
||||
model="claude-model",
|
||||
mock_response="The sky is blue.",
|
||||
)
|
||||
print("response=", response)
|
||||
|
||||
assert response._hidden_params["model_id"] == initial_model_id
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
|
@ -38,6 +39,7 @@ from litellm.router_utils.fallback_event_handlers import (
|
|||
run_async_fallback,
|
||||
)
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
||||
from typing import Dict
|
||||
import os
|
||||
|
||||
|
|
@ -2339,3 +2341,117 @@ class CustomTestLogger(CustomLogger):
|
|||
self.failure_fallback_events.append(
|
||||
(original_model_group, kwargs, original_exception)
|
||||
)
|
||||
|
||||
|
||||
def create_test_router_2():
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4",
|
||||
"api_key": "very-fake-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "fake-openai-endpoint-2",
|
||||
"litellm_params": {
|
||||
"model": "openai/fake-openai-endpoint-2",
|
||||
"api_key": "working-key-since-this-is-fake-endpoint",
|
||||
"api_base": FAKE_OPENAI_API_BASE,
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("function_name", ["_acompletion", "_atext_completion"])
|
||||
async def test_multiple_fallbacks(function_name, respx_mock: respx.MockRouter, monkeypatch):
|
||||
"""
|
||||
Tests that if multiple fallbacks passed:
|
||||
- fallback 1 = bad configured deployment / failing endpoint
|
||||
- fallback 2 = working deployment / working endpoint
|
||||
|
||||
Assert that:
|
||||
- a success response is received from the working endpoint (fallback 2)
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
invalid_key = httpx.Response(
|
||||
401,
|
||||
json={
|
||||
"error": {
|
||||
"message": "Incorrect API key provided",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
},
|
||||
)
|
||||
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(return_value=invalid_key)
|
||||
respx_mock.post("https://api.openai.com/v1/completions").mock(return_value=invalid_key)
|
||||
respx_mock.post(f"{FAKE_OPENAI_API_BASE}/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-fake",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": "fake-openai-endpoint-2",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
|
||||
},
|
||||
)
|
||||
)
|
||||
respx_mock.post(f"{FAKE_OPENAI_API_BASE}/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "cmpl-fake",
|
||||
"object": "text_completion",
|
||||
"created": 1700000000,
|
||||
"model": "fake-openai-endpoint-2",
|
||||
"choices": [{"index": 0, "text": "hi", "logprobs": None, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
|
||||
},
|
||||
)
|
||||
)
|
||||
router_2 = create_test_router_2()
|
||||
original_function = getattr(router_2, function_name)
|
||||
|
||||
fallback_model_group = ["gpt-4", "fake-openai-endpoint-2"]
|
||||
original_model_group = "gpt-3.5-turbo"
|
||||
original_exception = Exception("Simulated error")
|
||||
|
||||
request_kwargs: dict[str, object] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}}
|
||||
|
||||
if function_name == "_aembedding":
|
||||
request_kwargs["input"] = "hello this is a test for run_async_fallback"
|
||||
elif function_name == "_atext_completion":
|
||||
request_kwargs["prompt"] = "hello this is a test for run_async_fallback"
|
||||
elif function_name == "_acompletion":
|
||||
request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}]
|
||||
|
||||
result = await run_async_fallback(
|
||||
litellm_router=router_2,
|
||||
original_function=original_function,
|
||||
num_retries=1,
|
||||
fallback_model_group=fallback_model_group,
|
||||
original_model_group=original_model_group,
|
||||
original_exception=original_exception,
|
||||
max_fallbacks=5,
|
||||
fallback_depth=0,
|
||||
**request_kwargs,
|
||||
)
|
||||
|
||||
print(result)
|
||||
|
||||
print(result._hidden_params)
|
||||
|
||||
assert result._hidden_params["api_base"] == FAKE_OPENAI_API_BASE
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from openai.types.chat.chat_completion import ChatCompletion
|
|||
|
||||
import litellm
|
||||
from litellm import acompletion, completion
|
||||
from litellm import acompletion_with_retries, aresponses_with_retries, completion_with_retries, responses_with_retries
|
||||
from litellm import main as litellm_main
|
||||
from litellm.constants import CONTROL_OPTIONS_KEY
|
||||
from litellm.caching.base_cache import BaseCache
|
||||
|
|
@ -6104,3 +6105,84 @@ async def test_async_responses_still_retries_provider_server_errors(monkeypatch)
|
|||
)
|
||||
assert result.status == "completed"
|
||||
assert response.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_with_retries(sync_mode):
|
||||
"""
|
||||
If completion_with_retries is called with num_retries=3, and max_retries=0, then litellm.completion should receive num_retries , max_retries=0
|
||||
"""
|
||||
if sync_mode:
|
||||
target_function = "completion"
|
||||
else:
|
||||
target_function = "acompletion"
|
||||
|
||||
with patch.object(litellm, target_function) as mock_completion:
|
||||
if sync_mode:
|
||||
completion_with_retries(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"gm": "vibe", "role": "user"}],
|
||||
num_retries=3,
|
||||
original_function=mock_completion,
|
||||
)
|
||||
else:
|
||||
await acompletion_with_retries(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"gm": "vibe", "role": "user"}],
|
||||
num_retries=3,
|
||||
original_function=mock_completion,
|
||||
)
|
||||
mock_completion.assert_called_once()
|
||||
assert mock_completion.call_args.kwargs["num_retries"] == 0
|
||||
assert mock_completion.call_args.kwargs["max_retries"] == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_with_retries(sync_mode):
|
||||
"""
|
||||
Test that responses() and aresponses() properly handle num_retries parameter.
|
||||
If responses_with_retries is called with num_retries=3, and max_retries=0,
|
||||
then litellm.responses should receive num_retries=0, max_retries=0
|
||||
"""
|
||||
if sync_mode:
|
||||
target_function = "responses"
|
||||
retry_function = responses_with_retries
|
||||
else:
|
||||
target_function = "aresponses"
|
||||
retry_function = aresponses_with_retries
|
||||
|
||||
with patch(
|
||||
"litellm.responses.main.responses" if sync_mode else "litellm.responses.main.aresponses"
|
||||
) as mock_responses:
|
||||
if sync_mode:
|
||||
mock_responses.return_value = MagicMock()
|
||||
retry_function(
|
||||
model="gpt-4o",
|
||||
input="Hello, what's the weather?",
|
||||
num_retries=3,
|
||||
original_function=mock_responses,
|
||||
)
|
||||
else:
|
||||
mock_responses.return_value = AsyncMock()
|
||||
await retry_function(
|
||||
model="gpt-4o",
|
||||
input="Hello, what's the weather?",
|
||||
num_retries=3,
|
||||
original_function=mock_responses,
|
||||
)
|
||||
|
||||
mock_responses.assert_called_once()
|
||||
assert mock_responses.call_args.kwargs["num_retries"] == 0
|
||||
assert mock_responses.call_args.kwargs["max_retries"] == 0
|
||||
|
||||
|
||||
def test_azure_embedding_exceptions():
|
||||
with pytest.raises(Exception, match="Mock error") as exc_info:
|
||||
litellm.embedding(
|
||||
model="azure/text-embedding-ada-002",
|
||||
input="hello",
|
||||
mock_response="error",
|
||||
)
|
||||
assert str(exc_info.value) == "Mock error"
|
||||
|
|
|
|||
|
|
@ -89,6 +89,7 @@ from litellm.llms.base_llm.vector_store.transformation import(
|
|||
from litellm.types.utils import CallTypes, CredentialItem
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
||||
from tests.large_text import text
|
||||
import traceback
|
||||
import inspect
|
||||
|
|
@ -25040,3 +25041,83 @@ async def test_acompletion_keeps_include_fallback_errors_off_the_wire_and_return
|
|||
model="primary", messages=[{"role": "user", "content": "hi"}], include_fallback_errors=True
|
||||
)
|
||||
_assert_fallback_errors_reached_the_caller_and_not_the_wire(response, primary, backup)
|
||||
|
||||
|
||||
def _create_custom_routing_router():
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/very-special-endpoint",
|
||||
"api_base": FAKE_OPENAI_API_BASE,
|
||||
"api_key": "fake-key",
|
||||
},
|
||||
"model_info": {"id": "very-special-endpoint"},
|
||||
},
|
||||
{
|
||||
"model_name": "azure-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/fast-endpoint",
|
||||
"api_base": FAKE_OPENAI_API_BASE,
|
||||
"api_key": "fake-key",
|
||||
},
|
||||
"model_info": {"id": "fast-endpoint"},
|
||||
},
|
||||
],
|
||||
set_verbose=True,
|
||||
debug_level="DEBUG",
|
||||
)
|
||||
|
||||
|
||||
class SpecialEndpointRoutingStrategy(CustomRoutingStrategyBase):
|
||||
def __init__(self, router_instance: Router):
|
||||
self._router = router_instance
|
||||
|
||||
async def async_get_available_deployment(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, str]] | None = None,
|
||||
input: str | list | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
request_kwargs: dict | None = None,
|
||||
):
|
||||
print("In CUSTOM async get available deployment")
|
||||
model_list = self._router.model_list
|
||||
print("router model list=", model_list)
|
||||
for model in model_list:
|
||||
if isinstance(model, dict):
|
||||
if model["litellm_params"]["model"] == "openai/very-special-endpoint":
|
||||
return model
|
||||
pass
|
||||
|
||||
def get_available_deployment(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, str]] | None = None,
|
||||
input: str | list | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
request_kwargs: dict | None = None,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
def test_reset_custom_routing_strategy():
|
||||
"""
|
||||
Setting a custom routing strategy installs instance-level overrides for
|
||||
get_available_deployment / async_get_available_deployment. Re-initializing the
|
||||
routing strategy must clear them so the class implementations are used again.
|
||||
"""
|
||||
router = _create_custom_routing_router()
|
||||
router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router))
|
||||
|
||||
assert "get_available_deployment" in router.__dict__
|
||||
assert "async_get_available_deployment" in router.__dict__
|
||||
|
||||
router._reset_custom_routing_strategy()
|
||||
|
||||
assert "get_available_deployment" not in router.__dict__
|
||||
assert "async_get_available_deployment" not in router.__dict__
|
||||
assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment
|
||||
|
||||
router._reset_custom_routing_strategy()
|
||||
|
|
|
|||
|
|
@ -1,15 +1,21 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
import os
|
||||
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger, verbose_router_logger
|
||||
from litellm.llms.custom_httpx.async_client_cleanup import close_litellm_async_clients
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -504,3 +510,94 @@ class MyCustomHandler(CustomLogger):
|
|||
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
print(f"On Failure")
|
||||
|
||||
|
||||
def test_async_fallbacks(caplog, respx_mock: respx.MockRouter, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "set_verbose", False)
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
401,
|
||||
json={
|
||||
"error": {
|
||||
"message": "Incorrect API key provided: bad-key.",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
verbose_router_logger.setLevel(level=logging.INFO)
|
||||
verbose_logger.setLevel(logging.CRITICAL + 1)
|
||||
verbose_proxy_logger.setLevel(logging.CRITICAL + 1)
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "azure/gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-4.1-mini",
|
||||
"api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
"mock_response": "Hello world",
|
||||
},
|
||||
"tpm": 240000,
|
||||
"rpm": 1800,
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "bad-key",
|
||||
},
|
||||
"tpm": 1000000,
|
||||
"rpm": 9000,
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
fallbacks=[{"gpt-3.5-turbo": ["azure/gpt-3.5-turbo"]}],
|
||||
num_retries=1,
|
||||
)
|
||||
|
||||
user_message = "Hello, how are you?"
|
||||
messages = [{"content": user_message, "role": "user"}]
|
||||
|
||||
async def _make_request():
|
||||
try:
|
||||
await router.acompletion(model="gpt-3.5-turbo", messages=messages, max_tokens=1)
|
||||
router.reset()
|
||||
except litellm.Timeout:
|
||||
pass
|
||||
except Exception as e:
|
||||
pytest.fail(f"An exception occurred: {e}")
|
||||
finally:
|
||||
router.reset()
|
||||
await close_litellm_async_clients()
|
||||
|
||||
asyncio.run(_make_request())
|
||||
captured_logs = [rec.message for rec in caplog.records]
|
||||
|
||||
captured_logs = [
|
||||
log
|
||||
for log in captured_logs
|
||||
if "Task exception was never retrieved" not in log
|
||||
and "Task was destroyed but it is pending" not in log
|
||||
and "get_available_deployment" not in log
|
||||
and "Selected deployment for model" not in log
|
||||
and "in the Langfuse queue" not in log
|
||||
and "Unclosed client session" not in log
|
||||
and "Unclosed connector" not in log
|
||||
]
|
||||
|
||||
print("\n Captured caplog records - ", captured_logs)
|
||||
|
||||
expected_logs = [
|
||||
"Falling back to model_group = azure/gpt-3.5-turbo",
|
||||
"litellm.acompletion(model=azure/gpt-4.1-mini)\x1b[32m 200 OK\x1b[0m",
|
||||
"Successful fallback b/w models.",
|
||||
]
|
||||
|
||||
assert captured_logs[-3:] == expected_logs
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
|
|||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, headers
|
||||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
from litellm.proxy.utils import is_valid_api_key
|
||||
from litellm.responses.main import aresponses, responses
|
||||
from litellm.types.caching import CachingSupportedCallTypes
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
|
@ -9331,3 +9332,91 @@ def my_pre_call_rule(input: str):
|
|||
if len(input) > 10:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
async def test_completion_with_retry_policy(sync_mode):
|
||||
from litellm.types.router import RetryPolicy
|
||||
|
||||
retry_number = 1
|
||||
retry_policy = RetryPolicy(
|
||||
BadRequestErrorRetries=10,
|
||||
ContentPolicyViolationErrorRetries=retry_number,
|
||||
AuthenticationErrorRetries=0,
|
||||
)
|
||||
|
||||
target_function = "completion_with_retries"
|
||||
|
||||
with patch.object(litellm, target_function) as mock_completion_with_retries:
|
||||
data = {
|
||||
"model": "azure/gpt-3.5-turbo",
|
||||
"messages": [{"gm": "vibe", "role": "user"}],
|
||||
"retry_policy": retry_policy,
|
||||
"mock_response": "Exception: content_filter_policy",
|
||||
}
|
||||
try:
|
||||
if sync_mode:
|
||||
completion(**data)
|
||||
else:
|
||||
await completion(**data)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
mock_completion_with_retries.assert_called_once()
|
||||
assert mock_completion_with_retries.call_args.kwargs["num_retries"] == retry_number
|
||||
assert retry_policy.ContentPolicyViolationErrorRetries == retry_number
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
async def test_responses_retry_on_auth_error(sync_mode, respx_mock: respx.MockRouter, monkeypatch):
|
||||
"""
|
||||
Test that responses API actually retries when encountering authentication errors.
|
||||
This validates that the @client decorator properly handles responses/aresponses retries.
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx_mock.post("https://api.openai.com/v1/responses").mock(
|
||||
return_value=httpx.Response(
|
||||
401,
|
||||
json={
|
||||
"error": {
|
||||
"message": "Incorrect API key provided",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
},
|
||||
)
|
||||
)
|
||||
num_retries = 2
|
||||
|
||||
if sync_mode:
|
||||
with patch.object(litellm, "responses_with_retries") as mock_retry:
|
||||
mock_retry.return_value = None
|
||||
try:
|
||||
responses(
|
||||
model="gpt-4o",
|
||||
input="Test input",
|
||||
num_retries=num_retries,
|
||||
api_key="sk-invalid-key-12345",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
assert mock_retry.called
|
||||
assert mock_retry.call_args.kwargs.get("num_retries") == num_retries
|
||||
else:
|
||||
with patch.object(litellm, "aresponses_with_retries") as mock_retry:
|
||||
mock_retry.return_value = None
|
||||
try:
|
||||
await aresponses(
|
||||
model="gpt-4o",
|
||||
input="Test input",
|
||||
num_retries=num_retries,
|
||||
api_key="sk-invalid-key-12345",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
assert mock_retry.called
|
||||
assert mock_retry.call_args.kwargs.get("num_retries") == num_retries
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue