mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
f747bc6704
Second batch of the fifth basedpyright Any reduction round. Every change is typing-only and leaves runtime behavior identical. Guardrail hooks and file, vector store and usage endpoints move their payload, header and response annotations from Any to object, Mapping[str, object] or the concrete response model the call site already produces. Two private aggregation helpers in the guardrail usage endpoints take a key accessor function instead of an attribute name string, so the key they read is checked against the row type. The verification token repository reaches its two tables through Protocols that name the handles it calls, rather than reading them off an untyped prisma client, and the Azure AD credential wrapper describes the azure-identity credential it wraps the same way.
265 lines
8.2 KiB
Python
265 lines
8.2 KiB
Python
"""
|
|
Credential providers for proxy authentication.
|
|
|
|
This module provides a provider-agnostic interface for obtaining OAuth2/JWT tokens.
|
|
It follows the same TokenCredential protocol used by Azure SDK.
|
|
"""
|
|
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Final, Protocol, runtime_checkable
|
|
|
|
|
|
@dataclass
|
|
class AccessToken:
|
|
"""
|
|
Represents an OAuth2 access token with expiration.
|
|
|
|
This matches the structure used by azure.core.credentials.AccessToken.
|
|
|
|
Attributes:
|
|
token: The access token string (typically a JWT).
|
|
expires_on: Unix timestamp when the token expires.
|
|
"""
|
|
|
|
token: str
|
|
expires_on: int
|
|
|
|
|
|
@runtime_checkable
|
|
class TokenCredential(Protocol):
|
|
"""
|
|
Protocol for credential providers.
|
|
|
|
This matches the azure.core.credentials.TokenCredential interface,
|
|
allowing any Azure SDK credential to be used directly.
|
|
|
|
Any class implementing get_token(scope) -> AccessToken can be used.
|
|
"""
|
|
|
|
def get_token(self, scope: str) -> AccessToken:
|
|
"""
|
|
Get an access token for the specified scope.
|
|
|
|
Args:
|
|
scope: The OAuth2 scope to request (e.g., "api://my-app/.default")
|
|
|
|
Returns:
|
|
AccessToken with the token string and expiration timestamp.
|
|
"""
|
|
...
|
|
|
|
|
|
class _AzureAccessToken(Protocol):
|
|
"""The two attributes :class:`AzureADCredential` reads off an azure-identity token."""
|
|
|
|
@property
|
|
def token(self) -> str: ...
|
|
|
|
@property
|
|
def expires_on(self) -> int: ...
|
|
|
|
|
|
class _AzureTokenCredential(Protocol):
|
|
"""The single method :class:`AzureADCredential` calls on the credential it wraps."""
|
|
|
|
def get_token(self, *scopes: str) -> _AzureAccessToken: ...
|
|
|
|
|
|
class AzureADCredential:
|
|
"""
|
|
Wrapper for Azure Identity credentials.
|
|
|
|
This wraps any azure-identity credential (DefaultAzureCredential,
|
|
ClientSecretCredential, ManagedIdentityCredential, etc.) and converts
|
|
the token to our AccessToken format.
|
|
|
|
If no credential is provided, it will use DefaultAzureCredential
|
|
which tries multiple authentication methods automatically.
|
|
|
|
Example:
|
|
# Use default credential chain (env vars, managed identity, CLI, etc.)
|
|
cred = AzureADCredential()
|
|
|
|
# Or provide a specific credential
|
|
from azure.identity import ClientSecretCredential
|
|
azure_cred = ClientSecretCredential(tenant_id, client_id, client_secret)
|
|
cred = AzureADCredential(credential=azure_cred)
|
|
"""
|
|
|
|
def __init__(self, credential: _AzureTokenCredential | None = None):
|
|
"""
|
|
Initialize with an optional Azure credential.
|
|
|
|
Args:
|
|
credential: An azure-identity credential object. If None,
|
|
DefaultAzureCredential will be used on first token request.
|
|
"""
|
|
self._credential: _AzureTokenCredential | None = credential
|
|
self._initialized = credential is not None
|
|
|
|
def get_token(self, scope: str) -> AccessToken:
|
|
"""
|
|
Get an access token from Azure AD.
|
|
|
|
Args:
|
|
scope: The OAuth2 scope (e.g., "api://my-app/.default")
|
|
|
|
Returns:
|
|
AccessToken with the JWT and expiration.
|
|
|
|
Raises:
|
|
ImportError: If azure-identity is not installed.
|
|
"""
|
|
result: Final = self._resolve_credential().get_token(scope)
|
|
return AccessToken(token=result.token, expires_on=result.expires_on)
|
|
|
|
def _resolve_credential(self) -> _AzureTokenCredential:
|
|
"""Return the wrapped credential, building the Azure default chain on first use.
|
|
|
|
Raises:
|
|
ImportError: If azure-identity is not installed.
|
|
"""
|
|
existing: Final = self._credential
|
|
if existing is not None:
|
|
return existing
|
|
try:
|
|
from azure.identity import DefaultAzureCredential
|
|
|
|
created: Final = DefaultAzureCredential()
|
|
except ImportError:
|
|
raise ImportError(
|
|
"azure-identity is required for AzureADCredential. Install it with: pip install azure-identity"
|
|
)
|
|
self._credential = created
|
|
self._initialized = True
|
|
return created
|
|
|
|
|
|
class GenericOAuth2Credential:
|
|
"""
|
|
Generic OAuth2 client credentials flow.
|
|
|
|
This works with any OAuth2 provider (Okta, Auth0, Keycloak, etc.)
|
|
that supports the client_credentials grant type.
|
|
|
|
Example:
|
|
cred = GenericOAuth2Credential(
|
|
client_id="my-client-id",
|
|
client_secret="my-client-secret",
|
|
token_url="https://my-idp.com/oauth2/token"
|
|
)
|
|
"""
|
|
|
|
def __init__(self, client_id: str, client_secret: str, token_url: str):
|
|
"""
|
|
Initialize OAuth2 client credentials.
|
|
|
|
Args:
|
|
client_id: OAuth2 client ID
|
|
client_secret: OAuth2 client secret
|
|
token_url: Token endpoint URL (e.g., "https://idp.com/oauth2/token")
|
|
"""
|
|
self.client_id = client_id
|
|
self.client_secret = client_secret
|
|
self.token_url = token_url
|
|
self._cached_token: AccessToken | None = None
|
|
|
|
def get_token(self, scope: str) -> AccessToken:
|
|
"""
|
|
Get an access token using OAuth2 client credentials flow.
|
|
|
|
Tokens are cached and reused until they expire (with 60s buffer).
|
|
|
|
Args:
|
|
scope: The OAuth2 scope to request
|
|
|
|
Returns:
|
|
AccessToken with the token and expiration.
|
|
"""
|
|
# Return cached token if still valid (with 60s buffer)
|
|
if self._cached_token and self._cached_token.expires_on > time.time() + 60:
|
|
return self._cached_token
|
|
|
|
import httpx
|
|
|
|
response: Final = httpx.post(
|
|
self.token_url,
|
|
data={
|
|
"grant_type": "client_credentials",
|
|
"client_id": self.client_id,
|
|
"client_secret": self.client_secret,
|
|
"scope": scope,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
data: Final = response.json()
|
|
|
|
self._cached_token = AccessToken(
|
|
token=data["access_token"],
|
|
expires_on=int(time.time()) + data.get("expires_in", 3600),
|
|
)
|
|
return self._cached_token
|
|
|
|
|
|
class ProxyAuthHandler:
|
|
"""
|
|
Manages OAuth2/JWT token lifecycle for proxy authentication.
|
|
|
|
This handler:
|
|
- Obtains tokens from the configured credential provider
|
|
- Caches tokens to avoid unnecessary requests
|
|
- Automatically refreshes tokens before they expire (60s buffer)
|
|
- Generates Authorization headers for HTTP requests
|
|
|
|
Set this as litellm.proxy_auth to automatically inject auth headers
|
|
into all requests to your LiteLLM Proxy.
|
|
|
|
Example:
|
|
import litellm
|
|
from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler
|
|
|
|
litellm.proxy_auth = ProxyAuthHandler(
|
|
credential=AzureADCredential(),
|
|
scope="api://my-litellm-proxy/.default"
|
|
)
|
|
litellm.api_base = "https://my-proxy.example.com"
|
|
|
|
# Auth headers are now automatically injected
|
|
response = litellm.completion(model="gpt-4", messages=[...])
|
|
"""
|
|
|
|
def __init__(self, credential: TokenCredential, scope: str):
|
|
"""
|
|
Initialize the proxy auth handler.
|
|
|
|
Args:
|
|
credential: A TokenCredential implementation (AzureADCredential,
|
|
GenericOAuth2Credential, or any custom implementation)
|
|
scope: The OAuth2 scope to request tokens for
|
|
"""
|
|
self.credential = credential
|
|
self.scope = scope
|
|
self._cached_token: AccessToken | None = None
|
|
|
|
def get_token(self) -> AccessToken:
|
|
"""
|
|
Get a valid access token, refreshing if necessary.
|
|
|
|
Returns:
|
|
AccessToken that is valid for at least 60 more seconds.
|
|
"""
|
|
# Refresh if no token or token expires within 60 seconds
|
|
if not self._cached_token or self._cached_token.expires_on <= time.time() + 60:
|
|
self._cached_token = self.credential.get_token(self.scope)
|
|
return self._cached_token
|
|
|
|
def get_auth_headers(self) -> dict[str, str]:
|
|
"""
|
|
Get HTTP headers for authentication.
|
|
|
|
Returns:
|
|
Dict with Authorization header containing Bearer token.
|
|
"""
|
|
token: Final = self.get_token()
|
|
return {"Authorization": f"Bearer {token.token}"}
|