scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(openai): keep SDK 3 from sending OPENAI_ADMIN_KEY on model calls and send Azure Bearer on region and env-key calls

Browse source
This commit is contained in:
mateo-berri 2026-10-29 12:49:10 -07:00
parent a68052edc7
commit 3a3987bfd7
9 changed files with 164 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/llms/azure/azure.py
View file

@ -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

19
litellm/llms/azure/common_utils.py
View file

@ -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)

20
litellm/llms/openai/common_utils.py
View file

@ -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)

5
litellm/llms/openai/openai.py
View file

@ -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,

5
litellm/llms/openai/workload_identity.py
View file

@ -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,

16
litellm/main.py
View file

@ -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

25
tests/unit/llms/azure/test_azure_common_utils.py
View file

@ -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"

33
tests/unit/llms/openai/test_openai_common_utils.py
View file

@ -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

36
tests/unit/test_main.py
View file

@ -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")