diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 0f618aeadbd..4d6599fc1b5 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -14,6 +14,7 @@ from pydantic import BaseModel import litellm from litellm.cost_calculator import ( + completion_cost, handle_realtime_stream_cost_calculation, response_cost_calculator, ) @@ -22,6 +23,33 @@ from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.utils import TranscriptionResponse +def test_completion_cost_uses_response_model_for_dynamic_routing(): + """ + Test that completion_cost uses the model from the response object + when the input model (e.g., azure-model-router) is not in model_cost. + This supports Azure Model Router and similar dynamic routing scenarios. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + # Simulate Azure Model Router: input is generic router, response has actual model + response = ModelResponse( + id="test-id", + model="azure_ai/gpt-4o-2024-08-06", # Response contains actual model used + choices=[], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + + # Should calculate cost using the response model, not the input model + cost = completion_cost( + completion_response=response, + model="azure_ai/azure-model-router", # Input model doesn't exist in model_cost + custom_llm_provider="azure_ai", + ) + + assert cost > 0, "Cost should be calculated using response model" + + def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): _hidden_params = { diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index da52b7bcb39..6c9c1f31c09 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -21,6 +21,7 @@ from litellm.types.utils import ( from litellm.utils import ( ProviderConfigManager, TextCompletionStreamWrapper, + _check_provider_match, get_llm_provider, get_optional_params_image_gen, is_cached_message, @@ -29,6 +30,30 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_check_provider_match_azure_ai_allows_openai_and_azure(): + """ + Test that azure_ai provider can match openai and azure models. + This is needed for Azure Model Router which can route to OpenAI models. + """ + # azure_ai should match openai models + assert _check_provider_match( + model_info={"litellm_provider": "openai"}, + custom_llm_provider="azure_ai" + ) is True + + # azure_ai should match azure models + assert _check_provider_match( + model_info={"litellm_provider": "azure"}, + custom_llm_provider="azure_ai" + ) is True + + # azure_ai should NOT match other providers + assert _check_provider_match( + model_info={"litellm_provider": "anthropic"}, + custom_llm_provider="azure_ai" + ) is False + + def test_get_optional_params_image_gen(): from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig