litellm/litellm/proxy/client/credentials.py

from collections.abc import Mapping
from typing import Any, Final

import requests

from .exceptions import UnauthorizedError


class CredentialsManagementClient:
    def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
        """
        Initialize the CredentialsManagementClient.

        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.
            timeout (int): Request timeout in seconds (default: 30)
        """
        self._base_url = base_url.rstrip("/")  # Remove trailing slash if present
        self._api_key = api_key
        self._timeout = timeout

    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: Final = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        return headers

    def list(
        self,
        return_request: bool = False,
    ) -> dict[str, Any] | requests.Request:
        """
        List all credentials.

        Args:
            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: Final = f"{self._base_url}/credentials"

        request: Final = requests.Request("GET", url, headers=self._get_headers())

        if return_request:
            return request

        session: Final = requests.Session()
        try:
            response: Final = session.send(request.prepare(), timeout=self._timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            raise

    def create(
        self,
        credential_name: str,
        credential_info: Mapping[str, object],
        credential_values: Mapping[str, object],
        return_request: bool = False,
    ) -> dict[str, Any] | requests.Request:
        """
        Create a new credential.

        Args:
            credential_name (str): Name of the credential
            credential_info (Dict[str, Any]): Additional information about the credential
            credential_values (Dict[str, Any]): Values for the credential
            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: Final = f"{self._base_url}/credentials"

        data: Final = {
            "credential_name": credential_name,
            "credential_info": credential_info,
            "credential_values": credential_values,
        }

        request: Final = requests.Request("POST", url, headers=self._get_headers(), json=data)

        if return_request:
            return request

        session: Final = requests.Session()
        try:
            response: Final = session.send(request.prepare(), timeout=self._timeout)
            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,
        credential_name: str,
        return_request: bool = False,
    ) -> dict[str, Any] | requests.Request:
        """
        Delete a credential by name.

        Args:
            credential_name (str): Name of the credential 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: Final = f"{self._base_url}/credentials/{credential_name}"

        request: Final = requests.Request("DELETE", url, headers=self._get_headers())

        if return_request:
            return request

        session: Final = requests.Session()
        try:
            response: Final = session.send(request.prepare(), timeout=self._timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            raise

    def get(
        self,
        credential_name: str,
        return_request: bool = False,
    ) -> dict[str, Any] | requests.Request:
        """
        Get a credential by name.

        Args:
            credential_name (str): Name of the credential to retrieve
            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: Final = f"{self._base_url}/credentials/by_name/{credential_name}"

        request: Final = requests.Request("GET", url, headers=self._get_headers())

        if return_request:
            return request

        session: Final = requests.Session()
        try:
            response: Final = session.send(request.prepare(), timeout=self._timeout)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                raise UnauthorizedError(e)
            raise