feat(sdk): add proxy_auth for auto OAuth2/JWT token management

litellm/__init__.py

@@ -258,6 +258,8 @@
 in_memory_llm_clients_cache: "LLMClientCache"
 safe_memory_mode: bool = False
 enable_azure_ad_token_refresh: Optional[bool] = False
+# Proxy Authentication - auto-obtain/refresh OAuth2/JWT tokens for LiteLLM Proxy
+proxy_auth: Optional[Any] = None
 ### DEFAULT AZURE API VERSION ###
 AZURE_DEFAULT_API_VERSION = "2025-02-01-preview"  # this is updated to the latest
 ### DEFAULT WATSONX API VERSION ###

litellm/main.py

@@ -1213,6 +1213,13 @@ def completion(  # type: ignore # noqa: PLR0915
         headers = {}
     if extra_headers is not None:
         headers.update(extra_headers)
+    # Inject proxy auth headers if configured
+    if litellm.proxy_auth is not None:
+        try:
+            proxy_headers = litellm.proxy_auth.get_auth_headers()
+            headers.update(proxy_headers)
+        except Exception as e:
+            verbose_logger.warning(f"Failed to get proxy auth headers: {e}")
     num_retries = kwargs.get(
         "num_retries", None
     )  ## alt. param for 'max_retries'. Use this to pass retries w/ instructor.
@@ -4542,6 +4549,13 @@ def embedding(  # noqa: PLR0915
         headers = {}
     if extra_headers is not None:
         headers.update(extra_headers)
+    # Inject proxy auth headers if configured
+    if litellm.proxy_auth is not None:
+        try:
+            proxy_headers = litellm.proxy_auth.get_auth_headers()
+            headers.update(proxy_headers)
+        except Exception as e:
+            verbose_logger.warning(f"Failed to get proxy auth headers: {e}")
     ### CUSTOM MODEL COST ###
     input_cost_per_token = kwargs.get("input_cost_per_token", None)
     output_cost_per_token = kwargs.get("output_cost_per_token", None)

litellm/proxy_auth/__init__.py

@@ -0,0 +1,30 @@
+"""
+Proxy Authentication module for LiteLLM SDK.
+
+This module provides OAuth2/JWT token management for authenticating
+with LiteLLM Proxy or any OAuth2-protected endpoint.
+
+Usage:
+    from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler
+
+    litellm.proxy_auth = ProxyAuthHandler(
+        credential=AzureADCredential(),
+        scope="api://my-proxy/.default"
+    )
+"""
+
+from .credentials import (
+    AccessToken,
+    TokenCredential,
+    AzureADCredential,
+    GenericOAuth2Credential,
+    ProxyAuthHandler,
+)
+
+__all__ = [
+    "AccessToken",
+    "TokenCredential",
+    "AzureADCredential",
+    "GenericOAuth2Credential",
+    "ProxyAuthHandler",
+]

litellm/proxy_auth/credentials.py

@@ -0,0 +1,240 @@
+"""
+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 Any, Optional, 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 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: Optional[Any] = 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 = 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.
+        """
+        if not self._initialized:
+            try:
+                from azure.identity import DefaultAzureCredential
+
+                self._credential = DefaultAzureCredential()
+                self._initialized = True
+            except ImportError:
+                raise ImportError(
+                    "azure-identity is required for AzureADCredential. "
+                    "Install it with: pip install azure-identity"
+                )
+
+        result = self._credential.get_token(scope)
+        return AccessToken(token=result.token, expires_on=result.expires_on)
+
+
+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: Optional[AccessToken] = 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 = 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 = 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: Optional[AccessToken] = 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:
+        """
+        Get HTTP headers for authentication.
+
+        Returns:
+            Dict with Authorization header containing Bearer token.
+        """
+        token = self.get_token()
+        return {"Authorization": f"Bearer {token.token}"}