diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d60c3e9c0af..229d1eca3e8 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 16171 }, "reportArgumentType": { - "limit": 2226 + "limit": 2224 }, "reportAssignmentType": { "limit": 319 diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 1b1ab80e85d..4d774f6f165 100644 --- a/litellm/llms/openai/common_utils.py +++ b/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) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6e66c998acf..16fa0017b23 100644 --- a/litellm/llms/openai/openai.py +++ b/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 diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index eac844a790d..eadc087383a 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/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 diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py new file mode 100644 index 00000000000..ecec161ed46 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py new file mode 100644 index 00000000000..d8d9936e9a1 --- /dev/null +++ b/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"