Merge pull request #38995 from BerriAI/litellm_openai_wif

feat(openai): support workload identity federation (OIDC token exchange)

basedpyright-code-budget.json

@@ -3,7 +3,7 @@
     "limit": 16171
   },
   "reportArgumentType": {
-    "limit": 2226
+    "limit": 2224
   },
   "reportAssignmentType": {
     "limit": 319

litellm/llms/openai/common_utils.py

@@ -268,6 +268,7 @@ class BaseOpenAILLM:
             "max_retries",
             "organization",
             "api_base",
+            "workload_identity_config",
         )
         openai_client_fields: Final = (
             BaseOpenAILLM.get_openai_client_initialization_param_fields(client_type=client_type)

litellm/llms/openai/openai.py

@@ -51,6 +51,7 @@ from .common_utils import (
     drop_params_from_unprocessable_entity_error,
     is_output_token_limit_error,
 )
+from .workload_identity import resolve_openai_workload_identity_config
 
 openaiOSeriesConfig: Final = OpenAIOSeriesConfig()
 openAIGPT5Config: Final = OpenAIGPT5Config()
@@ -349,6 +350,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         client: OpenAI | AsyncOpenAI | None = None,
         shared_session: Optional["ClientSession"] = None,
     ) -> OpenAI | AsyncOpenAI | None:
+        workload_identity_config: Final = resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base)
         client_initialization_params: Final[dict] = locals()
         if client is None:
             if not isinstance(max_retries, int):
@@ -364,28 +366,49 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             if cached_client:
                 if isinstance(cached_client, OpenAI) or isinstance(cached_client, AsyncOpenAI):
                     return cached_client
-            http_client: Final[httpx.Client | httpx.AsyncClient | None] = (
-                OpenAIChatCompletion._get_async_http_client(shared_session=shared_session)
-                if is_async
-                else OpenAIChatCompletion._get_sync_http_client()
-            )
             if is_async:
