import requests from typing import Dict, Any, Optional, Union, List from .exceptions import UnauthorizedError class KeysManagementClient: def __init__(self, base_url: str, api_key: Optional[str] = None): """ Initialize the KeysManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key def _get_headers(self) -> Dict[str, str]: """ Get the headers for API requests, including authorization if api_key is set. Returns: Dict[str, str]: Headers to use for API requests """ headers = {"Content-Type": "application/json"} if self._api_key: headers["Authorization"] = f"Bearer {self._api_key}" return headers def list( self, page: Optional[int] = None, size: Optional[int] = None, user_id: Optional[str] = None, team_id: Optional[str] = None, organization_id: Optional[str] = None, key_hash: Optional[str] = None, key_alias: Optional[str] = None, return_full_object: Optional[bool] = None, include_team_keys: Optional[bool] = None, return_request: bool = False, ) -> Union[Dict[str, Any], requests.Request]: """ List all API keys with optional filtering and pagination. Args: page (Optional[int]): Page number for pagination size (Optional[int]): Number of items per page user_id (Optional[str]): Filter keys by user ID team_id (Optional[str]): Filter keys by team ID organization_id (Optional[str]): Filter keys by organization ID key_hash (Optional[str]): Filter by specific key hash key_alias (Optional[str]): Filter by key alias return_full_object (Optional[bool]): Whether to return the full key object include_team_keys (Optional[bool]): Whether to include team keys in the response return_request (bool): If True, returns the prepared request object instead of executing it Returns: Union[Dict[str, Any], requests.Request]: Either the response from the server or a prepared request object if return_request is True. The response contains a list of API keys with their configurations. Raises: UnauthorizedError: If the request fails with a 401 status code requests.exceptions.RequestException: If the request fails with any other error """ url = f"{self._base_url}/key/list" params: Dict[str, Any] = {} # Add optional query parameters if page is not None: params["page"] = page if size is not None: params["size"] = size if user_id is not None: params["user_id"] = user_id if team_id is not None: params["team_id"] = team_id if organization_id is not None: params["organization_id"] = organization_id if key_hash is not None: params["key_hash"] = key_hash if key_alias is not None: params["key_alias"] = key_alias if return_full_object is not None: params["return_full_object"] = str(return_full_object).lower() if include_team_keys is not None: params["include_team_keys"] = str(include_team_keys).lower() request = requests.Request("GET", url, headers=self._get_headers(), params=params) if return_request: return request session = requests.Session() try: response = session.send(request.prepare()) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) raise def generate( self, models: Optional[List[str]] = None, aliases: Optional[Dict[str, str]] = None, spend: Optional[float] = None, duration: Optional[str] = None, key_alias: Optional[str] = None, team_id: Optional[str] = None, user_id: Optional[str] = None, budget_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None, return_request: bool = False, ) -> Union[Dict[str, Any], requests.Request]: """ Generate an API key based on the provided data. Docs: https://docs.litellm.ai/docs/proxy/virtual_keys Args: models (Optional[List[str]]): List of allowed models for this key aliases (Optional[Dict[str, str]]): Model alias mappings spend (Optional[float]): Maximum spend limit for this key duration (Optional[str]): Duration for which the key is valid (e.g. "24h", "7d") key_alias (Optional[str]): Alias/name for the key for easier identification team_id (Optional[str]): Team ID to associate the key with user_id (Optional[str]): User ID to associate the key with budget_id (Optional[str]): Budget ID to associate the key with config (Optional[Dict[str, Any]]): Additional configuration parameters return_request (bool): If True, returns the prepared request object instead of executing it Returns: Union[Dict[str, Any], requests.Request]: Either the response from the server or a prepared request object if return_request is True Raises: UnauthorizedError: If the request fails with a 401 status code requests.exceptions.RequestException: If the request fails with any other error """ url = f"{self._base_url}/key/generate" data: Dict[str, Any] = {} if models is not None: data["models"] = models if aliases is not None: data["aliases"] = aliases if spend is not None: data["spend"] = spend if duration is not None: data["duration"] = duration if key_alias is not None: data["key_alias"] = key_alias if team_id is not None: data["team_id"] = team_id if user_id is not None: data["user_id"] = user_id if budget_id is not None: data["budget_id"] = budget_id if config is not None: data["config"] = config request = requests.Request("POST", url, headers=self._get_headers(), json=data) if return_request: return request session = requests.Session() try: response = session.send(request.prepare()) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) raise def delete( self, keys: Optional[List[str]] = None, key_aliases: Optional[List[str]] = None, return_request: bool = False, ) -> Union[Dict[str, Any], requests.Request]: """ Delete existing keys Args: keys (List[str]): List of API keys to delete key_aliases (List[str]): List of key aliases to delete return_request (bool): If True, returns the prepared request object instead of executing it Returns: Union[Dict[str, Any], requests.Request]: Either the response from the server or a prepared request object if return_request is True Raises: UnauthorizedError: If the request fails with a 401 status code requests.exceptions.RequestException: If the request fails with any other error """ url = f"{self._base_url}/key/delete" data = { "keys": keys, "key_aliases": key_aliases, } request = requests.Request("POST", url, headers=self._get_headers(), json=data) if return_request: return request session = requests.Session() try: response = session.send(request.prepare()) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) raise