feat(client): add support for short-lived tokens (#1608)

Co-authored-by: stainless-app[bot] <142633134+stainless-app[bot]@users.noreply.github.com>

Author: karpetrosyan

.stats.yml

@@ -1,4 +1,4 @@
 configured_endpoints: 152
-openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/openai%2Fopenai-dd99495ad509338e6de862802839360dfe394d5cd6d6ba6d13fec8fca92328b8.yml
+openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/openai%2Fopenai-a6eca1bd01e0c434af356fe5275c206057216a4e626d1051d294c27016cd6d05.yml
 openapi_spec_hash: 68abda9122013a9ae3f084cfdbe8e8c1
-config_hash: 5635033cdc8c930255f8b529a78de722
+config_hash: 4975e16a94e8f9901428022044131888

README.md

@@ -71,6 +71,109 @@ to add `OPENAI_API_KEY="My API Key"` to your `.env` file
 so that your API key is not stored in source control.
 [Get an API key here](https://platform.openai.com/settings/organization/api-keys).
 
+### Workload Identity Authentication
+
+For secure, automated environments like cloud-managed Kubernetes, Azure, and Google Cloud Platform, you can use workload identity authentication with short-lived tokens from cloud identity providers instead of long-lived API keys.
+
+#### Kubernetes (service account tokens)
+
+```python
+from openai import OpenAI
+from openai.auth import k8s_service_account_token_provider
+
+client = OpenAI(
+    workload_identity={
+        "client_id": "your-client-id",
+        "identity_provider_id": "idp-123",
+        "service_account_id": "sa-456",
+        "provider": k8s_service_account_token_provider(
+            "/var/run/secrets/kubernetes.io/serviceaccount/token"
+        ),
+    },
+    organization="org-xyz",
+    project="proj-abc",
+)
+
+response = client.chat.completions.create(
+    model="gpt-4",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+```
+
+#### Azure (managed identity)
+
+```python
+from openai import OpenAI
+from openai.auth import azure_managed_identity_token_provider
+
+client = OpenAI(
+    workload_identity={
+        "client_id": "your-client-id",
+        "identity_provider_id": "idp-123",
+        "service_account_id": "sa-456",
+        "provider": azure_managed_identity_token_provider(
+            resource="https://management.azure.com/",
+        ),
+    },
+)
+```
+
+#### Google Cloud Platform (compute engine metadata)
+
+```python
+from openai import OpenAI
+from openai.auth import gcp_id_token_provider
+
+client = OpenAI(
+    workload_identity={
+        "client_id": "your-client-id",
+        "identity_provider_id": "idp-123",
+        "service_account_id": "sa-456",
+        "provider": gcp_id_token_provider(audience="https://api.openai.com/v1"),
+    },
+)
+```
+
+#### Custom subject token provider
+
+```python
+from openai import OpenAI
+
+
+def get_custom_token() -> str:
+    return "your-jwt-token"
+
+
+client = OpenAI(
+    workload_identity={
+        "client_id": "your-client-id",
+        "identity_provider_id": "idp-123",
+        "service_account_id": "sa-456",
+        "provider": {
+            "token_type": "jwt",
+            "get_token": get_custom_token,
+        },
+    }
+)
+```
+
+You can also customize the token refresh buffer (default is 1200 seconds (20 minutes) before expiration):
+
+```python
+from openai import OpenAI
+from openai.auth import k8s_service_account_token_provider
+
+client = OpenAI(
+    workload_identity={
+        "client_id": "your-client-id",
+        "identity_provider_id": "idp-123",
+        "service_account_id": "sa-456",
+        "provider": k8s_service_account_token_provider("/var/token"),
+        "refresh_buffer_seconds": 120.0,
+    }
+)
+```
+
 ### Vision
 
 With an image URL:

api.md

@@ -11,6 +11,7 @@ from openai.types import (
     FunctionDefinition,
     FunctionParameters,
     Metadata,
+    OAuthErrorCode,
     Reasoning,
     ReasoningEffort,
     ResponseFormatJSONObject,

src/openai/__init__.py

@@ -16,6 +16,7 @@
 from ._constants import DEFAULT_TIMEOUT, DEFAULT_MAX_RETRIES, DEFAULT_CONNECTION_LIMITS
 from ._exceptions import (
     APIError,
+    OAuthError,
     OpenAIError,
     ConflictError,
     NotFoundError,
@@ -57,6 +58,7 @@
     "APIResponseValidationError",
     "BadRequestError",
     "AuthenticationError",
+    "OAuthError",
     "PermissionDeniedError",
     "NotFoundError",
     "ConflictError",

src/openai/_base_client.py

@@ -30,7 +30,7 @@
     cast,
     overload,
 )
-from typing_extensions import Literal, override, get_origin
+from typing_extensions import Unpack, Literal, override, get_origin
 
 import anyio
 import httpx
@@ -81,6 +81,7 @@
 )
 from ._streaming import Stream, SSEDecoder, AsyncStream, SSEBytesDecoder
 from ._exceptions import (
+    OpenAIError,
     APIStatusError,
     APITimeoutError,
     APIConnectionError,
@@ -936,6 +937,15 @@ def _prepare_request(
         """
         return None
 
+    def _send_request(
+        self,
+        request: httpx.Request,
+        *,
+        stream: bool,
+        **kwargs: Unpack[HttpxSendArgs],
+    ) -> httpx.Response:
+        return self._client.send(request, stream=stream, **kwargs)
+
     @overload
     def request(
         self,
@@ -1006,7 +1016,7 @@ def request(
 
             response = None
             try:
-                response = self._client.send(
+                response = self._send_request(
                     request,
                     stream=stream or self._should_stream_response_body(request=request),
                     **kwargs,
@@ -1025,6 +1035,9 @@ def request(
 
                 log.debug("Raising timeout error")
                 raise APITimeoutError(request=request) from err
+            except OpenAIError as err:
+                # Propagate OpenAIErrors as-is, without retrying or wrapping in APIConnectionError
+                raise err
             except Exception as err:
                 log.debug("Encountered Exception", exc_info=True)
 
@@ -1530,6 +1543,15 @@ async def _prepare_request(
         """
         return None
 
+    async def _send_request(
+        self,
+        request: httpx.Request,
+        *,
+        stream: bool,
+        **kwargs: Unpack[HttpxSendArgs],
+    ) -> httpx.Response:
+        return await self._client.send(request, stream=stream, **kwargs)
+
     @overload
     async def request(
         self,
@@ -1605,7 +1627,7 @@ async def request(
 
             response = None
             try:
-                response = await self._client.send(
+                response = await self._send_request(
                     request,
                     stream=stream or self._should_stream_response_body(request=request),
                     **kwargs,
@@ -1624,6 +1646,9 @@ async def request(
 
                 log.debug("Raising timeout error")
                 raise APITimeoutError(request=request) from err
+            except OpenAIError as err:
+                # Propagate OpenAIErrors as-is, without retrying or wrapping in APIConnectionError
+                raise err
             except Exception as err:
                 log.debug("Encountered Exception", exc_info=True)