From 3a3987bfd786c3cc525414b422da404aa6fa9fc2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:49:10 -0700 Subject: [PATCH] fix(openai): keep SDK 3 from sending OPENAI_ADMIN_KEY on model calls and send Azure Bearer on region and env-key calls --- litellm/llms/azure/azure.py | 29 +++++++++------ litellm/llms/azure/common_utils.py | 19 ++++++---- litellm/llms/openai/common_utils.py | 20 +++++++++++ litellm/llms/openai/openai.py | 5 +-- litellm/llms/openai/workload_identity.py | 5 +-- litellm/main.py | 16 +++++++-- .../llms/azure/test_azure_common_utils.py | 25 ++++++++++++- .../llms/openai/test_openai_common_utils.py | 33 ++++++++++++++++- tests/unit/test_main.py | 36 +++++++++++++++++++ 9 files changed, 164 insertions(+), 24 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 8bf44159a47..680a84d6c3d 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -51,6 +51,7 @@ from .common_utils import ( process_azure_headers, redact_azure_auth_headers, select_azure_base_url_or_endpoint, + with_one_azure_sdk_credential, ) from .image_generation import ( AzureFoundryMAIImageGenerationConfig, @@ -1330,21 +1331,29 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_base += "/" api_base += f"{model}" client = AzureOpenAI( - base_url=api_base, - api_version=api_version, - api_key=api_key, - timeout=timeout, - http_client=client_session, + **with_one_azure_sdk_credential( + { + "base_url": api_base, + "api_version": api_version, + "api_key": api_key, + "timeout": timeout, + "http_client": client_session, + } + ) ) model = None # cloudflare ai gateway, needs model=None else: client = AzureOpenAI( - api_version=api_version, - azure_endpoint=api_base, - api_key=api_key, - timeout=timeout, - http_client=client_session, + **with_one_azure_sdk_credential( + { + "api_version": api_version, + "azure_endpoint": api_base, + "api_key": api_key, + "timeout": timeout, + "http_client": client_session, + } + ) ) # only run this check if it's not cloudflare ai gateway diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index e746a8e0d0c..df088d03ae2 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -22,7 +22,12 @@ from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient +from litellm.llms.openai.common_utils import ( + BaseOpenAILLM, + OpenAIAsyncHTTPClient, + OpenAIHTTPClient, + openai_sdk_key_params, +) from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, single_flight_bearer_token_provider, @@ -693,10 +698,12 @@ class BaseAzureLLM(BaseOpenAILLM): else v1_token_provider ) - v1_params: Final[_AzureV1ClientParams] = { - "api_key": v1_api_key, - "base_url": f"{ensure_api_route(api_base or '', 'openai/v1')}/", - } + v1_base_url: Final = f"{ensure_api_route(api_base or '', 'openai/v1')}/" + v1_params: Final[_AzureV1ClientParams] = ( + {"api_key": v1_api_key, "base_url": v1_base_url} + if callable(v1_api_key) + else {**openai_sdk_key_params(v1_api_key), "base_url": v1_base_url} + ) if "timeout" in azure_client_params: v1_params["timeout"] = azure_client_params["timeout"] if "max_retries" in azure_client_params: @@ -704,7 +711,7 @@ class BaseAzureLLM(BaseOpenAILLM): if "http_client" in azure_client_params: v1_params["http_client"] = azure_client_params["http_client"] - verbose_logger.debug("Using Azure v1 API with base_url: %s", v1_params["base_url"]) + verbose_logger.debug("Using Azure v1 API with base_url: %s", v1_base_url) if _is_async is True: openai_client = AsyncOpenAI(**v1_params) diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index a9d2dcdffb8..507368bba9e 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -24,6 +24,7 @@ from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta from openai.types.completion_usage import CompletionUsage from pydantic import BaseModel, TypeAdapter +from typing_extensions import ReadOnly, TypedDict if TYPE_CHECKING: from aiohttp import ClientSession @@ -148,6 +149,25 @@ class OpenAIError(BaseLLMException): ) +_SDK_READS_ADMIN_KEY_ENV: Final = "admin_api_key" in _OPENAI_INIT_PARAMS +_MISSING_OPENAI_API_KEY_MESSAGE: Final = ( + "Missing credentials. Please pass an `api_key` or set the `OPENAI_API_KEY` environment variable." +) + + +class OpenAISdkKeyParams(TypedDict, total=False): + api_key: ReadOnly[str | None] + admin_api_key: ReadOnly[str] + + +def openai_sdk_key_params(api_key: str | None) -> OpenAISdkKeyParams: + if not _SDK_READS_ADMIN_KEY_ENV: + return {"api_key": api_key} + if not api_key: + raise OpenAIError(status_code=401, message=_MISSING_OPENAI_API_KEY_MESSAGE) + return {"api_key": api_key, "admin_api_key": api_key} + + _DECODED_JSON: Final = TypeAdapter(object) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index af2e85c5f20..5e9d4f2af7b 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -64,6 +64,7 @@ from .common_utils import ( is_openai_backed_api_base, is_output_token_limit_error, model_dump_provider_response, + openai_sdk_key_params, ) from .workload_identity import ( build_async_openai_client, @@ -454,7 +455,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if workload_identity_config is not None else AsyncOpenAI( - api_key=openai_static_api_key(api_key), + **openai_sdk_key_params(openai_static_api_key(api_key)), base_url=api_base, http_client=async_http_client, timeout=timeout, @@ -476,7 +477,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if workload_identity_config is not None else OpenAI( - api_key=openai_static_api_key(api_key), + **openai_sdk_key_params(openai_static_api_key(api_key)), base_url=api_base, http_client=sync_http_client, timeout=timeout, diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index 5952ae3b3ef..30c3c83f12d 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -21,6 +21,7 @@ from .common_utils import ( OpenAIError, OpenAIHTTPClient, is_openai_backed_api_base, + openai_sdk_key_params, ) if TYPE_CHECKING: @@ -185,7 +186,7 @@ def build_openai_client( None if static_key_http_client_factory is None else static_key_http_client_factory() ) return OpenAI( - api_key=openai_static_api_key(api_key), + **openai_sdk_key_params(openai_static_api_key(api_key)), base_url=api_base, timeout=timeout, max_retries=retries, @@ -241,7 +242,7 @@ def build_async_openai_client( None if static_key_http_client_factory is None else static_key_http_client_factory() ) return AsyncOpenAI( - api_key=openai_static_api_key(api_key), + **openai_sdk_key_params(openai_static_api_key(api_key)), base_url=api_base, timeout=timeout, max_retries=retries, diff --git a/litellm/main.py b/litellm/main.py index d1af2de3a13..d45d088c132 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6851,7 +6851,13 @@ def embedding( azure_ad_token: Final = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") - api_key = api_key or litellm.api_key or litellm.azure_key or get_secret_str("AZURE_API_KEY") + api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_API_KEY") + or get_secret_str("AZURE_OPENAI_API_KEY") + ) if api_base is None: raise ValueError("No API Base provided for Azure OpenAI LLM provider. Set 'AZURE_API_BASE' in .env") @@ -8441,7 +8447,13 @@ def transcription( azure_ad_token: Final = kwargs.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") - api_key = api_key or litellm.api_key or litellm.azure_key or get_secret_str("AZURE_API_KEY") + api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_API_KEY") + or get_secret_str("AZURE_OPENAI_API_KEY") + ) optional_params["extra_headers"] = extra_headers diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 7953f2b27e8..ecd2823b503 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -27,11 +27,12 @@ from litellm.llms.azure.common_utils import ( from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, ) -from litellm.types.router import GenericLiteLLMParams +from litellm.types.router import GenericLiteLLMParams, LiteLLM_Params from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, ) from litellm.types.utils import CallTypes +from litellm.utils import get_model_region # Mock the necessary dependencies @@ -2629,3 +2630,25 @@ async def test_azure_sdk_client_sends_the_auth_headers_sdk_2_sent( ) assert json.loads(response.choices[0].message.content or "{}") == expected_headers + + +def _region_only_for_key_and_bearer(request: httpx.Request) -> httpx.Response: + if (request.headers.get("api-key"), request.headers.get("authorization")) != ("azure-key", "Bearer azure-key"): + return httpx.Response(401, json={"error": {"message": "denied"}}) + return httpx.Response( + 200, + headers={"x-ms-region": "eastus"}, + json={"id": "chatcmpl-region", "object": "chat.completion", "created": 0, "model": "gpt", "choices": []}, + ) + + +@pytest.mark.parametrize( + "api_base", + ["https://region-lookup.openai.azure.com", "https://gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/res"], +) +def test_azure_region_lookup_sends_the_key_and_bearer_headers(monkeypatch: pytest.MonkeyPatch, api_base: str) -> None: + session: Final = httpx.Client(transport=httpx.MockTransport(_region_only_for_key_and_bearer)) + monkeypatch.setattr(litellm, "client_session", session) + params: Final = LiteLLM_Params(model="azure/gpt-region", api_key="azure-key", api_base=api_base) + + assert get_model_region(litellm_params=params, mode="chat") == "eastus" diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 206fd13260c..820c00c015d 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -78,10 +78,11 @@ API_FUNCTION_PARAMS = [ @pytest.mark.parametrize("function_name,is_async,args", API_FUNCTION_PARAMS) @pytest.mark.asyncio -async def test_openai_client_reuse(function_name, is_async, args): +async def test_openai_client_reuse(function_name, is_async, args, monkeypatch: pytest.MonkeyPatch): """ Test that multiple API calls reuse the same OpenAI client """ + monkeypatch.setenv("OPENAI_API_KEY", "sk-client-reuse") # Determine which client class to mock based on whether the test is async client_path = "litellm.llms.openai.openai.AsyncOpenAI" if is_async else "litellm.llms.openai.openai.OpenAI" @@ -565,3 +566,33 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session( ) assert response.data[0].url == "https://example.com/image.png" assert not session.is_closed + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_unkeyed_openai_deployment_never_sends_the_openai_admin_key( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, is_async: bool +) -> None: + monkeypatch.setenv("OPENAI_ADMIN_KEY", "sk-admin-key-for-org-apis") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + api_base: Final = f"https://unkeyed-{is_async}.example.com/v1" + route: Final = respx_mock.post(f"{api_base}/chat/completions").respond( + 200, + json={ + "id": "chatcmpl-unkeyed", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + }, + ) + request: Final = { + "model": "openai/gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "api_base": api_base, + } + + with pytest.raises(litellm.AuthenticationError, match="Missing credentials"): + await litellm.acompletion(**request) if is_async else litellm.completion(**request) + + assert not route.called diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index e119fe133d5..a7a4cf09793 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10410,3 +10410,39 @@ async def test_bedrock_strict_tools_limit_rejection_is_not_sent_again_when_no_to assert error is not None and "Too many strict tools" in str(error) assert len(sent) == 1 + + +_AZURE_ENV_KEY_RESPONSES: Final = { + "embedding": { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, + "transcription": {"text": "Four score and seven years ago"}, +} + + +@pytest.mark.parametrize("call", ["embedding", "transcription"]) +def test_azure_key_from_azure_openai_api_key_env_sends_the_bearer_header( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, call: str +) -> None: + monkeypatch.delenv("AZURE_API_KEY", raising=False) + monkeypatch.setenv("AZURE_OPENAI_API_KEY", "azure-env-key") + api_base: Final = f"https://env-key-{call}.openai.azure.com" + route: Final = respx_mock.post(url__regex=rf"{api_base}/openai/deployments/.*").respond( + 200, json=_AZURE_ENV_KEY_RESPONSES[call] + ) + target: Final = {"api_base": api_base, "api_version": "2024-06-01", "max_retries": 0} + + if call == "embedding": + litellm.embedding(model="azure/text-embedding-3-small", input=["hello"], **target) + else: + litellm.transcription( + model="azure/whisper-1", + file=("gettysburg.wav", (Path(__file__).parents[1] / "gettysburg.wav").read_bytes(), "audio/wav"), + **target, + ) + + sent: Final = route.calls.last.request.headers + assert (sent.get("api-key"), sent.get("authorization")) == ("azure-env-key", "Bearer azure-env-key")