-                _new_client: OpenAI | AsyncOpenAI = AsyncOpenAI(
-                    api_key=api_key,
-                    base_url=api_base,
-                    http_client=http_client,
-                    timeout=timeout,
-                    max_retries=max_retries,
-                    organization=organization,
+                async_http_client: Final = OpenAIChatCompletion._get_async_http_client(shared_session=shared_session)
+                http_client: httpx.Client | httpx.AsyncClient | None = async_http_client
+                _new_client: OpenAI | AsyncOpenAI = (
+                    AsyncOpenAI(
+                        workload_identity=workload_identity_config.to_sdk_workload_identity(),
+                        base_url=api_base,
+                        http_client=async_http_client,
+                        timeout=timeout,
+                        max_retries=max_retries,
+                        organization=organization,
+                    )
+                    if workload_identity_config is not None
+                    else AsyncOpenAI(
+                        api_key=api_key,
+                        base_url=api_base,
+                        http_client=async_http_client,
+                        timeout=timeout,
+                        max_retries=max_retries,
+                        organization=organization,
+                    )
                 )
             else:
-                _new_client = OpenAI(
-                    api_key=api_key,
-                    base_url=api_base,
-                    http_client=http_client,
-                    timeout=timeout,
-                    max_retries=max_retries,
-                    organization=organization,
+                sync_http_client: Final = OpenAIChatCompletion._get_sync_http_client()
+                http_client = sync_http_client
+                _new_client = (
+                    OpenAI(
+                        workload_identity=workload_identity_config.to_sdk_workload_identity(),
+                        base_url=api_base,
+                        http_client=sync_http_client,
+                        timeout=timeout,
+                        max_retries=max_retries,
+                        organization=organization,
+                    )
+                    if workload_identity_config is not None
+                    else OpenAI(
+                        api_key=api_key,
+                        base_url=api_base,
+                        http_client=sync_http_client,
+                        timeout=timeout,
+                        max_retries=max_retries,
+                        organization=organization,
+                    )
                 )
 
             ## SAVE CACHE KEY

litellm/llms/openai/responses/transformation.py

@@ -21,6 +21,7 @@ from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import LlmProviders
 
 from ..common_utils import OpenAIError
+from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config
 
 OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: Final = 16
 
@@ -392,6 +393,14 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
         litellm_params = litellm_params or GenericLiteLLMParams()
         api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
         headers.setdefault("Content-Type", "application/json")
+        workload_identity_config: Final = (
+            resolve_openai_workload_identity_config(api_key=api_key, api_base=litellm_params.api_base)
+            if self.custom_llm_provider is LlmProviders.OPENAI
+            else None
+        )
+        if workload_identity_config is not None:
+            headers["Authorization"] = f"Bearer {get_workload_identity_bearer_token(workload_identity_config)}"
+            return headers
         headers["Authorization"] = f"Bearer {api_key}"
         return headers
 

litellm/llms/openai/workload_identity.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING, Final
+from urllib.parse import urlparse
+
+import litellm
+from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
+
+from .common_utils import OpenAIError
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth
+
+OPENAI_WIF_CLIENT_ID: Final = "litellm"
+_OPENAI_API_HOST: Final = "api.openai.com"
+_SDK_UPGRADE_MESSAGE: Final = (
+    "OpenAI workload identity federation requires openai>=2.32.0. "
+    "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / "
+    "OPENAI_SERVICE_ACCOUNT_ID / OPENAI_IDENTITY_TOKEN_FILE."
+)
+
+
+@dataclass(frozen=True, slots=True)
+class OpenAIWorkloadIdentityConfig:
+    identity_provider_id: str
+    service_account_id: str
+    token_file: str
+
+    def to_sdk_workload_identity(self) -> WorkloadIdentity:
+        k8s_token_provider: Final = _load_sdk_k8s_token_provider()
+        workload_identity: Final[WorkloadIdentity] = {
+            "client_id": OPENAI_WIF_CLIENT_ID,
+            "identity_provider_id": self.identity_provider_id,
+            "service_account_id": self.service_account_id,
+            "provider": k8s_token_provider(self.token_file),
+        }
+        return workload_identity
+
+
+def resolve_openai_workload_identity_config(
+    api_key: str | None,
+    api_base: str | None,
+) -> OpenAIWorkloadIdentityConfig | None:
+    static_api_key: Final = normalize_nonempty_secret_str(api_key) or normalize_nonempty_secret_str(
+        get_secret_str("OPENAI_API_KEY")
+    )
+    if static_api_key is not None:
+        return None
+    effective_api_base: Final = (
+        api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE")
+    )
+    if not _targets_openai_api(effective_api_base):
+        return None
+    identity_provider_id: Final = get_secret_str("OPENAI_IDENTITY_PROVIDER_ID")
+    service_account_id: Final = get_secret_str("OPENAI_SERVICE_ACCOUNT_ID")
+    token_file: Final = get_secret_str("OPENAI_IDENTITY_TOKEN_FILE")
+    if not identity_provider_id or not service_account_id or not token_file:
+        return None
+    return OpenAIWorkloadIdentityConfig(
+        identity_provider_id=identity_provider_id,
+        service_account_id=service_account_id,
+        token_file=token_file,
+    )
+
+
+def get_workload_identity_bearer_token(config: OpenAIWorkloadIdentityConfig) -> str:
+    return _workload_identity_auth(config).get_token()
+
+
+def _targets_openai_api(api_base: str | None) -> bool:
+    if api_base is None:
+        return True
+    parsed: Final = urlparse(api_base)
+    return parsed.scheme == "https" and parsed.hostname == _OPENAI_API_HOST
+
+
+@lru_cache(maxsize=16)
+def _workload_identity_auth(config: OpenAIWorkloadIdentityConfig) -> WorkloadIdentityAuth:
+    sdk_workload_identity_auth: Final = _load_sdk_workload_identity_auth()
+    return sdk_workload_identity_auth(workload_identity=config.to_sdk_workload_identity())
+
+
+def _load_sdk_workload_identity_auth() -> type[WorkloadIdentityAuth]:
+    try:
+        from openai.auth import WorkloadIdentityAuth as sdk_workload_identity_auth
+    except ImportError as e:
+        raise OpenAIError(status_code=500, message=_SDK_UPGRADE_MESSAGE) from e
+    return sdk_workload_identity_auth
+
+
+def _load_sdk_k8s_token_provider() -> Callable[[str], SubjectTokenProvider]:
+    try:
+        from openai.auth import k8s_service_account_token_provider
+    except ImportError as e:
+        raise OpenAIError(status_code=500, message=_SDK_UPGRADE_MESSAGE) from e
+    return k8s_service_account_token_provider

tests/test_litellm/llms/openai/test_openai_workload_identity.py

@@ -0,0 +1,238 @@
+import json
+import sys
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import respx
+from openai import AsyncOpenAI, OpenAI
+
+import litellm
+from litellm.llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig
+from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError
+from litellm.llms.openai.openai import OpenAIChatCompletion
+from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+from litellm.llms.openai.workload_identity import (
+    OpenAIWorkloadIdentityConfig,
+    _workload_identity_auth,
+    get_workload_identity_bearer_token,
+    resolve_openai_workload_identity_config,
+)
+from litellm.types.router import GenericLiteLLMParams
+
+TOKEN_EXCHANGE_URL: Final = "https://auth.openai.com/oauth/token"
+
+
+@pytest.fixture
+def wif_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> OpenAIWorkloadIdentityConfig:
+    token_file: Final = tmp_path / "subject_token.jwt"
+    token_file.write_text("subject-token-from-file")
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
+    monkeypatch.delenv("OPENAI_API_BASE", raising=False)
+    monkeypatch.setattr(litellm, "api_base", None)
+    monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_test123")
+    monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "user-test456")
+    monkeypatch.setenv("OPENAI_IDENTITY_TOKEN_FILE", str(token_file))
+    _workload_identity_auth.cache_clear()
+    litellm.in_memory_llm_clients_cache.flush_cache()
+    return OpenAIWorkloadIdentityConfig(
+        identity_provider_id="idp_test123",
+        service_account_id="user-test456",
+        token_file=str(token_file),
+    )
+
+
+def mock_token_exchange(access_token: str = "exchanged-bearer-token") -> respx.Route:
+    return respx.post(TOKEN_EXCHANGE_URL).mock(
+        return_value=httpx.Response(200, json={"access_token": access_token, "expires_in": 3600})
+    )
+
+
+class TestResolveConfig:
+    def test_resolves_from_env(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env
+
+    def test_static_api_key_wins(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        assert resolve_openai_workload_identity_config(api_key="sk-static", api_base=None) is None
+
+    def test_env_openai_api_key_wins(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env")
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None
+
+    @pytest.mark.parametrize("empty_key", ["", "   "])
+    def test_empty_api_key_arg_does_not_disable_wif(
+        self, wif_env: OpenAIWorkloadIdentityConfig, empty_key: str
+    ) -> None:
+        assert resolve_openai_workload_identity_config(api_key=empty_key, api_base=None) == wif_env
+
+    @pytest.mark.parametrize("empty_key", ["", "   "])
+    def test_empty_env_openai_api_key_does_not_disable_wif(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch, empty_key: str
+    ) -> None:
+        monkeypatch.setenv("OPENAI_API_KEY", empty_key)
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env
+
+    def test_foreign_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        assert resolve_openai_workload_identity_config(api_key=None, api_base="https://my-vllm.internal/v1") is None
+
+    def test_openai_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        assert resolve_openai_workload_identity_config(api_key=None, api_base="https://api.openai.com/v1") == wif_env
+
+    def test_plaintext_http_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        assert resolve_openai_workload_identity_config(api_key=None, api_base="http://api.openai.com/v1") is None
+
+    def test_foreign_env_base_url_disables(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("OPENAI_BASE_URL", "https://my-vllm.internal/v1")
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None
+
+    def test_openai_env_base_url_allows(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env
+
+    def test_foreign_litellm_api_base_disables(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setattr(litellm, "api_base", "https://my-vllm.internal/v1")
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None
+
+    @pytest.mark.parametrize(
+        "missing_var",
+        ["OPENAI_IDENTITY_PROVIDER_ID", "OPENAI_SERVICE_ACCOUNT_ID", "OPENAI_IDENTITY_TOKEN_FILE"],
+    )
+    def test_partial_env_disables(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch, missing_var: str
+    ) -> None:
+        monkeypatch.delenv(missing_var)
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None
+
+
+class TestTokenExchange:
+    @respx.mock
+    def test_exchanges_subject_token_for_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        route: Final = mock_token_exchange()
+        assert get_workload_identity_bearer_token(wif_env) == "exchanged-bearer-token"
+        request_body: Final = json.loads(route.calls.last.request.content)
+        assert request_body["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange"
+        assert request_body["subject_token"] == "subject-token-from-file"
+        assert request_body["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt"
+        assert request_body["identity_provider_id"] == "idp_test123"
+        assert request_body["service_account_id"] == "user-test456"
+
+    @respx.mock
+    def test_token_cached_across_mints(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        route: Final = mock_token_exchange()
+        first: Final = get_workload_identity_bearer_token(wif_env)
+        second: Final = get_workload_identity_bearer_token(wif_env)
+        assert first == second == "exchanged-bearer-token"
+        assert route.call_count == 1
+
+    def test_old_sdk_raises_upgrade_error(
+        self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        import openai as openai_module
+
+        monkeypatch.delattr(openai_module, "auth", raising=False)
+        monkeypatch.setitem(sys.modules, "openai.auth", None)
+        with pytest.raises(OpenAIError, match=r"openai>=2\.32\.0"):
+            wif_env.to_sdk_workload_identity()
+
+
+class TestClientConstruction:
+    def test_sync_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key=None, api_base=None)
+        assert isinstance(client, OpenAI)
+        assert client.api_key == "workload-identity-auth"
+        assert client._workload_identity_auth is not None
+
+    def test_async_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        client: Final = OpenAIChatCompletion()._get_openai_client(is_async=True, api_key=None, api_base=None)
+        assert isinstance(client, AsyncOpenAI)
+        assert client.api_key == "workload-identity-auth"
+        assert client._workload_identity_auth is not None
+
+    def test_static_key_client_unaffected(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key="sk-static", api_base=None)
+        assert isinstance(client, OpenAI)
+        assert client.api_key == "sk-static"
+        assert client._workload_identity_auth is None
+
+    def test_cache_key_separates_wif_identities(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        other_config: Final = OpenAIWorkloadIdentityConfig(
+            identity_provider_id="idp_other",
+            service_account_id="user-other",
+            token_file=wif_env.token_file,
+        )
+        keys: Final = tuple(
+            BaseOpenAILLM.get_openai_client_cache_key(
+                client_initialization_params={"api_key": None, "is_async": False, "workload_identity_config": config},
+                client_type="openai",
+            )
+            for config in (wif_env, other_config, None)
+        )
+        assert len(set(keys)) == 3
+
+    @respx.mock
+    def test_request_carries_exchanged_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        mock_token_exchange()
+        completion_route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock(
+            return_value=httpx.Response(
+                200,
+                json={
+                    "id": "chatcmpl-wif",
+                    "object": "chat.completion",
+                    "created": 1,
+                    "model": "gpt-4o-mini",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "message": {"role": "assistant", "content": "ok"},
+                            "finish_reason": "stop",
+                        }
+                    ],
+                    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+                },
+            )
+        )
+        client = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key=None, api_base=None)
+        assert isinstance(client, OpenAI)
+        client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}])
+        auth_header: Final = completion_route.calls.last.request.headers["Authorization"]
+        assert auth_header == "Bearer exchanged-bearer-token"
+
+
+class TestResponsesValidateEnvironment:
+    @respx.mock
+    def test_mints_bearer_when_wif_configured(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        mock_token_exchange()
+        headers: Final = OpenAIResponsesAPIConfig().validate_environment(
+            headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams()
+        )
+        assert headers["Authorization"] == "Bearer exchanged-bearer-token"
+
+    def test_static_key_wins(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        headers: Final = OpenAIResponsesAPIConfig().validate_environment(
+            headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams(api_key="sk-responses")
+        )
+        assert headers["Authorization"] == "Bearer sk-responses"
+
+    def test_foreign_api_base_skips_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        headers: Final = OpenAIResponsesAPIConfig().validate_environment(
+            headers={},
+            model="gpt-4o-mini",
+            litellm_params=GenericLiteLLMParams(api_base="https://my-vllm.internal/v1"),
+        )
+        assert headers["Authorization"] == "Bearer None"
+
+    def test_litellm_proxy_subclass_never_mints_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        headers: Final = LiteLLMProxyResponsesAPIConfig().validate_environment(
+            headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams()
+        )
+        assert headers["Authorization"] == "Bearer None"