diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index bebf5e2ee5a..0ab386d544e 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -28,7 +28,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelResponse, ProviderField -from litellm.utils import _add_path_to_api_base, supports_tool_choice +from litellm.utils import _add_path_to_api_base, supports_reasoning, supports_tool_choice if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer @@ -111,7 +111,9 @@ class AzureAIStudioConfig(OpenAIConfig): drop_params=drop_params, ) - if "max_completion_tokens" not in optional_params: + if "max_completion_tokens" not in optional_params or supports_reasoning( + model=model, custom_llm_provider="azure_ai" + ): return optional_params max_tokens: Final = optional_params["max_completion_tokens"] diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 5a99a4ef95a..28568958569 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -169,6 +169,17 @@ def test_azure_ai_keeps_max_completion_tokens_for_gpt_5(): assert "max_tokens" not in mapped_params +def test_azure_ai_keeps_max_completion_tokens_for_reasoning_models(_local_model_cost_map: None) -> None: + mapped_params: Final = litellm.get_optional_params( + model="o3", + custom_llm_provider="azure_ai", + max_completion_tokens=256, + ) + + assert mapped_params["max_completion_tokens"] == 256 + assert "max_tokens" not in mapped_params + + def test_azure_ai_keeps_params_without_max_completion_tokens(): mapped_params: Final = litellm.get_optional_params( model="mistral-large-3",