feat: CV Application - AI-powered CV management, parsing, matching, and generation
This commit is contained in:
root
309
venv/lib/python3.12/site-packages/openai/auth/_workload.py
Normal file
309
venv/lib/python3.12/site-packages/openai/auth/_workload.py
Normal file
@@ -0,0 +1,309 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import threading
|
||||
from typing import Any, Callable, TypedDict, cast
|
||||
from pathlib import Path
|
||||
from typing_extensions import Literal, NotRequired
|
||||
|
||||
import httpx
|
||||
|
||||
from .._exceptions import OAuthError, OpenAIError, SubjectTokenProviderError
|
||||
from .._utils._sync import to_thread
|
||||
|
||||
TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
|
||||
DEFAULT_TOKEN_EXCHANGE_URL = "https://auth.openai.com/oauth/token"
|
||||
DEFAULT_REFRESH_BUFFER_SECONDS = 1200
|
||||
|
||||
SUBJECT_TOKEN_TYPES = {
|
||||
"jwt": "urn:ietf:params:oauth:token-type:jwt",
|
||||
"id": "urn:ietf:params:oauth:token-type:id_token",
|
||||
}
|
||||
|
||||
|
||||
class SubjectTokenProvider(TypedDict):
|
||||
token_type: Literal["jwt", "id"]
|
||||
get_token: Callable[[], str]
|
||||
|
||||
|
||||
class WorkloadIdentity(TypedDict):
|
||||
"""Identity provider resource id in WIFAPI."""
|
||||
|
||||
identity_provider_id: str
|
||||
|
||||
"""Service account id to bind the verified external identity to."""
|
||||
service_account_id: str
|
||||
|
||||
"""The provider configuration for obtaining the subject token."""
|
||||
provider: SubjectTokenProvider
|
||||
|
||||
"""Optional buffer time in seconds to refresh the OpenAI token before it expires. Defaults to 1200 seconds (20 minutes)."""
|
||||
refresh_buffer_seconds: NotRequired[float]
|
||||
|
||||
|
||||
def k8s_service_account_token_provider(
|
||||
token_file_path: str | Path = "/var/run/secrets/kubernetes.io/serviceaccount/token",
|
||||
) -> SubjectTokenProvider:
|
||||
"""
|
||||
Get a subject token provider for Kubernetes clusters with Workload Identity configured.
|
||||
|
||||
Cloud providers typically mount the subject token as a file in the container.
|
||||
|
||||
Args:
|
||||
token_file_path: path to the mounted service account token file. Defaults to `/var/run/secrets/kubernetes.io/serviceaccount/token`.
|
||||
"""
|
||||
|
||||
def get_token() -> str:
|
||||
try:
|
||||
with open(token_file_path, "r") as f:
|
||||
token = f.read().strip()
|
||||
if not token:
|
||||
raise SubjectTokenProviderError(f"The token file at {token_file_path} is empty.")
|
||||
return token
|
||||
except Exception as e:
|
||||
raise SubjectTokenProviderError(f"Failed to read the token file at {token_file_path}: {e}") from e
|
||||
|
||||
return {"token_type": "jwt", "get_token": get_token}
|
||||
|
||||
|
||||
def azure_managed_identity_token_provider(
|
||||
resource: str = "https://management.azure.com/",
|
||||
*,
|
||||
object_id: str | None = None,
|
||||
client_id: str | None = None,
|
||||
msi_res_id: str | None = None,
|
||||
api_version: str = "2018-02-01",
|
||||
timeout: float = 10.0,
|
||||
http_client: httpx.Client | None = None,
|
||||
) -> SubjectTokenProvider:
|
||||
"""
|
||||
Get a subject token provider for Azure Managed Identities.
|
||||
|
||||
See: https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http
|
||||
|
||||
Args:
|
||||
resource: the resource URI to request a token for. Defaults to `https://management.azure.com/` (Azure Resource Manager).
|
||||
object_id: the object ID of the managed identity to use, when multiple are assigned.
|
||||
client_id: the client ID of the managed identity to use, when multiple are assigned.
|
||||
msi_res_id: the ARM resource ID of the managed identity to use, when multiple are assigned.
|
||||
api_version: the Azure IMDS API version. Defaults to `2018-02-01`.
|
||||
timeout: the request timeout in seconds. Defaults to 10.0.
|
||||
http_client: optional httpx.Client instance to use for requests. If not provided, a new client will be created for each request.
|
||||
"""
|
||||
|
||||
def get_token() -> str:
|
||||
try:
|
||||
url = "http://169.254.169.254/metadata/identity/oauth2/token"
|
||||
params: dict[str, str] = {"api-version": api_version, "resource": resource}
|
||||
if object_id is not None:
|
||||
params["object_id"] = object_id
|
||||
if client_id is not None:
|
||||
params["client_id"] = client_id
|
||||
if msi_res_id is not None:
|
||||
params["msi_res_id"] = msi_res_id
|
||||
|
||||
if http_client is not None:
|
||||
response = http_client.get(url, params=params, headers={"Metadata": "true"}, timeout=timeout)
|
||||
else:
|
||||
with httpx.Client() as client:
|
||||
response = client.get(url, params=params, headers={"Metadata": "true"}, timeout=timeout)
|
||||
|
||||
if response.is_error:
|
||||
raise SubjectTokenProviderError(
|
||||
f"Failed to fetch Azure subject token from IMDS: HTTP {response.status_code}",
|
||||
response=response,
|
||||
)
|
||||
data = response.json()
|
||||
token = data.get("access_token")
|
||||
if not token:
|
||||
raise SubjectTokenProviderError(
|
||||
"Azure IMDS response did not include an access_token", response=response
|
||||
)
|
||||
return cast(str, token)
|
||||
except Exception as e:
|
||||
raise SubjectTokenProviderError(f"Failed to fetch Azure subject token from IMDS: {e}") from e
|
||||
|
||||
return {"token_type": "jwt", "get_token": get_token}
|
||||
|
||||
|
||||
def gcp_id_token_provider(
|
||||
audience: str = "https://api.openai.com/v1",
|
||||
*,
|
||||
timeout: float = 10.0,
|
||||
http_client: httpx.Client | None = None,
|
||||
) -> SubjectTokenProvider:
|
||||
"""
|
||||
Get a subject token provider for GCP VM instances using the instance metadata server.
|
||||
|
||||
See: https://cloud.google.com/compute/docs/instances/verifying-instance-identity
|
||||
|
||||
Args:
|
||||
audience: the unique URI agreed upon by both the instance and the system verifying
|
||||
the instance's identity. Defaults to `https://api.openai.com/v1`.
|
||||
timeout: the request timeout in seconds. Defaults to 10.0.
|
||||
http_client: optional httpx.Client instance to use for requests. If not provided, a new client will be created for each request.
|
||||
"""
|
||||
|
||||
def get_token() -> str:
|
||||
try:
|
||||
url = "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity"
|
||||
params = {"audience": audience}
|
||||
|
||||
if http_client is not None:
|
||||
response = http_client.get(url, params=params, headers={"Metadata-Flavor": "Google"}, timeout=timeout)
|
||||
else:
|
||||
with httpx.Client() as client:
|
||||
response = client.get(url, params=params, headers={"Metadata-Flavor": "Google"}, timeout=timeout)
|
||||
|
||||
if response.is_error:
|
||||
raise SubjectTokenProviderError(
|
||||
f"Failed to fetch GCP subject token from metadata server: HTTP {response.status_code}",
|
||||
response=response,
|
||||
)
|
||||
token = response.text.strip()
|
||||
if not token:
|
||||
raise SubjectTokenProviderError("GCP metadata server returned an empty token", response=response)
|
||||
return token
|
||||
except Exception as e:
|
||||
raise SubjectTokenProviderError(f"Failed to fetch GCP subject token from metadata server: {e}") from e
|
||||
|
||||
return {"token_type": "id", "get_token": get_token}
|
||||
|
||||
|
||||
class WorkloadIdentityAuth:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
workload_identity: WorkloadIdentity,
|
||||
token_exchange_url: str = DEFAULT_TOKEN_EXCHANGE_URL,
|
||||
):
|
||||
self.workload_identity = workload_identity
|
||||
self.token_exchange_url = token_exchange_url
|
||||
|
||||
self._cached_token: str | None = None
|
||||
self._cached_token_expires_at_monotonic: float | None = None
|
||||
self._cached_token_refresh_at_monotonic: float | None = None
|
||||
self._refreshing: bool = False
|
||||
self._lock = threading.Lock()
|
||||
self._condition = threading.Condition(self._lock)
|
||||
|
||||
def get_token(self) -> str:
|
||||
with self._lock:
|
||||
while self._refreshing and self._token_unusable():
|
||||
self._condition.wait()
|
||||
|
||||
if not self._token_unusable() and not self._needs_refresh():
|
||||
return cast(str, self._cached_token)
|
||||
|
||||
if self._refreshing:
|
||||
while self._refreshing:
|
||||
self._condition.wait()
|
||||
token = self._cached_token # type: ignore[unreachable]
|
||||
if self._token_unusable():
|
||||
raise RuntimeError("Token is unusable after refresh completed")
|
||||
return cast(str, token)
|
||||
|
||||
self._refreshing = True
|
||||
|
||||
try:
|
||||
self._perform_refresh()
|
||||
with self._lock:
|
||||
if self._token_unusable():
|
||||
raise RuntimeError("Token is unusable after refresh completed")
|
||||
return cast(str, self._cached_token)
|
||||
finally:
|
||||
with self._lock:
|
||||
self._refreshing = False
|
||||
self._condition.notify_all()
|
||||
|
||||
async def get_token_async(self) -> str:
|
||||
return await to_thread(self.get_token)
|
||||
|
||||
def invalidate_token(self) -> None:
|
||||
with self._lock:
|
||||
self._cached_token = None
|
||||
self._cached_token_expires_at_monotonic = None
|
||||
self._cached_token_refresh_at_monotonic = None
|
||||
|
||||
def _perform_refresh(self) -> None:
|
||||
token_data = self._fetch_token_from_exchange()
|
||||
now = time.monotonic()
|
||||
expires_in = token_data["expires_in"]
|
||||
|
||||
with self._lock:
|
||||
self._cached_token = token_data["access_token"]
|
||||
self._cached_token_expires_at_monotonic = now + expires_in
|
||||
self._cached_token_refresh_at_monotonic = now + self._refresh_delay_seconds(expires_in)
|
||||
|
||||
def _fetch_token_from_exchange(self) -> dict[str, Any]:
|
||||
subject_token = self._get_subject_token()
|
||||
|
||||
token_type = self.workload_identity["provider"]["token_type"]
|
||||
subject_token_type = SUBJECT_TOKEN_TYPES.get(token_type)
|
||||
if subject_token_type is None:
|
||||
raise OpenAIError(
|
||||
f"Unsupported token type: {token_type!r}. Supported types: {', '.join(SUBJECT_TOKEN_TYPES.keys())}"
|
||||
)
|
||||
|
||||
with httpx.Client() as client:
|
||||
response = client.post(
|
||||
self.token_exchange_url,
|
||||
json={
|
||||
"grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
|
||||
"subject_token": subject_token,
|
||||
"subject_token_type": subject_token_type,
|
||||
"identity_provider_id": self.workload_identity["identity_provider_id"],
|
||||
"service_account_id": self.workload_identity["service_account_id"],
|
||||
},
|
||||
timeout=10.0,
|
||||
)
|
||||
return self._handle_token_response(response)
|
||||
|
||||
def _handle_token_response(self, response: httpx.Response) -> dict[str, Any]:
|
||||
try:
|
||||
body = response.json() if response.content else None
|
||||
except ValueError:
|
||||
body = None
|
||||
|
||||
if response.status_code in (400, 401, 403):
|
||||
raise OAuthError(response=response, body=body)
|
||||
|
||||
if response.is_success:
|
||||
if body is None:
|
||||
raise OpenAIError("Token exchange succeeded but response body was empty")
|
||||
access_token = body.get("access_token")
|
||||
expires_in = body.get("expires_in")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
raise OpenAIError("Token exchange response did not include a valid access_token")
|
||||
if not isinstance(expires_in, (int, float)):
|
||||
raise OpenAIError("Token exchange response did not include a valid expires_in")
|
||||
return {"access_token": access_token, "expires_in": float(expires_in)}
|
||||
|
||||
raise OpenAIError(
|
||||
f"Token exchange failed with status {response.status_code}",
|
||||
)
|
||||
|
||||
def _get_subject_token(self) -> str:
|
||||
provider = self.workload_identity["provider"]
|
||||
subject_token = provider["get_token"]()
|
||||
if not subject_token:
|
||||
raise OpenAIError("The workload identity provider returned an empty subject token")
|
||||
return subject_token
|
||||
|
||||
def _token_unusable(self) -> bool:
|
||||
return self._cached_token is None or self._token_expired()
|
||||
|
||||
def _token_expired(self) -> bool:
|
||||
if self._cached_token_expires_at_monotonic is None:
|
||||
return True
|
||||
return time.monotonic() >= self._cached_token_expires_at_monotonic
|
||||
|
||||
def _needs_refresh(self) -> bool:
|
||||
if self._cached_token_refresh_at_monotonic is None:
|
||||
return False
|
||||
return time.monotonic() >= self._cached_token_refresh_at_monotonic
|
||||
|
||||
def _refresh_delay_seconds(self, expires_in: float) -> float:
|
||||
configured_buffer = self.workload_identity.get("refresh_buffer_seconds", DEFAULT_REFRESH_BUFFER_SECONDS)
|
||||
effective_buffer = min(configured_buffer, expires_in / 2)
|
||||
return max(expires_in - effective_buffer, 0.0)
|
||||
Reference in New Issue
Block a user