diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 8ca45ba3f75..837832795b1 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -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 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 6dee2bb605f..4c3c3aa1cfa 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -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" diff --git a/tests/local_testing/test_anthropic_prompt_caching.py b/tests/local_testing/test_anthropic_prompt_caching.py index 9fb16808f32..24ffe1b093e 100644 --- a/tests/local_testing/test_anthropic_prompt_caching.py +++ b/tests/local_testing/test_anthropic_prompt_caching.py @@ -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 diff --git a/tests/local_testing/test_completion_with_retries.py b/tests/local_testing/test_completion_with_retries.py index ede07a15225..1f530682c7b 100644 --- a/tests/local_testing/test_completion_with_retries.py +++ b/tests/local_testing/test_completion_with_retries.py @@ -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 diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index feb9bc2d45b..d42dfded27e 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -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() diff --git a/tests/local_testing/test_fake_openai_endpoint.py b/tests/local_testing/test_fake_openai_endpoint.py index d5236d3de1b..5487831b31e 100644 --- a/tests/local_testing/test_fake_openai_endpoint.py +++ b/tests/local_testing/test_fake_openai_endpoint.py @@ -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", diff --git a/tests/local_testing/test_lowest_latency_routing.py b/tests/local_testing/test_lowest_latency_routing.py index aba4500199a..51e3058c90c 100644 --- a/tests/local_testing/test_lowest_latency_routing.py +++ b/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]}" diff --git a/tests/local_testing/test_router_custom_routing.py b/tests/local_testing/test_router_custom_routing.py index bd624f7a19f..39043ad4ccc 100644 --- a/tests/local_testing/test_router_custom_routing.py +++ b/tests/local_testing/test_router_custom_routing.py @@ -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 diff --git a/tests/local_testing/test_router_debug_logs.py b/tests/local_testing/test_router_debug_logs.py deleted file mode 100644 index b3c26a7689f..00000000000 --- a/tests/local_testing/test_router_debug_logs.py +++ /dev/null @@ -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 diff --git a/tests/local_testing/test_router_fallback_handlers.py b/tests/local_testing/test_router_fallback_handlers.py deleted file mode 100644 index d427285f750..00000000000 --- a/tests/local_testing/test_router_fallback_handlers.py +++ /dev/null @@ -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 diff --git a/tests/unit/batch_completion/__init__.py b/tests/unit/batch_completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/local_testing/test_batch_completion_return_exceptions.py b/tests/unit/batch_completion/test_main.py similarity index 58% rename from tests/local_testing/test_batch_completion_return_exceptions.py rename to tests/unit/batch_completion/test_main.py index e04879382f9..4ff7adebe1e 100644 --- a/tests/local_testing/test_batch_completion_return_exceptions.py +++ b/tests/unit/batch_completion/test_main.py @@ -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], ( diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 56e97173eb3..b698f69d7c1 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -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'}" + ) diff --git a/tests/unit/router_strategy/test_lowest_latency.py b/tests/unit/router_strategy/test_lowest_latency.py index 1a8614e3fca..4e41d6f3ade 100644 --- a/tests/unit/router_strategy/test_lowest_latency.py +++ b/tests/unit/router_strategy/test_lowest_latency.py @@ -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]}" diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 18b23a9a703..c6ab259991e 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -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 diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 38c13decaa7..c23a8c95906 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -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 diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index c1b72bf8ff1..ab64fabce7e 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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" diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4356488a2f8..d603a8cabe2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -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() diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index 751352ecfb0..dc7ba5b4894 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -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 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 58fd2250d2a..63e2280dc86 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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