diff --git a/litellm/__init__.py b/litellm/__init__.py index e5c09702b9bc..5fb3bd025051 100644 --- a/litellm/__init__.py +++ b/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 ### diff --git a/litellm/main.py b/litellm/main.py index ce84c8988e02..97065e59601a 100644 --- a/litellm/main.py +++ b/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) diff --git a/litellm/proxy_auth/__init__.py b/litellm/proxy_auth/__init__.py new file mode 100644 index 000000000000..27624a94fb9d --- /dev/null +++ b/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", +] diff --git a/litellm/proxy_auth/credentials.py b/litellm/proxy_auth/credentials.py new file mode 100644 index 000000000000..cddaf1278f9a --- /dev/null +++ b/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}"} diff --git a/tests/litellm/test_proxy_auth.py b/tests/litellm/test_proxy_auth.py new file mode 100644 index 000000000000..1d73e143e102 --- /dev/null +++ b/tests/litellm/test_proxy_auth.py @@ -0,0 +1,204 @@ +""" +Unit tests for litellm.proxy_auth module. + +Tests the OAuth2/JWT token management for LiteLLM Proxy authentication. +""" + +import time +from unittest.mock import Mock, patch + +import pytest + +from litellm.proxy_auth import ( + AccessToken, + AzureADCredential, + GenericOAuth2Credential, + ProxyAuthHandler, +) + + +class TestAccessToken: + """Tests for AccessToken dataclass.""" + + def test_access_token_creation(self): + """Test AccessToken can be created with required fields.""" + token = AccessToken(token="test-token", expires_on=1234567890) + assert token.token == "test-token" + assert token.expires_on == 1234567890 + + def test_access_token_equality(self): + """Test AccessToken equality comparison.""" + token1 = AccessToken(token="test", expires_on=123) + token2 = AccessToken(token="test", expires_on=123) + assert token1 == token2 + + +class MockCredential: + """Mock credential for testing.""" + + def __init__(self, expires_in_seconds: int = 3600): + self.call_count = 0 + self.expires_in = expires_in_seconds + + def get_token(self, scope: str) -> AccessToken: + self.call_count += 1 + return AccessToken( + token=f"mock-token-{self.call_count}", + expires_on=int(time.time()) + self.expires_in, + ) + + +class TestProxyAuthHandler: + """Tests for ProxyAuthHandler.""" + + def test_get_auth_headers_returns_bearer_token(self): + """Test that get_auth_headers returns correct Authorization header.""" + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + headers = handler.get_auth_headers() + + assert "Authorization" in headers + assert headers["Authorization"].startswith("Bearer ") + assert "mock-token-1" in headers["Authorization"] + + def test_token_caching(self): + """Test that tokens are cached and not re-requested.""" + cred = MockCredential(expires_in_seconds=3600) # Long expiry + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + # Multiple calls should only request token once + handler.get_auth_headers() + handler.get_auth_headers() + handler.get_auth_headers() + + assert cred.call_count == 1 + + def test_token_refresh_when_about_to_expire(self): + """Test that tokens are refreshed when about to expire (within 60s buffer).""" + cred = MockCredential(expires_in_seconds=30) # Expires in 30s (< 60s buffer) + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + # First call gets token + handler.get_auth_headers() + # Second call should refresh because token expires within 60s buffer + handler.get_auth_headers() + + assert cred.call_count == 2 + + def test_get_token_method(self): + """Test the get_token method returns AccessToken.""" + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + token = handler.get_token() + + assert isinstance(token, AccessToken) + assert token.token == "mock-token-1" + + +class TestAzureADCredential: + """Tests for AzureADCredential.""" + + def test_lazy_initialization(self): + """Test that azure-identity is not imported until get_token is called.""" + # This should not raise ImportError even if azure-identity is not installed + cred = AzureADCredential(credential=None) + # _initialized should be False until get_token is called + assert cred._initialized is False + + def test_wraps_azure_credential(self): + """Test that AzureADCredential wraps an azure-identity credential.""" + # Mock Azure credential + mock_azure_cred = Mock() + mock_azure_cred.get_token.return_value = Mock( + token="azure-token", expires_on=9999999999 + ) + + cred = AzureADCredential(credential=mock_azure_cred) + token = cred.get_token("https://graph.microsoft.com/.default") + + assert token.token == "azure-token" + assert token.expires_on == 9999999999 + mock_azure_cred.get_token.assert_called_once_with( + "https://graph.microsoft.com/.default" + ) + + +class TestGenericOAuth2Credential: + """Tests for GenericOAuth2Credential.""" + + def test_token_request(self): + """Test that GenericOAuth2Credential makes correct OAuth2 request.""" + with patch("httpx.post") as mock_post: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "oauth2-token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + cred = GenericOAuth2Credential( + client_id="test-client", + client_secret="test-secret", + token_url="https://example.com/oauth2/token", + ) + token = cred.get_token("test-scope") + + assert token.token == "oauth2-token" + mock_post.assert_called_once() + call_kwargs = mock_post.call_args + assert call_kwargs[1]["data"]["grant_type"] == "client_credentials" + assert call_kwargs[1]["data"]["client_id"] == "test-client" + assert call_kwargs[1]["data"]["client_secret"] == "test-secret" + assert call_kwargs[1]["data"]["scope"] == "test-scope" + + def test_token_caching(self): + """Test that GenericOAuth2Credential caches tokens.""" + with patch("httpx.post") as mock_post: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "oauth2-token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + cred = GenericOAuth2Credential( + client_id="test-client", + client_secret="test-secret", + token_url="https://example.com/oauth2/token", + ) + + # Multiple calls should only make one HTTP request + cred.get_token("test-scope") + cred.get_token("test-scope") + cred.get_token("test-scope") + + assert mock_post.call_count == 1 + + +class TestLiteLLMIntegration: + """Tests for integration with litellm module.""" + + def test_proxy_auth_variable_exists(self): + """Test that litellm.proxy_auth variable exists.""" + import litellm + + # Should be None by default + assert hasattr(litellm, "proxy_auth") + + def test_proxy_auth_can_be_set(self): + """Test that litellm.proxy_auth can be set to a ProxyAuthHandler.""" + import litellm + + original_value = litellm.proxy_auth + try: + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test") + litellm.proxy_auth = handler + + assert litellm.proxy_auth is handler + finally: + litellm.proxy_auth = original_value