From aac9907ab8cd8acb06da79650fa0f73c0a34579c Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Feb 2026 04:09:33 +0530 Subject: [PATCH] feat: Per Model Budgets and Limits (#20076) --- litellm/proxy/_types.py | 8008 +++++++-------- .../proxy/hooks/model_max_budget_limiter.py | 400 +- .../budget_management_endpoints.py | 673 +- .../key_management_endpoints.py | 8961 +++++++++-------- tests/proxy_unit_tests/test_proxy_utils.py | 145 +- ...test_unit_test_max_model_budget_limiter.py | 63 +- .../test_budget_endpoints.py | 54 +- 7 files changed, 9250 insertions(+), 9054 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bf99347ef6e..c91a8ffeaa6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,4000 +1,4008 @@ -import enum -import json -from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union - -import httpx -from pydantic import ( - BaseModel, - ConfigDict, - Field, - Json, - field_validator, - model_validator, -) -from typing_extensions import Required, TypedDict - -from litellm._uuid import uuid -from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.llms.openai import ( - AllMessageValues, - OpenAIFileObject, - ResponsesAPIResponse, -) -from litellm.types.mcp import ( - MCPAuth, - MCPAuthType, - MCPCredentials, - MCPTransport, - MCPTransportType, -) -from litellm.types.mcp_server.mcp_server_manager import MCPInfo -from litellm.types.router import RouterErrors, UpdateRouterConfig -from litellm.types.secret_managers.main import KeyManagementSystem -from litellm.types.utils import ( - CallTypes, - CostBreakdown, - EmbeddingResponse, - GenericBudgetConfigType, - ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, - LiteLLMPydanticObjectBase, - ModelResponse, - ProviderField, - StandardCallbackDynamicParams, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingVectorStoreRequest, - StandardPassThroughResponseObject, - TextCompletionResponse, -) -from litellm.types.videos.main import VideoObject - -from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type - -if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - Span = Union[_Span, Any] -else: - Span = Any - - -class SupportedDBObjectType(str, enum.Enum): - """ - Supported database object types for fine-grained DB storage control. - Use in general_settings.supported_db_objects to specify which objects to load from DB. - """ - - MODELS = "models" - MCP = "mcp" - GUARDRAILS = "guardrails" - POLICIES = "policies" - VECTOR_STORES = "vector_stores" - PASS_THROUGH_ENDPOINTS = "pass_through_endpoints" - PROMPTS = "prompts" - MODEL_COST_MAP = "model_cost_map" - - def __str__(self): - return str(self.value) - - -class LiteLLMTeamRoles(enum.Enum): - # team admin - TEAM_ADMIN = "admin" - # team member - TEAM_MEMBER = "user" - - -class LitellmUserRoles(str, enum.Enum): - """ - Admin Roles: - PROXY_ADMIN: admin over the platform - PROXY_ADMIN_VIEW_ONLY: can login, view all own keys, view all spend - ORG_ADMIN: admin over a specific organization, can create teams, users only within their organization - - Internal User Roles: - INTERNAL_USER: can login, view/create/delete their own keys, view their spend - INTERNAL_USER_VIEW_ONLY: can login, view their own keys, view their own spend - - - Team Roles: - TEAM: used for JWT auth - - - Customer Roles: - CUSTOMER: External users -> these are customers - - """ - - # Admin Roles - PROXY_ADMIN = "proxy_admin" - PROXY_ADMIN_VIEW_ONLY = "proxy_admin_viewer" - - # Organization admins - ORG_ADMIN = "org_admin" - - # Internal User Roles - INTERNAL_USER = "internal_user" - INTERNAL_USER_VIEW_ONLY = "internal_user_viewer" - - # Team Roles - TEAM = "team" - - # Customer Roles - External users of proxy - CUSTOMER = "customer" - - def __str__(self): - return str(self.value) - - def values(self) -> List[str]: - return list(self.__annotations__.keys()) - - @property - def description(self): - """ - Descriptions for the enum values - """ - descriptions = { - "proxy_admin": "admin over litellm proxy, has all permissions", - "proxy_admin_viewer": "view all keys, view all spend", - "internal_user": "view/create/delete their own keys, view their own spend", - "internal_user_viewer": "view their own keys, view their own spend", - "team": "team scope used for JWT auth", - "customer": "customer", - } - return descriptions.get(self.value, "") - - @property - def ui_label(self): - """ - UI labels for the enum values - """ - ui_labels = { - "proxy_admin": "Admin (All Permissions)", - "proxy_admin_viewer": "Admin (View Only)", - "internal_user": "Internal User (Create/Delete/View)", - "internal_user_viewer": "Internal User (View Only)", - "team": "Team", - "customer": "Customer", - } - return ui_labels.get(self.value, "") - - @property - def is_internal_user_role(self) -> bool: - """returns true if this role is an `internal_user` or `internal_user_viewer` role""" - return self.value in [ - self.INTERNAL_USER, - self.INTERNAL_USER_VIEW_ONLY, - ] - - -class LitellmTableNames(str, enum.Enum): - """ - Enum for Table Names used by LiteLLM - """ - - TEAM_TABLE_NAME = "LiteLLM_TeamTable" - USER_TABLE_NAME = "LiteLLM_UserTable" - KEY_TABLE_NAME = "LiteLLM_VerificationToken" - PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" - MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable" - - -class Litellm_EntityType(enum.Enum): - """ - Enum for types of entities on litellm - - This enum allows specifying the type of entity that is being tracked in the database. - """ - - KEY = "key" - USER = "user" - END_USER = "end_user" - TEAM = "team" - TEAM_MEMBER = "team_member" - ORGANIZATION = "organization" - TAG = "tag" - - # global proxy level entity - PROXY = "proxy" - - -def hash_token(token: str): - import hashlib - - # Hash the string using SHA-256 - hashed_token = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - -class KeyManagementRoutes(str, enum.Enum): - """ - Enum for key management routes - """ - - # write routes - KEY_GENERATE = "/key/generate" - KEY_UPDATE = "/key/update" - KEY_DELETE = "/key/delete" - KEY_REGENERATE = "/key/regenerate" - KEY_GENERATE_SERVICE_ACCOUNT = "/key/service-account/generate" - KEY_REGENERATE_WITH_PATH_PARAM = "/key/{key_id}/regenerate" - KEY_BLOCK = "/key/block" - KEY_UNBLOCK = "/key/unblock" - KEY_BULK_UPDATE = "/key/bulk_update" - - # info and health routes - KEY_INFO = "/key/info" - KEY_HEALTH = "/key/health" - - # list routes - KEY_LIST = "/key/list" - - -class LiteLLMRoutes(enum.Enum): - openai_route_names = [ - "chat_completion", - "completion", - "embeddings", - "image_generation", - "video_generation", - "audio_transcriptions", - "moderations", - "model_list", # OpenAI /v1/models route - ] - openai_routes = [ - # chat completions - "/engines/{model}/chat/completions", - "/openai/deployments/{model}/chat/completions", - "/chat/completions", - "/v1/chat/completions", - "/cursor/chat/completions", - # completions - "/engines/{model}/completions", - "/openai/deployments/{model}/completions", - "/completions", - "/v1/completions", - # embeddings - "/engines/{model}/embeddings", - "/openai/deployments/{model}/embeddings", - "/embeddings", - "/v1/embeddings", - # image generation - "/images/generations", - "/v1/images/generations", - # image edit - "/images/edits", - "/v1/images/edits", - # video generation - "/videos", - "/v1/videos", - "/videos/{video_id}", - "/v1/videos/{video_id}", - "/videos/{video_id}/content", - "/v1/videos/{video_id}/content", - "/videos/{video_id}/remix", - "/v1/videos/{video_id}/remix", - # audio transcription - "/audio/transcriptions", - "/v1/audio/transcriptions", - # audio Speech - "/audio/speech", - "/v1/audio/speech", - # moderations - "/moderations", - "/v1/moderations", - # batches - "/v1/batches", - "/batches", - "/v1/batches/{batch_id}", - "/batches/{batch_id}", - "/v1/batches/{batch_id}/cancel", - "/batches/{batch_id}/cancel", - # files - "/v1/files", - "/files", - "/v1/files/{file_id}", - "/files/{file_id}", - "/v1/files/{file_id}/content", - "/files/{file_id}/content", - # fine_tuning - "/fine_tuning/jobs", - "/v1/fine_tuning/jobs", - "/fine_tuning/jobs/{fine_tuning_job_id}/cancel", - "/v1/fine_tuning/jobs/{fine_tuning_job_id}/cancel", - # assistants-related routes - "/assistants", - "/v1/assistants", - "/v1/assistants/{assistant_id}", - "/assistants/{assistant_id}", - "/threads", - "/v1/threads", - "/threads/{thread_id}", - "/v1/threads/{thread_id}", - "/threads/{thread_id}/messages", - "/v1/threads/{thread_id}/messages", - "/threads/{thread_id}/runs", - "/v1/threads/{thread_id}/runs", - # models - "/models", - "/v1/models", - # token counter - "/utils/token_counter", - "/utils/transform_request", - # rerank - "/rerank", - "/v1/rerank", - "/v2/rerank", - # realtime - "/realtime", - "/v1/realtime", - "/realtime?{model}", - "/v1/realtime?{model}", - # responses API - "/responses", - "/v1/responses", - "/responses/{response_id}", - "/v1/responses/{response_id}", - "/responses/{response_id}/input_items", - "/v1/responses/{response_id}/input_items", - "/responses/{response_id}/cancel", - "/v1/responses/{response_id}/cancel", - # vector stores - "/vector_stores", - "/v1/vector_stores", - "/vector_stores/{vector_store_id}/search", - "/v1/vector_stores/{vector_store_id}/search", - "/vector_stores/{vector_store_id}/files", - "/v1/vector_stores/{vector_store_id}/files", - "/vector_stores/{vector_store_id}/files/{file_id}", - "/v1/vector_stores/{vector_store_id}/files/{file_id}", - "/vector_stores/{vector_store_id}/files/{file_id}/content", - "/v1/vector_stores/{vector_store_id}/files/{file_id}/content", - "/vector_store/list", - "/v1/vector_store/list", - - # search - "/search", - "/v1/search", - "/search/{search_tool_name}", - "/v1/search/{search_tool_name}", - # OCR - "/ocr", - "/v1/ocr", - # containers API - "/containers", - "/v1/containers", - "/containers/*", - "/v1/containers/*", - ] - - mapped_pass_through_routes = [ - "/bedrock", - "/vertex-ai", - "/vertex_ai", - "/cohere", - "/gemini", - "/anthropic", - "/langfuse", - "/azure", - "/azure_ai", - "/openai", - "/openai_passthrough", - "/assemblyai", - "/eu.assemblyai", - "/vllm", - "/mistral", - "/milvus", - ] - - ######################################################### - # e.g /vllm/*, anthropic/*, etc. - # allows using /anthropic/v1/messages, /vllm/v1/chat/completions, etc. - ######################################################### - passthrough_routes_wildcard = [f"{route}/*" for route in mapped_pass_through_routes] - - litellm_native_routes = [ - "/rag/ingest", - "/v1/rag/ingest", - "/rag/query", - "/v1/rag/query", - ] - - anthropic_routes = [ - "/v1/messages", - "/v1/messages/count_tokens", - "/v1/skills", - "/v1/skills/{skill_id}", - ] - - mcp_routes = [ - "/mcp", - "/mcp/", - "/mcp/{subpath}", - "/mcp/tools", - "/mcp/tools/list", - "/mcp/tools/call", - ] - - agent_routes = [ - "/v1/agents", - "/agents", - "/a2a/{agent_id}", - "/a2a/{agent_id}/message/send", - "/a2a/{agent_id}/message/stream", - "/a2a/{agent_id}/.well-known/agent-card.json", - ] - - google_routes = [ - "/v1beta/models/{model_name:path}:countTokens", - "/v1beta/models/{model_name:path}:generateContent", - "/v1beta/models/{model_name:path}:streamGenerateContent", - "/models/{model_name:path}:countTokens", - "/models/{model_name:path}:generateContent", - "/models/{model_name:path}:streamGenerateContent", - # Google Interactions API - "/interactions", - "/v1beta/interactions", - "/interactions/{interaction_id}", - "/v1beta/interactions/{interaction_id}", - "/interactions/{interaction_id}/cancel", - "/v1beta/interactions/{interaction_id}/cancel", - ] - - apply_guardrail_routes = [ - "/guardrails/apply_guardrail", - ] - - llm_api_routes = ( - openai_routes - + anthropic_routes - + google_routes - + mapped_pass_through_routes - + passthrough_routes_wildcard - + apply_guardrail_routes - + mcp_routes - + litellm_native_routes - + agent_routes - ) - info_routes = [ - "/key/info", - "/key/health", - "/team/info", - "/team/list", - "/v2/team/list", - "/organization/list", - "/team/available", - "/user/info", - "/model/info", - "/v1/model/info", - "/v2/model/info", - "/v2/key/info", - "/model_group/info", - "/health", - "/key/list", - "/user/filter/ui", - "/models", - "/v1/models", - ] - - # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend - master_key_only_routes = [ - "/global/spend/reset", - "/memory-usage-in-mem-cache", - "/memory-usage-in-mem-cache-items", - ] - - key_management_routes = [ - KeyManagementRoutes.KEY_GENERATE.value, - KeyManagementRoutes.KEY_UPDATE.value, - KeyManagementRoutes.KEY_DELETE.value, - KeyManagementRoutes.KEY_INFO.value, - KeyManagementRoutes.KEY_REGENERATE.value, - KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT.value, - KeyManagementRoutes.KEY_REGENERATE_WITH_PATH_PARAM.value, - KeyManagementRoutes.KEY_LIST.value, - KeyManagementRoutes.KEY_BLOCK.value, - KeyManagementRoutes.KEY_UNBLOCK.value, - KeyManagementRoutes.KEY_BULK_UPDATE.value, - ] - - management_routes = [ - # user - "/user/new", - "/user/update", - "/user/delete", - "/user/info", - "/user/list", - # team - "/team/new", - "/team/update", - "/team/delete", - "/team/list", - "/v2/team/list", - "/team/info", - "/team/block", - "/team/unblock", - "/team/available", - "/team/permissions_list", - "/team/permissions_update", - # model - "/model/new", - "/model/update", - "/model/delete", - "/model/info", - ] + key_management_routes - - spend_tracking_routes = [ - # spend - "/spend/keys", - "/spend/users", - "/spend/tags", - "/spend/calculate", - "/spend/logs", - "/cost/estimate", - ] - - global_spend_tracking_routes = [ - # global spend - "/global/spend/logs", - "/global/spend", - "/global/spend/keys", - "/global/spend/teams", - "/global/spend/end_users", - "/global/spend/models", - "/global/predict/spend/logs", - "/global/spend/report", - "/global/spend/provider", - "/global/spend/tags", - ] - - public_routes = set( - [ - "/routes", - "/", - "/health/liveliness", - "/health/liveness", - "/health/readiness", - "/test", - "/config/yaml", - "/metrics", - "/litellm/.well-known/litellm-ui-config", - "/.well-known/litellm-ui-config", - "/public/model_hub", - "/public/agent_hub", - "/public/mcp_hub", - "/public/litellm_model_cost_map", - ] - ) - - ui_routes = [ - "/sso", - "/sso/get/ui_settings", - "/get/ui_settings", - "/login", - "/key/info", - "/config", - "/spend", - "/model/info", - "/v2/model/info", - "/v2/key/info", - "/models", - "/v1/models", - "/global/spend", - "/global/spend/logs", - "/global/spend/keys", - "/global/spend/models", - "/global/spend/tags", - "/global/predict/spend/logs", - "/global/activity", - "/health/services", - ] + info_routes - - internal_user_routes = ( - [ - "/global/spend/tags", - "/global/spend/keys", - "/global/spend/models", - "/global/spend/provider", - "/global/spend/end_users", - "/global/activity", - "/global/activity/model", - "/v1/models/{model_id}", - "/models/{model_id}", - ] - + spend_tracking_routes - + key_management_routes - ) - - internal_user_view_only_routes = ( - spend_tracking_routes + global_spend_tracking_routes - ) - - self_managed_routes = [ - "/team/member_add", - "/team/member_delete", - "/team/member_update", - "/team/permissions_list", - "/team/permissions_update", - "/team/daily/activity", - "/model/new", - "/model/update", - "/model/delete", - "/user/daily/activity", - "/model/{model_id}/update", - "/prompt/list", - "/prompt/info", - ] # routes that manage their own allowed/disallowed logic - - ## Org Admin Routes ## - - # Routes only an Org Admin Can Access - org_admin_only_routes = [ - "/organization/info", - "/organization/delete", - "/organization/member_add", - "/organization/member_update", - ] - - # Routes accessible by Admin Viewer (read-only admin access) - admin_viewer_routes = [ - "/user/list", - "/user/available_users", - "/user/available_roles", - "/user/daily/activity", - "/team/daily/activity", - "/tag/daily/activity", - "/tag/list", - ] + info_routes - - # All routes accesible by an Org Admin - org_admin_allowed_routes = ( - org_admin_only_routes - + management_routes - + self_managed_routes - + admin_viewer_routes - ) - - -class LiteLLMPromptInjectionParams(LiteLLMPydanticObjectBase): - heuristics_check: bool = False - vector_db_check: bool = False - llm_api_check: bool = False - llm_api_name: Optional[str] = None - llm_api_system_prompt: Optional[str] = None - llm_api_fail_call_string: Optional[str] = None - reject_as_response: Optional[bool] = Field( - default=False, - description="Return rejected request error message as a string to the user. Default behaviour is to raise an exception.", - ) - - @model_validator(mode="before") - @classmethod - def check_llm_api_params(cls, values): - llm_api_check = values.get("llm_api_check") - if llm_api_check is True: - if "llm_api_name" not in values or not values["llm_api_name"]: - raise ValueError( - "If llm_api_check is set to True, llm_api_name must be provided" - ) - if ( - "llm_api_system_prompt" not in values - or not values["llm_api_system_prompt"] - ): - raise ValueError( - "If llm_api_check is set to True, llm_api_system_prompt must be provided" - ) - if ( - "llm_api_fail_call_string" not in values - or not values["llm_api_fail_call_string"] - ): - raise ValueError( - "If llm_api_check is set to True, llm_api_fail_call_string must be provided" - ) - return values - - -######### Request Class Definition ###### -class ProxyChatCompletionRequest(LiteLLMPydanticObjectBase): - """ - Pydantic model for chat completion requests that includes both OpenAI standard fields - and LiteLLM-specific parameters. This replaces the previous TypedDict version. - """ - - # Required fields (from ChatCompletionRequest) - model: str - messages: List[AllMessageValues] - - # Standard OpenAI completion parameters (all optional) - frequency_penalty: Optional[float] = None - logit_bias: Optional[Dict[str, float]] = None - logprobs: Optional[bool] = None - top_logprobs: Optional[int] = None - max_tokens: Optional[int] = None - n: Optional[int] = None - presence_penalty: Optional[float] = None - response_format: Optional[Dict[str, Any]] = None - seed: Optional[int] = None - service_tier: Optional[str] = None - stop: Optional[Union[str, List[str]]] = None - stream_options: Optional[Dict[str, Any]] = None - temperature: Optional[float] = None - top_p: Optional[float] = None - tools: Optional[List[Dict[str, Any]]] = None - tool_choice: Optional[Union[str, Dict[str, Any]]] = None - parallel_tool_calls: Optional[bool] = None - function_call: Optional[Union[str, Dict[str, Any]]] = None - functions: Optional[List[Dict[str, Any]]] = None - user: Optional[str] = None - stream: Optional[bool] = None - - # LiteLLM-specific metadata param (from original ChatCompletionRequest) - metadata: Optional[Dict[str, Any]] = None - - # Optional LiteLLM params - guardrails: Optional[List[str]] = None - caching: Optional[bool] = None - num_retries: Optional[int] = None - context_window_fallback_dict: Optional[Dict[str, str]] = None - fallbacks: Optional[List[str]] = None - - -class ModelInfoDelete(LiteLLMPydanticObjectBase): - id: str - - -class ModelInfo(LiteLLMPydanticObjectBase): - id: Optional[str] - mode: Optional[Literal["embedding", "chat", "completion"]] - input_cost_per_token: Optional[float] = 0.0 - output_cost_per_token: Optional[float] = 0.0 - max_tokens: Optional[int] = 2048 # assume 2048 if not set - - # for azure models we need users to specify the base model, one azure you can call deployments - azure/my-random-model - # we look up the base model in model_prices_and_context_window.json - base_model: Optional[ - Literal[ - "gpt-4-1106-preview", - "gpt-4-32k", - "gpt-4", - "gpt-3.5-turbo-16k", - "gpt-3.5-turbo", - "text-embedding-ada-002", - ] - ] - - model_config = ConfigDict(protected_namespaces=(), extra="allow") - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("id") is None: - values.update({"id": str(uuid.uuid4())}) - if values.get("mode") is None: - values.update({"mode": None}) - if values.get("input_cost_per_token") is None: - values.update({"input_cost_per_token": None}) - if values.get("output_cost_per_token") is None: - values.update({"output_cost_per_token": None}) - if values.get("max_tokens") is None: - values.update({"max_tokens": None}) - if values.get("base_model") is None: - values.update({"base_model": None}) - return values - - -class ProviderInfo(LiteLLMPydanticObjectBase): - name: str - fields: List[ProviderField] - - -class BlockUsers(LiteLLMPydanticObjectBase): - user_ids: List[str] # required - - -class ModelParams(LiteLLMPydanticObjectBase): - model_name: str - litellm_params: dict - model_info: ModelInfo - - model_config = ConfigDict(protected_namespaces=()) - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("model_info") is None: - values.update( - {"model_info": ModelInfo(id=None, mode="chat", base_model=None)} - ) - return values - - -class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): - mcp_servers: Optional[List[str]] = None - mcp_access_groups: Optional[List[str]] = None - mcp_tool_permissions: Optional[Dict[str, List[str]]] = None - vector_stores: Optional[List[str]] = None - agents: Optional[List[str]] = None - agent_access_groups: Optional[List[str]] = None - - -class GenerateRequestBase(LiteLLMPydanticObjectBase): - """ - Overlapping schema between key and user generate/update requests - """ - - key_alias: Optional[str] = None - duration: Optional[str] = None - models: Optional[list] = [] - spend: Optional[float] = 0 - max_budget: Optional[float] = None - user_id: Optional[str] = None - team_id: Optional[str] = None - max_parallel_requests: Optional[int] = None - metadata: Optional[dict] = {} - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - - budget_duration: Optional[str] = None - allowed_cache_controls: Optional[list] = [] - config: Optional[dict] = {} - permissions: Optional[dict] = {} - model_max_budget: Optional[ - dict - ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} - - model_config = ConfigDict(protected_namespaces=()) - model_rpm_limit: Optional[dict] = None - model_tpm_limit: Optional[dict] = None - guardrails: Optional[List[str]] = None - policies: Optional[List[str]] = None - prompts: Optional[List[str]] = None - blocked: Optional[bool] = None - aliases: Optional[dict] = {} - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - - @field_validator("max_budget", mode="before") - @classmethod - def check_max_budget(cls, v): - if v == "": - return None - return v - - -class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): - index_name: str - index_permissions: List[Literal["read", "write"]] - - -class KeyRequestBase(GenerateRequestBase): - key: Optional[str] = None - budget_id: Optional[str] = None - tags: Optional[List[str]] = None - enforced_params: Optional[List[str]] = None - allowed_routes: Optional[list] = [] - allowed_passthrough_routes: Optional[list] = None - allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None - rpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm - tpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm - router_settings: Optional[UpdateRouterConfig] = None - - -class LiteLLMKeyType(str, enum.Enum): - """ - Enum for key types that determine what routes a key can access - """ - - LLM_API = "llm_api" # Can call LLM API routes (chat/completions, embeddings, etc.) - MANAGEMENT = "management" # Can call management routes (user/team/key management) - READ_ONLY = "read_only" # Can only call info/read routes - DEFAULT = "default" # Uses default allowed routes - - -class GenerateKeyRequest(KeyRequestBase): - soft_budget: Optional[float] = None - send_invite_email: Optional[bool] = None - key_type: Optional[LiteLLMKeyType] = Field( - default=LiteLLMKeyType.DEFAULT, - description="Type of key that determines default allowed routes.", - ) - auto_rotate: Optional[bool] = Field( - default=False, description="Whether this key should be automatically rotated" - ) - rotation_interval: Optional[str] = Field( - default=None, - description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True", - ) - organization_id: Optional[str] = None - - -class GenerateKeyResponse(KeyRequestBase): - key: str # type: ignore - key_name: Optional[str] = None - expires: Optional[datetime] = None - user_id: Optional[str] = None - token_id: Optional[str] = None - organization_id: Optional[str] = None - litellm_budget_table: Optional[Any] = None - token: Optional[str] = None - created_by: Optional[str] = None - updated_by: Optional[str] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("token") is not None: - values.update({"key": values.get("token")}) - dict_fields = [ - "metadata", - "aliases", - "config", - "permissions", - "model_max_budget", - "router_settings", - ] - for field in dict_fields: - value = values.get(field) - if value is not None and isinstance(value, str): - try: - values[field] = json.loads(value) - except json.JSONDecodeError: - raise ValueError(f"Field {field} should be a valid dictionary") - - return values - - -class UpdateKeyRequest(KeyRequestBase): - # Note: the defaults of all Params here MUST BE NONE - # else they will get overwritten - key: str # type: ignore - duration: Optional[str] = None - spend: Optional[float] = None - metadata: Optional[dict] = None - temp_budget_increase: Optional[float] = None - temp_budget_expiry: Optional[datetime] = None - auto_rotate: Optional[bool] = None - rotation_interval: Optional[str] = None - - @model_validator(mode="after") - def validate_temp_budget(self) -> "UpdateKeyRequest": - if self.temp_budget_increase is not None or self.temp_budget_expiry is not None: - if self.temp_budget_increase is None or self.temp_budget_expiry is None: - raise ValueError( - "temp_budget_increase and temp_budget_expiry must be set together" - ) - return self - - -class RegenerateKeyRequest(GenerateKeyRequest): - # This needs to be different from UpdateKeyRequest, because "key" is optional for this - key: Optional[str] = None - new_key: Optional[str] = None - duration: Optional[str] = None - spend: Optional[float] = None - metadata: Optional[dict] = None - new_master_key: Optional[str] = None - - -class KeyRequest(LiteLLMPydanticObjectBase): - keys: Optional[List[str]] = None - key_aliases: Optional[List[str]] = None - - @model_validator(mode="before") - @classmethod - def validate_at_least_one(cls, values): - if not values.get("keys") and not values.get("key_aliases"): - raise ValueError( - "At least one of 'keys' or 'key_aliases' must be provided." - ) - return values - - -class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): - id: Optional[int] = None - model_aliases: Optional[Union[str, dict]] = None # json dump the dict - created_by: str - updated_by: str - team: Optional["LiteLLM_TeamTable"] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): - model_id: str - model_name: str - litellm_params: dict - model_info: dict - created_at: Optional[datetime] = None - created_by: str - updated_at: Optional[datetime] = None - updated_by: str - - @model_validator(mode="before") - @classmethod - def check_potential_json_str(cls, values): - if isinstance(values.get("litellm_params"), str): - try: - values["litellm_params"] = json.loads(values["litellm_params"]) - except json.JSONDecodeError: - pass - if isinstance(values.get("model_info"), str): - try: - values["model_info"] = json.loads(values["model_info"]) - except json.JSONDecodeError: - pass - return values - - -# MCP Types -class SpecialMCPServerName(str, enum.Enum): - all_team_servers = "all-team-mcpservers" - all_proxy_servers = "all-proxy-mcpservers" - - -# MCP Proxy Request Types -class NewMCPServerRequest(LiteLLMPydanticObjectBase): - server_id: Optional[str] = None - server_name: Optional[str] = None - alias: Optional[str] = None - description: Optional[str] = None - transport: MCPTransportType = MCPTransport.sse - auth_type: Optional[MCPAuthType] = None - credentials: Optional[MCPCredentials] = None - url: Optional[str] = None - mcp_info: Optional[MCPInfo] = None - mcp_access_groups: List[str] = Field(default_factory=list) - allowed_tools: Optional[List[str]] = None - extra_headers: Optional[List[str]] = None - static_headers: Optional[Dict[str, str]] = None - # Stdio-specific fields - command: Optional[str] = None - args: List[str] = Field(default_factory=list) - env: Dict[str, str] = Field(default_factory=dict) - authorization_url: Optional[str] = None - token_url: Optional[str] = None - registration_url: Optional[str] = None - allow_all_keys: bool = False - - @model_validator(mode="before") - @classmethod - def validate_transport_fields(cls, values): - if isinstance(values, dict): - transport = values.get("transport") - if transport == MCPTransport.stdio: - if not values.get("command"): - raise ValueError("command is required for stdio transport") - if not values.get("args"): - raise ValueError("args is required for stdio transport") - elif transport in [MCPTransport.http, MCPTransport.sse]: - if not values.get("url"): - raise ValueError("url is required for HTTP/SSE transport") - return values - - @model_validator(mode="before") - @classmethod - def validate_credentials_requirements(cls, values): - if not isinstance(values, dict): - return values - - auth_type = values.get("auth_type") - if auth_type in {MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic}: - credentials = values.get("credentials") - auth_value = None - if isinstance(credentials, dict): - auth_value = credentials.get("auth_value") - elif hasattr(credentials, "get"): - auth_value = credentials.get("auth_value") # type: ignore[attr-defined] - - if not auth_value: - raise ValueError( - "auth_value is required when auth_type is api_key, bearer_token, or basic" - ) - - return values - - -class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): - server_id: str - server_name: Optional[str] = None - alias: Optional[str] = None - description: Optional[str] = None - transport: MCPTransportType = MCPTransport.sse - auth_type: Optional[MCPAuthType] = None - credentials: Optional[MCPCredentials] = None - url: Optional[str] = None - mcp_info: Optional[MCPInfo] = None - mcp_access_groups: List[str] = Field(default_factory=list) - allowed_tools: Optional[List[str]] = None - extra_headers: Optional[List[str]] = None - static_headers: Optional[Dict[str, str]] = None - # Stdio-specific fields - command: Optional[str] = None - args: List[str] = Field(default_factory=list) - env: Dict[str, str] = Field(default_factory=dict) - authorization_url: Optional[str] = None - token_url: Optional[str] = None - registration_url: Optional[str] = None - allow_all_keys: bool = False - - @model_validator(mode="before") - @classmethod - def validate_transport_fields(cls, values): - if isinstance(values, dict): - transport = values.get("transport") - if transport == MCPTransport.stdio: - if not values.get("command"): - raise ValueError("command is required for stdio transport") - if not values.get("args"): - raise ValueError("args is required for stdio transport") - elif transport in [MCPTransport.http, MCPTransport.sse]: - if not values.get("url"): - raise ValueError("url is required for HTTP/SSE transport") - return values - - -class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_MCPServerTable record""" - - server_id: str - server_name: Optional[str] = None - alias: Optional[str] = None - description: Optional[str] = None - url: Optional[str] = None - transport: MCPTransportType - auth_type: Optional[MCPAuthType] = None - credentials: Optional[MCPCredentials] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) - mcp_access_groups: List[str] = Field(default_factory=list) - allowed_tools: List[str] = Field(default_factory=list) - extra_headers: List[str] = Field(default_factory=list) - mcp_info: Optional[MCPInfo] = None - static_headers: Optional[Dict[str, str]] = None - # Health check status - status: Optional[Literal["healthy", "unhealthy", "unknown"]] = Field( - default="unknown", - description="Health status: 'healthy', 'unhealthy', 'unknown'", - ) - last_health_check: Optional[datetime] = None - health_check_error: Optional[str] = None - # Stdio-specific fields - command: Optional[str] = None - args: List[str] = Field(default_factory=list) - env: Dict[str, str] = Field(default_factory=dict) - authorization_url: Optional[str] = None - token_url: Optional[str] = None - registration_url: Optional[str] = None - allow_all_keys: bool = False - - -class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): - mcp_server_ids: List[str] - - -######## Skills API Types ######## - - -class NewSkillRequest(LiteLLMPydanticObjectBase): - """Request to create a new skill in LiteLLM database""" - - display_title: Optional[str] = None - description: Optional[str] = None - instructions: Optional[str] = None - file_content: Optional[bytes] = None # Binary content of skill files (zip) - file_name: Optional[str] = None # Original filename - file_type: Optional[str] = None # MIME type (e.g., "application/zip") - metadata: Optional[Dict[str, Any]] = None - authorization_url: Optional[str] = None - token_url: Optional[str] = None - registration_url: Optional[str] = None - - -class UpdateSkillRequest(LiteLLMPydanticObjectBase): - """Request to update an existing skill""" - - skill_id: str - display_title: Optional[str] = None - description: Optional[str] = None - instructions: Optional[str] = None - file_content: Optional[bytes] = None # Binary content of skill files (zip) - file_name: Optional[str] = None # Original filename - file_type: Optional[str] = None # MIME type - metadata: Optional[Dict[str, Any]] = None - - -class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_SkillsTable record""" - - skill_id: str - display_title: Optional[str] = None - description: Optional[str] = None - instructions: Optional[str] = None - source: str = "custom" - latest_version: Optional[str] = None - file_content: Optional[bytes] = None # Binary content of skill files (zip) - file_name: Optional[str] = None # Original filename - file_type: Optional[str] = None # MIME type - metadata: Optional[Dict[str, Any]] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - - -class ListSkillsRequest(LiteLLMPydanticObjectBase): - """Request to list skills from LiteLLM database""" - - limit: Optional[int] = 20 - offset: Optional[int] = 0 - - -class NewUserRequestTeam(LiteLLMPydanticObjectBase): - team_id: str - max_budget_in_team: Optional[float] = None - user_role: Literal["user", "admin"] = "user" - - -class NewUserRequest(GenerateRequestBase): - max_budget: Optional[float] = None - user_email: Optional[str] = None - user_alias: Optional[str] = None - user_role: Optional[ - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - ] = None - teams: Optional[Union[List[str], List[NewUserRequestTeam]]] = None - auto_create_key: bool = ( - True # flag used for returning a key as part of the /user/new response - ) - send_invite_email: Optional[bool] = None - sso_user_id: Optional[str] = None - organizations: Optional[List[str]] = None - - -class NewUserResponse(GenerateKeyResponse): - max_budget: Optional[float] = None - user_email: Optional[str] = None - user_role: Optional[ - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - ] = None - teams: Optional[list] = None - user_alias: Optional[str] = None - model_max_budget: Optional[dict] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - - -class UpdateUserRequestNoUserIDorEmail( - GenerateRequestBase -): # shared with BulkUpdateUserRequest - password: Optional[str] = None - spend: Optional[float] = None - metadata: Optional[dict] = None - user_alias: Optional[str] = None - user_role: Optional[ - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - ] = None - max_budget: Optional[float] = None - - -class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): - # Note: the defaults of all Params here MUST BE NONE - # else they will get overwritten - user_id: Optional[str] = None - user_email: Optional[str] = None - - @model_validator(mode="before") - @classmethod - def check_user_info(cls, values): - if values.get("user_id") is None and values.get("user_email") is None: - raise ValueError("Either user id or user email must be provided") - return values - - -class DeleteUserRequest(LiteLLMPydanticObjectBase): - user_ids: List[str] # required - - -AllowedModelRegion = Literal["eu", "us"] - - -class BudgetNewRequest(LiteLLMPydanticObjectBase): - budget_id: Optional[str] = Field(default=None, description="The unique budget id.") - max_budget: Optional[float] = Field( - default=None, - description="Requests will fail if this budget (in USD) is exceeded.", - ) - soft_budget: Optional[float] = Field( - default=None, - description="Requests will NOT fail if this is exceeded. Will fire alerting though.", - ) - max_parallel_requests: Optional[int] = Field( - default=None, description="Max concurrent requests allowed for this budget id." - ) - tpm_limit: Optional[int] = Field( - default=None, description="Max tokens per minute, allowed for this budget id." - ) - rpm_limit: Optional[int] = Field( - default=None, description="Max requests per minute, allowed for this budget id." - ) - budget_duration: Optional[str] = Field( - default=None, - description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')", - ) - model_max_budget: Optional[GenericBudgetConfigType] = Field( - default=None, - description="Max budget for each model (e.g. {'gpt-4o': {'max_budget': '0.0000001', 'budget_duration': '1d', 'tpm_limit': 1000, 'rpm_limit': 1000}})", - ) - budget_reset_at: Optional[datetime] = Field( - default=None, - description="Datetime when the budget is reset", - ) - - -class BudgetRequest(LiteLLMPydanticObjectBase): - budgets: List[str] - - -class BudgetDeleteRequest(LiteLLMPydanticObjectBase): - id: str - - -class CustomerBase(LiteLLMPydanticObjectBase): - user_id: str - alias: Optional[str] = None - spend: float = 0.0 - allowed_model_region: Optional[AllowedModelRegion] = None - default_model: Optional[str] = None - budget_id: Optional[str] = None - litellm_budget_table: Optional[BudgetNewRequest] = None - blocked: bool = False - - -class NewCustomerRequest(BudgetNewRequest): - """ - Create a new customer, allocate a budget to them - """ - - user_id: str - alias: Optional[str] = None # human-friendly alias - blocked: bool = False # allow/disallow requests for this end-user - budget_id: Optional[str] = None # give either a budget_id or max_budget - spend: Optional[float] = None - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model - - @model_validator(mode="before") - @classmethod - def check_user_info(cls, values): - if values.get("max_budget") is not None and values.get("budget_id") is not None: - raise ValueError("Set either 'max_budget' or 'budget_id', not both.") - - return values - - -class UpdateCustomerRequest(LiteLLMPydanticObjectBase): - """ - Update a Customer, use this to update customer budgets etc - - """ - - user_id: str - alias: Optional[str] = None # human-friendly alias - blocked: bool = False # allow/disallow requests for this end-user - max_budget: Optional[float] = None - budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model - - -class DeleteCustomerRequest(LiteLLMPydanticObjectBase): - """ - Delete multiple Customers - """ - - user_ids: List[str] - - -class MemberBase(LiteLLMPydanticObjectBase): - user_id: Optional[str] = Field( - default=None, - description="The unique ID of the user to add. Either user_id or user_email must be provided", - ) - user_email: Optional[str] = Field( - default=None, - description="The email address of the user to add. Either user_id or user_email must be provided", - ) - - @model_validator(mode="before") - @classmethod - def check_user_info(cls, values): - if not isinstance(values, dict): - raise ValueError("input needs to be a dictionary") - if values.get("user_id") is None and values.get("user_email") is None: - raise ValueError("Either user id or user email must be provided") - return values - - -class Member(MemberBase): - role: Literal[ - "admin", - "user", - ] = Field( - description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" - ) - - -class OrgMember(MemberBase): - role: Literal[ - LitellmUserRoles.ORG_ADMIN, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - - -class TeamBase(LiteLLMPydanticObjectBase): - team_alias: Optional[str] = None - team_id: Optional[str] = None - organization_id: Optional[str] = None - admins: list = [] - members: list = [] - members_with_roles: List[Member] = [] - team_member_permissions: Optional[List[str]] = None - metadata: Optional[dict] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - - # Budget fields - max_budget: Optional[float] = None - budget_duration: Optional[str] = None - - models: list = [] - blocked: bool = False - router_settings: Optional[dict] = None - - -class NewTeamRequest(TeamBase): - model_aliases: Optional[dict] = None - tags: Optional[list] = None - guardrails: Optional[List[str]] = None - policies: Optional[List[str]] = None - prompts: Optional[List[str]] = None - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - allowed_passthrough_routes: Optional[list] = None - secret_manager_settings: Optional[dict] = None - model_rpm_limit: Optional[Dict[str, int]] = None - rpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm - tpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm - - model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[ - float - ] = None # allow user to set a budget for all team members - team_member_rpm_limit: Optional[ - int - ] = None # allow user to set RPM limit for all team members - team_member_tpm_limit: Optional[ - int - ] = None # allow user to set TPM limit for all team members - team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" - allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class GlobalEndUsersSpend(LiteLLMPydanticObjectBase): - api_key: Optional[str] = None - startTime: Optional[datetime] = None - endTime: Optional[datetime] = None - - -class UpdateTeamRequest(LiteLLMPydanticObjectBase): - """ - UpdateTeamRequest, used by /team/update when you need to update a team - - team_id: str - team_alias: Optional[str] = None - organization_id: Optional[str] = None - metadata: Optional[dict] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - max_budget: Optional[float] = None - models: Optional[list] = None - blocked: Optional[bool] = None - budget_duration: Optional[str] = None - guardrails: Optional[List[str]] = None - policies: Optional[List[str]] = None - """ - - team_id: str # required - team_alias: Optional[str] = None - organization_id: Optional[str] = None - metadata: Optional[dict] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - max_budget: Optional[float] = None - models: Optional[list] = None - blocked: Optional[bool] = None - budget_duration: Optional[str] = None - tags: Optional[list] = None - model_aliases: Optional[dict] = None - guardrails: Optional[List[str]] = None - policies: Optional[List[str]] = None - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - team_member_budget: Optional[float] = None - team_member_budget_duration: Optional[str] = None - team_member_rpm_limit: Optional[int] = None - team_member_tpm_limit: Optional[int] = None - team_member_key_duration: Optional[str] = None - allowed_passthrough_routes: Optional[list] = None - secret_manager_settings: Optional[dict] = None - prompts: Optional[List[str]] = None - model_rpm_limit: Optional[Dict[str, int]] = None - model_tpm_limit: Optional[Dict[str, int]] = None - allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None - router_settings: Optional[dict] = None - - -class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): - """ - internal type used to reset the budget on a team - used by reset_budget() - - team_id: str - spend: float - budget_reset_at: datetime - """ - - team_id: str - spend: float - budget_reset_at: datetime - updated_at: datetime - - -class DeleteTeamRequest(LiteLLMPydanticObjectBase): - team_ids: List[str] # required - - -class BlockTeamRequest(LiteLLMPydanticObjectBase): - team_id: str # required - - -class BlockKeyRequest(LiteLLMPydanticObjectBase): - key: str # required - - -class AddTeamCallback(LiteLLMPydanticObjectBase): - callback_name: str - callback_type: Optional[ - Literal["success", "failure", "success_and_failure"] - ] = "success_and_failure" - callback_vars: Dict[str, str] - - @model_validator(mode="before") - @classmethod - def validate_callback_vars(cls, values): - callback_vars = values.get("callback_vars", {}) - valid_keys = set(StandardCallbackDynamicParams.__annotations__.keys()) - for key, value in callback_vars.items(): - if key not in valid_keys: - raise ValueError( - f"Invalid callback variable: {key}. Must be one of {valid_keys}" - ) - if not isinstance(value, str): - callback_vars[key] = str(value) - return values - - -class TeamCallbackMetadata(LiteLLMPydanticObjectBase): - success_callback: Optional[List[str]] = [] - failure_callback: Optional[List[str]] = [] - callbacks: Optional[List[str]] = [] - # for now - only supported for langfuse - callback_vars: Optional[Dict[str, str]] = {} - - @model_validator(mode="before") - @classmethod - def validate_callback_vars(cls, values): - success_callback = values.get("success_callback", []) - if success_callback is None: - values.pop("success_callback", None) - failure_callback = values.get("failure_callback", []) - if failure_callback is None: - values.pop("failure_callback", None) - callbacks = values.get("callbacks", []) - if callbacks is None: - values.pop("callbacks", None) - - callback_vars = values.get("callback_vars", {}) - if callback_vars is None: - values.pop("callback_vars", None) - if all(val is None for val in values.values()): - return { - "success_callback": [], - "failure_callback": [], - "callbacks": [], - "callback_vars": {}, - } - valid_keys = set(StandardCallbackDynamicParams.__annotations__.keys()) - if callback_vars is not None: - for key in callback_vars: - if key not in valid_keys: - raise ValueError( - f"Invalid callback variable: {key}. Must be one of {valid_keys}" - ) - return values - - -class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_ObjectPermissionTable record""" - - object_permission_id: str - mcp_servers: Optional[List[str]] = [] - mcp_access_groups: Optional[List[str]] = [] - mcp_tool_permissions: Optional[Dict[str, List[str]]] = None - """ - Mapping - server_id -> list of tools - - Enforces allowed tools for a specific key/team/organization - { - "1234567890": ["tool_name_1", "tool_name_2"] - } - """ - - vector_stores: Optional[List[str]] = [] - agents: Optional[List[str]] = [] - agent_access_groups: Optional[List[str]] = [] - - -class LiteLLM_TeamTable(TeamBase): - team_id: str # type: ignore - spend: Optional[float] = None - max_parallel_requests: Optional[int] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - model_id: Optional[int] = None - litellm_model_table: Optional[LiteLLM_ModelTable] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - updated_at: Optional[datetime] = None - created_at: Optional[datetime] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission_id: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - dict_fields = [ - "metadata", - "aliases", - "config", - "permissions", - "model_max_budget", - "model_aliases", - "router_settings", - ] - - if isinstance(values, BaseModel): - values = values.model_dump() - - if ( - isinstance(values.get("members_with_roles"), dict) - and not values["members_with_roles"] - ): - values["members_with_roles"] = [] - - for field in dict_fields: - value = values.get(field) - if value is not None and isinstance(value, str): - try: - values[field] = json.loads(value) - except json.JSONDecodeError: - raise ValueError(f"Field {field} should be a valid dictionary") - - return values - - -class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): - last_refreshed_at: Optional[float] = None - - -class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): - """ - Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class TeamRequest(LiteLLMPydanticObjectBase): - teams: List[str] - - -class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): - """Represents user-controllable params for a LiteLLM_BudgetTable record""" - - budget_id: Optional[str] = None - soft_budget: Optional[float] = None - max_budget: Optional[float] = None - max_parallel_requests: Optional[int] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - model_max_budget: Optional[dict] = None - budget_duration: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): - """Represents all params for a LiteLLM_BudgetTable record""" - - budget_reset_at: Optional[datetime] = None - created_at: datetime - - -class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): - """ - Used to track spend of a user_id within a team_id - """ - - spend: Optional[float] = None - user_id: Optional[str] = None - team_id: Optional[str] = None - budget_id: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class NewOrganizationRequest(LiteLLM_BudgetTable): - organization_id: Optional[str] = None - organization_alias: str - models: List = [] - budget_id: Optional[str] = None - metadata: Optional[dict] = None - model_rpm_limit: Optional[Dict[str, int]] = None - model_tpm_limit: Optional[Dict[str, int]] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - - -class OrganizationRequest(LiteLLMPydanticObjectBase): - organizations: List[str] - - -class DeleteOrganizationRequest(LiteLLMPydanticObjectBase): - organization_ids: List[str] # required - - -class TeamDefaultSettings(LiteLLMPydanticObjectBase): - team_id: str - - model_config = ConfigDict( - extra="allow" - ) # allow params not defined here, these fall in litellm.completion(**kwargs) - - -class DynamoDBArgs(LiteLLMPydanticObjectBase): - billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"] - read_capacity_units: Optional[int] = None - write_capacity_units: Optional[int] = None - ssl_verify: Optional[bool] = None - region_name: str - user_table_name: str = "LiteLLM_UserTable" - key_table_name: str = "LiteLLM_VerificationToken" - config_table_name: str = "LiteLLM_Config" - spend_table_name: str = "LiteLLM_SpendLogs" - aws_role_name: Optional[str] = None - aws_session_name: Optional[str] = None - aws_web_identity_token: Optional[str] = None - aws_provider_id: Optional[str] = None - aws_policy_arns: Optional[List[str]] = None - aws_policy: Optional[str] = None - aws_duration_seconds: Optional[int] = None - assume_role_aws_role_name: Optional[str] = None - assume_role_aws_session_name: Optional[str] = None - - -class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): - """ - Settings for a specific guardrail on a passthrough endpoint. - - Allows field-level targeting for guardrail execution. - """ - - request_fields: Optional[List[str]] = Field( - default=None, - description="JSONPath expressions for input field targeting (pre_call). Examples: 'query', 'documents[*].text', 'messages[*].content'. If not specified, guardrail runs on entire request payload.", - ) - response_fields: Optional[List[str]] = Field( - default=None, - description="JSONPath expressions for output field targeting (post_call). Examples: 'results[*].text', 'output'. If not specified, guardrail runs on entire response payload.", - ) - - -# Type alias for the guardrails dict: guardrail_name -> settings (or None for defaults) -PassThroughGuardrailsConfig = Dict[str, Optional[PassThroughGuardrailSettings]] - - -class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): - id: Optional[str] = Field( - default=None, - description="Optional unique identifier for the pass-through endpoint. If not provided, endpoints will be identified by path for backwards compatibility.", - ) - path: str = Field(description="The route to be added to the LiteLLM Proxy Server.") - target: str = Field( - description="The URL to which requests for this path should be forwarded." - ) - headers: dict = Field( - default={}, - description="Key-value pairs of headers to be forwarded with the request. You can set any key value pair here and it will be forwarded to your target endpoint", - ) - include_subpath: bool = Field( - default=False, - description="If True, requests to subpaths of the path will be forwarded to the target endpoint. For example, if the path is /bria and include_subpath is True, requests to /bria/v1/text-to-image/base/2.3 will be forwarded to the target endpoint.", - ) - cost_per_request: float = Field( - default=0.0, - description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.", - ) - auth: bool = Field( - default=False, - description="Whether authentication is required for the pass-through endpoint. If True, requests to the endpoint will require a valid LiteLLM API key.", - ) - guardrails: Optional[PassThroughGuardrailsConfig] = Field( - default=None, - description="Guardrails configuration for this passthrough endpoint. Dict keys are guardrail names, values are optional settings for field targeting. When set, all org/team/key level guardrails will also execute. Defaults to None (no guardrails execute).", - ) - - -class PassThroughEndpointResponse(LiteLLMPydanticObjectBase): - endpoints: List[PassThroughGenericEndpoint] - - -class ConfigFieldUpdate(LiteLLMPydanticObjectBase): - field_name: str - field_value: Any - config_type: Literal["general_settings"] - - -class ConfigFieldDelete(LiteLLMPydanticObjectBase): - config_type: Literal["general_settings"] - field_name: str - - -class CallbackDelete(LiteLLMPydanticObjectBase): - callback_name: str - - -class FieldDetail(BaseModel): - field_name: str - field_type: str - field_description: str - field_default_value: Any = None - stored_in_db: Optional[bool] - - -class ConfigList(LiteLLMPydanticObjectBase): - field_name: str - field_type: str - field_description: str - field_value: Any - stored_in_db: Optional[bool] - field_default_value: Any - premium_field: bool = False - nested_fields: Optional[ - List[FieldDetail] - ] = None # For nested dictionary or Pydantic fields - - -class UserHeaderMapping(LiteLLMPydanticObjectBase): - """ - Map an incoming HTTP header to a LiteLLM user role. - """ - - header_name: str - litellm_user_role: Literal[ - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.CUSTOMER, - ] - - model_config = { - "extra": "forbid", - } - - -UserMCPManagementMode = Literal["restricted", "view_all"] - - -class ConfigGeneralSettings(LiteLLMPydanticObjectBase): - """ - Documents all the fields supported by `general_settings` in config.yaml - """ - - completion_model: Optional[str] = Field( - None, description="proxy level default model for all chat completion calls" - ) - key_management_system: Optional[KeyManagementSystem] = Field( - None, description="key manager to load keys from / decrypt keys with" - ) - use_google_kms: Optional[bool] = Field( - None, description="decrypt keys with google kms" - ) - use_azure_key_vault: Optional[bool] = Field( - None, description="load keys from azure key vault" - ) - master_key: Optional[str] = Field( - None, description="require a key for all calls to proxy" - ) - database_url: Optional[str] = Field( - None, - description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", - ) - database_connection_pool_limit: Optional[int] = Field( - 10, - description="default connection pool for prisma client connecting to postgres db", - ) - database_connection_timeout: Optional[float] = Field( - 60, description="default timeout for a connection to the database" - ) - database_type: Optional[Literal["dynamo_db"]] = Field( - None, description="to use dynamodb instead of postgres db" - ) - database_args: Optional[DynamoDBArgs] = Field( - None, - description="custom args for instantiating dynamodb client - e.g. billing provision", - ) - otel: Optional[bool] = Field( - None, - description="[BETA] OpenTelemetry support - this might change, use with caution.", - ) - custom_auth: Optional[str] = Field( - None, - description="override user_api_key_auth with your own auth script - https://docs.litellm.ai/docs/proxy/virtual_keys#custom-auth", - ) - max_parallel_requests: Optional[int] = Field( - None, - description="maximum parallel requests for each api key", - ) - global_max_parallel_requests: Optional[int] = Field( - None, description="global max parallel requests to allow for a proxy instance." - ) - max_request_size_mb: Optional[int] = Field( - None, - description="max request size in MB, if a request is larger than this size it will be rejected", - ) - max_response_size_mb: Optional[int] = Field( - None, - description="max response size in MB, if a response is larger than this size it will be rejected", - ) - infer_model_from_keys: Optional[bool] = Field( - None, - description="for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY)", - ) - background_health_checks: Optional[bool] = Field( - None, description="run health checks in background" - ) - health_check_interval: int = Field( - 300, description="background health check interval in seconds" - ) - alerting: Optional[List] = Field( - None, - description="List of alerting integrations. Today, just slack - `alerting: ['slack']`", - ) - alert_types: Optional[List[AlertType]] = Field( - None, - description="List of alerting types. By default it is all alerts", - ) - alert_to_webhook_url: Optional[Dict] = Field( - None, - description="Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`", - ) - alerting_args: Optional[Dict] = Field( - None, description="Controllable params for slack alerting - e.g. ttl in cache." - ) - alerting_threshold: Optional[int] = Field( - None, - description="sends alerts if requests hang for 5min+", - ) - ui_access_mode: Optional[Literal["admin_only", "all"]] = Field( - "all", description="Control access to the Proxy UI" - ) - allowed_routes: Optional[List] = Field( - None, description="Proxy API Endpoints you want users to be able to access" - ) - reject_clientside_metadata_tags: Optional[bool] = Field( - None, - description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", - ) - enable_public_model_hub: bool = Field( - default=False, - description="Public model hub for users to see what models they have access to, supported openai params, etc.", - ) - pass_through_endpoints: Optional[List[PassThroughGenericEndpoint]] = Field( - default=None, - description="Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through", - ) - user_header_name: Optional[str] = Field( - None, - description="[DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings.", - ) - user_header_mappings: Optional[List[UserHeaderMapping]] = None - supported_db_objects: Optional[List[SupportedDBObjectType]] = Field( - None, - description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map'. If not set, all objects are loaded (default behavior).", - ) - user_mcp_management_mode: Optional[UserMCPManagementMode] = Field( - None, - description="Controls how non-admin users interact with MCP servers in the dashboard. 'restricted' shows only accessible servers, 'view_all' lists every server in read-only mode.", - ) - store_prompts_in_spend_logs: Optional[bool] = Field( - None, - description="If True, stores request messages and responses in spend logs. Default is False.", - ) - maximum_spend_logs_retention_period: Optional[str] = Field( - None, - description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", - ) - - -class ConfigYAML(LiteLLMPydanticObjectBase): - """ - Documents all the fields supported by the config.yaml - """ - - environment_variables: Optional[dict] = Field( - None, - description="Object to pass in additional environment variables via POST request", - ) - model_list: Optional[List[ModelParams]] = Field( - None, - description="List of supported models on the server, with model-specific configs", - ) - litellm_settings: Optional[dict] = Field( - None, - description="litellm Module settings. See __init__.py for all, example litellm.drop_params=True, litellm.set_verbose=True, litellm.api_base, litellm.cache", - ) - general_settings: Optional[ConfigGeneralSettings] = None - router_settings: Optional[UpdateRouterConfig] = Field( - None, - description="litellm router object settings. See router.py __init__ for all, example router.num_retries=5, router.timeout=5, router.max_retries=5, router.retry_after=5", - ) - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): - token: Optional[str] = None - key_name: Optional[str] = None - key_alias: Optional[str] = None - spend: float = 0.0 - max_budget: Optional[float] = None - expires: Optional[Union[str, datetime]] = None - models: List = [] - aliases: Dict = {} - config: Dict = {} - user_id: Optional[str] = None - team_id: Optional[str] = None - max_parallel_requests: Optional[int] = None - metadata: Dict = {} - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - allowed_cache_controls: Optional[list] = [] - allowed_routes: Optional[list] = [] - permissions: Dict = {} - model_spend: Dict = {} - model_max_budget: Dict = {} - soft_budget_cooldown: bool = False - blocked: Optional[bool] = None - litellm_budget_table: Optional[dict] = None - org_id: Optional[str] = None # org id for a given key - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - object_permission_id: Optional[str] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - rotation_count: Optional[int] = 0 # Number of times key has been rotated - auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated - rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") - last_rotation_at: Optional[datetime] = None # When this key was last rotated - key_rotation_at: Optional[datetime] = None # When this key should next be rotated - router_settings: Optional[ - Dict - ] = None # Router settings for this key (Key > Team > Global precedence) - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): - """ - Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): - """ - Combined view of litellm verification token + litellm team table (select values) - """ - - team_spend: Optional[float] = None - team_alias: Optional[str] = None - team_tpm_limit: Optional[int] = None - team_rpm_limit: Optional[int] = None - team_max_budget: Optional[float] = None - team_models: List = [] - team_blocked: bool = False - soft_budget: Optional[float] = None - team_model_aliases: Optional[Dict] = None - team_member: Optional[Member] = None - team_metadata: Optional[Dict] = None - team_object_permission_id: Optional[str] = None - - # Team Member Specific Params - team_member_spend: Optional[float] = None - team_member_tpm_limit: Optional[int] = None - team_member_rpm_limit: Optional[int] = None - - # End User Params - end_user_id: Optional[str] = None - end_user_tpm_limit: Optional[int] = None - end_user_rpm_limit: Optional[int] = None - end_user_max_budget: Optional[float] = None - - # Organization Params - organization_max_budget: Optional[float] = None - organization_tpm_limit: Optional[int] = None - organization_rpm_limit: Optional[int] = None - organization_metadata: Optional[dict] = None - - # Time stamps - last_refreshed_at: Optional[float] = None # last time joint view was pulled from db - - def __init__(self, **kwargs): - # Handle litellm_budget_table_* keys - for key, value in list(kwargs.items()): - if key.startswith("litellm_budget_table_") and value is not None: - # Extract the corresponding attribute name - attr_name = key.replace("litellm_budget_table_", "") - # Check if the value is None and set the corresponding attribute - if getattr(self, attr_name, None) is None: - kwargs[attr_name] = value - if key == "end_user_id" and value is not None and isinstance(value, int): - kwargs[key] = str(value) - - if kwargs.get("organization_id") is not None: - kwargs["org_id"] = kwargs.pop("organization_id") - # Initialize the superclass - super().__init__(**kwargs) - - -class UserAPIKeyAuth( - LiteLLM_VerificationTokenView -): # the expected response object for user api key auth - """ - Return the row in the db - """ - - api_key: Optional[str] = None - user_role: Optional[LitellmUserRoles] = None - allowed_model_region: Optional[AllowedModelRegion] = None - parent_otel_span: Optional[Span] = None - rpm_limit_per_model: Optional[Dict[str, int]] = None - tpm_limit_per_model: Optional[Dict[str, int]] = None - user_tpm_limit: Optional[int] = None - user_rpm_limit: Optional[int] = None - user_email: Optional[str] = None - user_spend: Optional[float] = None - user_max_budget: Optional[float] = None - request_route: Optional[str] = None - user: Optional[Any] = None # Expanded user object when expand=user is used - - model_config = ConfigDict(arbitrary_types_allowed=True) - - @model_validator(mode="before") - @classmethod - def check_api_key(cls, values): - # If values is already an instance (not a dict), return it as-is - if not isinstance(values, dict): - return values - if values.get("api_key") is not None: - values.update( - {"token": cls._safe_hash_litellm_api_key(values.get("api_key"))} - ) - if isinstance(values.get("api_key"), str): - values.update( - {"api_key": cls._safe_hash_litellm_api_key(values.get("api_key"))} - ) - return values - - @classmethod - def _safe_hash_litellm_api_key(cls, api_key: str) -> str: - """ - Helper to ensure all logged keys are hashed - Covers: - 1. Regular API keys from LiteLLM DB - 2. JWT tokens used for connecting to LiteLLM API - """ - if api_key.startswith("sk-"): - return hash_token(api_key) - from litellm.proxy.auth.handle_jwt import JWTHandler - - if JWTHandler.is_jwt(token=api_key): - return f"hashed-jwt-{hash_token(token=api_key)}" - return api_key - - @classmethod - def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth": - """ - Returns a `UserAPIKeyAuth` object for the litellm internal health check service account. - - This is used to track number of requests/spend for health check calls. - """ - from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME - - return cls( - api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, - team_id=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, - key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, - team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, - ) - - @classmethod - def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth": - """ - Returns a `UserAPIKeyAuth` object for the litellm internal health check service account. - - This is used to track number of requests/spend for health check calls. - """ - from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME - - return cls( - api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME, - team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME, - key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME, - team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME, - ) - - @classmethod - def get_litellm_internal_jobs_user_api_key_auth(cls) -> "UserAPIKeyAuth": - """ - Returns a `UserAPIKeyAuth` object for internal LiteLLM jobs like key rotation. - - This is used to track actions performed by automated system jobs. - """ - from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME - - return cls( - api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, - team_id="system", - key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, - team_alias="system", - user_id="system", - ) - - -class UserInfoResponse(LiteLLMPydanticObjectBase): - user_id: Optional[str] - user_info: Optional[Union[dict, BaseModel]] - keys: List - teams: List - - -class LiteLLM_Config(LiteLLMPydanticObjectBase): - param_name: str - param_value: Dict - - -class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): - """ - This is the table that track what organizations a user belongs to and users spend within the organization - """ - - user_id: str - organization_id: str - user_role: Optional[str] = None - spend: float = 0.0 - budget_id: Optional[str] = None - created_at: datetime - updated_at: datetime - user: Optional[ - Any - ] = None # You might want to replace 'Any' with a more specific type if available - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable): - """Represents user-controllable params for a LiteLLM_OrganizationTable record""" - - organization_id: Optional[str] = None - organization_alias: Optional[str] = None - budget_id: Optional[str] = None - spend: Optional[float] = None - metadata: Optional[dict] = None - models: Optional[List[str]] = None - updated_by: Optional[str] = None - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - model_tpm_limit: Optional[Dict[str, int]] = None - model_rpm_limit: Optional[Dict[str, int]] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - for field in LiteLLM_ManagementEndpoint_MetadataFields: - if values.get(field) is not None: - # add to metadata - if values.get("metadata") is None: - values.update({"metadata": {}}) - values["metadata"][field] = values.get(field) - values.pop(field) - return values - - -class LiteLLM_UserTable(LiteLLMPydanticObjectBase): - user_id: str - max_budget: Optional[float] = None - spend: float = 0.0 - model_max_budget: Optional[Dict] = {} - model_spend: Optional[Dict] = {} - user_email: Optional[str] = None - user_alias: Optional[str] = None - models: list = [] - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - user_role: Optional[str] = None - organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None - teams: List[str] = [] - sso_user_id: Optional[str] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - metadata: Optional[dict] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - if values.get("models") is None: - values.update({"models": []}) - if values.get("teams") is None: - values.update({"teams": []}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): - """Represents user-controllable params for a LiteLLM_OrganizationTable record""" - - organization_id: Optional[str] = None - organization_alias: Optional[str] = None - budget_id: str - spend: float = 0.0 - metadata: Optional[dict] = None - models: List[str] - created_by: str - updated_by: str - users: Optional[List[LiteLLM_UserTable]] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - object_permission_id: Optional[str] = None - - -class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable): - """Returned by the /organization/info endpoint and /organization/list endpoint""" - - members: List[LiteLLM_OrganizationMembershipTable] = [] - teams: List[LiteLLM_TeamTable] = [] - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - created_at: datetime - updated_at: datetime - - -class NewOrganizationResponse(LiteLLM_OrganizationTable): - organization_id: str # type: ignore - created_at: datetime - updated_at: datetime - - -class LiteLLM_UserTableFiltered(BaseModel): # done to avoid exposing sensitive data - user_id: str - user_email: Optional[str] = None - - -class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable): - key_count: int = 0 - - -class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): - user_id: str - blocked: bool - alias: Optional[str] = None - spend: float = 0.0 - allowed_model_region: Optional[AllowedModelRegion] = None - default_model: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_TagTable(LiteLLMPydanticObjectBase): - tag_name: str - description: Optional[str] = None - models: List[str] = [] - model_info: Optional[dict] = None - spend: float = 0.0 - budget_id: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - if values.get("models") is None: - values.update({"models": []}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_SpendLogs(LiteLLMPydanticObjectBase): - request_id: str - api_key: str - model: Optional[str] = "" - api_base: Optional[str] = "" - call_type: str - spend: Optional[float] = 0.0 - total_tokens: Optional[int] = 0 - prompt_tokens: Optional[int] = 0 - completion_tokens: Optional[int] = 0 - startTime: Union[str, datetime, None] - endTime: Union[str, datetime, None] - user: Optional[str] = "" - metadata: Optional[Json] = {} - cache_hit: Optional[str] = "False" - cache_key: Optional[str] = None - request_tags: Optional[Json] = None - requester_ip_address: Optional[str] = None - messages: Optional[Union[str, list, dict]] - response: Optional[Union[str, list, dict]] - - -class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): - request_id: Optional[str] = str(uuid.uuid4()) - api_base: Optional[str] = "" - model_group: Optional[str] = "" - litellm_model_name: Optional[str] = "" - model_id: Optional[str] = "" - request_kwargs: Optional[dict] = {} - exception_type: Optional[str] = "" - status_code: Optional[str] = "" - exception_string: Optional[str] = "" - startTime: Union[str, datetime, None] - endTime: Union[str, datetime, None] - - -AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "rotated"] - - -class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): - id: str - updated_at: datetime - changed_by: Optional[Any] = None - changed_by_api_key: Optional[str] = None - action: AUDIT_ACTIONS - table_name: LitellmTableNames - object_id: str - before_value: Optional[Json] = None - updated_values: Optional[Json] = None - - @model_validator(mode="before") - @classmethod - def cast_changed_by_to_str(cls, values): - if values.get("changed_by") is not None: - values["changed_by"] = str(values["changed_by"]) - return values - - @model_validator(mode="after") - def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker - - masker = SensitiveDataMasker(sensitive_patterns={"key"}) - - if self.before_value is not None: - json_before_value: Optional[dict] = None - if isinstance(self.before_value, str): - json_before_value = json.loads(self.before_value) - elif isinstance(self.before_value, dict): - json_before_value = self.before_value - - if json_before_value is not None: - json_before_value = masker.mask_dict(json_before_value) - self.before_value = json.dumps(json_before_value, default=str) - - if self.updated_values is not None: - json_updated_values: Optional[dict] = None - if isinstance(self.updated_values, str): - json_updated_values = json.loads(self.updated_values) - elif isinstance(self.updated_values, dict): - json_updated_values = self.updated_values - - if json_updated_values is not None: - json_updated_values = masker.mask_dict(json_updated_values) - self.updated_values = json.dumps(json_updated_values, default=str) - - return self - - -class LiteLLM_SpendLogs_ResponseObject(LiteLLMPydanticObjectBase): - response: Optional[List[Union[LiteLLM_SpendLogs, Any]]] = None - - -class TokenCountRequest(LiteLLMPydanticObjectBase): - model: str - prompt: Optional[str] = None - messages: Optional[List[dict]] = None - """ - Anthropic token counting endpoint uses /messages - """ - - contents: Optional[List[dict]] = None - """ - Google /countTokens endpoint expects contents to be a list of dicts with the following structure: - """ - - -class CallInfo(LiteLLMPydanticObjectBase): - """Used for slack budget alerting""" - - spend: float - max_budget: Optional[float] = None - soft_budget: Optional[float] = None - token: Optional[str] = Field(default=None, description="Hashed value of that key") - customer_id: Optional[str] = None - user_id: Optional[str] = None - team_id: Optional[str] = None - team_alias: Optional[str] = None - organization_id: Optional[str] = None - user_email: Optional[str] = None - key_alias: Optional[str] = None - projected_exceeded_date: Optional[str] = None - projected_spend: Optional[float] = None - event_group: Litellm_EntityType - - -class WebhookEvent(CallInfo): - event: Literal[ - "budget_crossed", - "max_budget_alert", - "soft_budget_crossed", - "threshold_crossed", - "projected_limit_exceeded", - "key_created", - "key_rotated", - "internal_user_created", - "spend_tracked", - ] - event_message: str # human-readable description of event - event_group: Litellm_EntityType - - -class SpecialModelNames(enum.Enum): - all_team_models = "all-team-models" - all_proxy_models = "all-proxy-models" - no_default_models = "no-default-models" - - -class SpecialProxyStrings(enum.Enum): - default_user_id = "default_user_id" # global proxy admin - - -class InvitationNew(LiteLLMPydanticObjectBase): - user_id: str - - -class InvitationUpdate(LiteLLMPydanticObjectBase): - invitation_id: str - is_accepted: bool - - -class InvitationDelete(LiteLLMPydanticObjectBase): - invitation_id: str - - -class InvitationModel(LiteLLMPydanticObjectBase): - id: str - user_id: str - is_accepted: bool - accepted_at: Optional[datetime] - expires_at: datetime - created_at: datetime - created_by: str - updated_at: datetime - updated_by: str - - -class InvitationClaim(LiteLLMPydanticObjectBase): - invitation_link: str - user_id: str - password: str - - -class ConfigFieldInfo(LiteLLMPydanticObjectBase): - field_name: str - field_value: Any - - -class CallbackOnUI(LiteLLMPydanticObjectBase): - litellm_callback_name: str - litellm_callback_params: Optional[list] - ui_callback_name: str - - -class AllCallbacks(LiteLLMPydanticObjectBase): - langfuse: CallbackOnUI = CallbackOnUI( - litellm_callback_name="langfuse", - ui_callback_name="Langfuse", - litellm_callback_params=[ - "LANGFUSE_PUBLIC_KEY", - "LANGFUSE_SECRET_KEY", - "LANGFUSE_HOST", - ], - ) - - otel: CallbackOnUI = CallbackOnUI( - litellm_callback_name="otel", - ui_callback_name="OpenTelemetry", - litellm_callback_params=[ - "OTEL_EXPORTER", - "OTEL_ENDPOINT", - "OTEL_HEADERS", - ], - ) - - s3: CallbackOnUI = CallbackOnUI( - litellm_callback_name="s3", - ui_callback_name="s3 Bucket (AWS)", - litellm_callback_params=[ - "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", - "AWS_REGION_NAME", - ], - ) - - openmeter: CallbackOnUI = CallbackOnUI( - litellm_callback_name="openmeter", - ui_callback_name="OpenMeter", - litellm_callback_params=[ - "OPENMETER_API_ENDPOINT", - "OPENMETER_API_KEY", - ], - ) - - custom_callback_api: CallbackOnUI = CallbackOnUI( - litellm_callback_name="custom_callback_api", - litellm_callback_params=["GENERIC_LOGGER_ENDPOINT", "GENERIC_LOGGER_HEADERS"], - ui_callback_name="Custom Callback API", - ) - - generic_api: CallbackOnUI = CallbackOnUI( - litellm_callback_name="generic_api", - litellm_callback_params=["GENERIC_LOGGER_ENDPOINT", "GENERIC_LOGGER_HEADERS"], - ui_callback_name="Custom Callback API", - ) - - datadog: CallbackOnUI = CallbackOnUI( - litellm_callback_name="datadog", - litellm_callback_params=["DD_API_KEY", "DD_SITE"], - ui_callback_name="Datadog", - ) - - braintrust: CallbackOnUI = CallbackOnUI( - litellm_callback_name="braintrust", - litellm_callback_params=["BRAINTRUST_API_KEY", "BRAINTRUST_API_BASE"], - ui_callback_name="Braintrust", - ) - - langsmith: CallbackOnUI = CallbackOnUI( - litellm_callback_name="langsmith", - litellm_callback_params=[ - "LANGSMITH_API_KEY", - "LANGSMITH_PROJECT", - "LANGSMITH_DEFAULT_RUN_NAME", - ], - ui_callback_name="Langsmith", - ) - - lago: CallbackOnUI = CallbackOnUI( - litellm_callback_name="lago", - litellm_callback_params=[ - "LAGO_API_BASE", - "LAGO_API_KEY", - "LAGO_API_EVENT_CODE", - "LAGO_API_CHARGE_BY", - ], - ui_callback_name="Lago Billing", - ) - - traceloop: CallbackOnUI = CallbackOnUI( - litellm_callback_name="traceloop", - litellm_callback_params=[ - "TRACELOOP_API_KEY", - ], - ui_callback_name="Traceloop", - ) - - -class SpendLogsMetadata(TypedDict): - """ - Specific metadata k,v pairs logged to spendlogs for easier cost tracking - """ - - additional_usage_values: Optional[ - dict - ] # covers provider-specific usage information - e.g. prompt caching - user_api_key: Optional[str] - user_api_key_alias: Optional[str] - user_api_key_team_id: Optional[str] - user_api_key_org_id: Optional[str] - user_api_key_user_id: Optional[str] - user_api_key_team_alias: Optional[str] - spend_logs_metadata: Optional[ - dict - ] # special param to log k,v pairs to spendlogs for a call - requester_ip_address: Optional[str] - applied_guardrails: Optional[List[str]] - mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] - vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] - guardrail_information: Optional[List[StandardLoggingGuardrailInformation]] - status: StandardLoggingPayloadStatus - proxy_server_request: Optional[str] - batch_models: Optional[List[str]] - error_information: Optional[StandardLoggingPayloadErrorInformation] - usage_object: Optional[dict] - model_map_information: Optional[StandardLoggingModelInformation] - cold_storage_object_key: Optional[ - str - ] # S3/GCS object key for cold storage retrieval - litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds - cost_breakdown: Optional[ - CostBreakdown - ] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) - - -class SpendLogsPayload(TypedDict): - request_id: str - call_type: str - api_key: str - spend: float - total_tokens: int - prompt_tokens: int - completion_tokens: int - startTime: Union[datetime, str] - endTime: Union[datetime, str] - completionStartTime: Optional[Union[datetime, str]] - model: str - model_id: Optional[str] - model_group: Optional[str] - mcp_namespaced_tool_name: Optional[str] - agent_id: Optional[str] - api_base: str - user: str - metadata: str # json str - cache_hit: str - cache_key: str - request_tags: str # json str - team_id: Optional[str] - organization_id: Optional[str] - end_user: Optional[str] - requester_ip_address: Optional[str] - custom_llm_provider: Optional[str] - messages: Optional[Union[str, list, dict]] - response: Optional[Union[str, list, dict]] - proxy_server_request: Optional[str] - session_id: Optional[str] - status: Literal["success", "failure"] - - -class SpanAttributes(str, enum.Enum): - # Note: We've taken this from opentelemetry-semantic-conventions-ai - # I chose to not add a new dependency to litellm for this - - # Semantic Conventions for LLM requests, this needs to be removed after - # OpenTelemetry Semantic Conventions support Gen AI. - # Issue at https://github.com/open-telemetry/opentelemetry-python/issues/3868 - # Refer to https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/llm-spans.md - - LLM_SYSTEM = "gen_ai.system" - LLM_REQUEST_MODEL = "gen_ai.request.model" - LLM_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens" - LLM_REQUEST_TEMPERATURE = "gen_ai.request.temperature" - LLM_REQUEST_TOP_P = "gen_ai.request.top_p" - LLM_PROMPTS = "gen_ai.prompt" - LLM_COMPLETIONS = "gen_ai.completion" - LLM_RESPONSE_MODEL = "gen_ai.response.model" - LLM_USAGE_COMPLETION_TOKENS = "gen_ai.usage.completion_tokens" - LLM_USAGE_PROMPT_TOKENS = "gen_ai.usage.prompt_tokens" - - # OTEL 1.38 attributes - GEN_AI_INPUT_MESSAGES = "gen_ai.input.messages" - GEN_AI_OUTPUT_MESSAGES = "gen_ai.output.messages" - GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens" - GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens" - GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens" - GEN_AI_OPERATION_NAME = "gen_ai.operation.name" - GEN_AI_REQUEST_ID = "gen_ai.request.id" - GEN_AI_SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions" - GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons" - - LLM_TOKEN_TYPE = "gen_ai.token.type" - # To be added - # LLM_RESPONSE_FINISH_REASON = "gen_ai.response.finish_reasons" - # LLM_RESPONSE_ID = "gen_ai.response.id" - - # LLM - LLM_REQUEST_TYPE = "llm.request.type" - LLM_USAGE_TOTAL_TOKENS = "llm.usage.total_tokens" - LLM_USAGE_TOKEN_TYPE = "llm.usage.token_type" - LLM_USER = "llm.user" - LLM_HEADERS = "llm.headers" - LLM_TOP_K = "llm.top_k" - LLM_IS_STREAMING = "llm.is_streaming" - LLM_FREQUENCY_PENALTY = "llm.frequency_penalty" - LLM_PRESENCE_PENALTY = "llm.presence_penalty" - LLM_CHAT_STOP_SEQUENCES = "llm.chat.stop_sequences" - LLM_REQUEST_FUNCTIONS = "llm.request.functions" - LLM_REQUEST_REPETITION_PENALTY = "llm.request.repetition_penalty" - LLM_RESPONSE_FINISH_REASON = "llm.response.finish_reason" - LLM_RESPONSE_STOP_REASON = "llm.response.stop_reason" - LLM_CONTENT_COMPLETION_CHUNK = "llm.content.completion.chunk" - - # OpenAI - LLM_OPENAI_RESPONSE_SYSTEM_FINGERPRINT = "gen_ai.openai.system_fingerprint" - LLM_OPENAI_API_BASE = "gen_ai.openai.api_base" - LLM_OPENAI_API_VERSION = "gen_ai.openai.api_version" - LLM_OPENAI_API_TYPE = "gen_ai.openai.api_type" - - -class ManagementEndpointLoggingPayload(LiteLLMPydanticObjectBase): - route: str - request_data: dict - response: Optional[dict] = None - exception: Optional[Any] = None - start_time: Optional[datetime] = None - end_time: Optional[datetime] = None - - -class ProxyException(Exception): - # NOTE: DO NOT MODIFY THIS - # This is used to map exactly to OPENAI Exceptions - def __init__( - self, - message: str, - type: str, - param: Optional[str], - code: Optional[Union[int, str]] = None, # maps to status code - headers: Optional[Dict[str, str]] = None, - openai_code: Optional[str] = None, # maps to 'code' in openai - provider_specific_fields: Optional[dict] = None, - ): - self.message = str(message) - self.type = type - self.param = param - self.openai_code = openai_code or code - # If we look on official python OpenAI lib, the code should be a string: - # https://github.com/openai/openai-python/blob/195c05a64d39c87b2dfdf1eca2d339597f1fce03/src/openai/types/shared/error_object.py#L11 - # Related LiteLLM issue: https://github.com/BerriAI/litellm/discussions/4834 - self.code = str(code) - if headers is not None: - for k, v in headers.items(): - if not isinstance(v, str): - headers[k] = str(v) - self.headers = headers or {} - self.provider_specific_fields = provider_specific_fields - # rules for proxyExceptions - # Litellm router.py returns "No healthy deployment available" when there are no deployments available - # Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487 - if ( - "No healthy deployment available" in self.message - or "No deployments available" in self.message - ): - self.code = "429" - elif RouterErrors.no_deployments_with_tag_routing.value in self.message: - self.code = "401" - - def to_dict(self) -> dict: - """Converts the ProxyException instance to a dictionary.""" - error_dict: Dict[str, Optional[Union[str, Dict]]] = { - "message": self.message, - "type": self.type, - "param": self.param, - "code": self.code, - } - if self.provider_specific_fields: - error_dict["provider_specific_fields"] = self.provider_specific_fields - return error_dict - - -class CommonProxyErrors(str, enum.Enum): - db_not_connected_error = ( - "DB not connected. See https://docs.litellm.ai/docs/proxy/virtual_keys" - ) - no_llm_router = "No models configured on proxy" - not_allowed_access = "Admin-only endpoint. Not allowed to access this." - not_premium_user = "You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/enterprise#trial. \nPricing: https://www.litellm.ai/#pricing" - max_parallel_request_limit_reached = ( - "Crossed TPM / RPM / Max Parallel Request Limit" - ) - missing_enterprise_package = "Missing litellm-enterprise package. Please install it to use this feature. Run `pip install litellm-enterprise`" - missing_enterprise_package_docker = ( - "This uses the enterprise folder - only available on the Docker image." - ) - - -class SpendCalculateRequest(LiteLLMPydanticObjectBase): - model: Optional[str] = None - messages: Optional[List] = None - completion_response: Optional[dict] = None - - -class ProxyErrorTypes(str, enum.Enum): - budget_exceeded = "budget_exceeded" - """ - Object was over budget - """ - no_db_connection = "no_db_connection" - """ - No database connection - """ - - token_not_found_in_db = "token_not_found_in_db" - """ - Requested token was not found in the database - """ - - key_model_access_denied = "key_model_access_denied" - """ - Key does not have access to the model - """ - - team_model_access_denied = "team_model_access_denied" - """ - Team does not have access to the model - """ - - user_model_access_denied = "user_model_access_denied" - """ - User does not have access to the model - """ - - org_model_access_denied = "org_model_access_denied" - """ - Organization does not have access to the model - """ - - expired_key = "expired_key" - """ - Key has expired - """ - - auth_error = "auth_error" - """ - General authentication error - """ - - internal_server_error = "internal_server_error" - """ - Internal server error - """ - - bad_request_error = "bad_request_error" - """ - Bad request error - """ - - not_found_error = "not_found_error" - """ - Not found error - """ - - validation_error = "validation_error" - """ - Validation error - """ - - cache_ping_error = "cache_ping_error" - """ - Cache ping error - """ - - team_member_permission_error = "team_member_permission_error" - """ - Team member permission error - """ - - key_vector_store_access_denied = "key_vector_store_access_denied" - """ - Key does not have access to the vector store - """ - - team_vector_store_access_denied = "team_vector_store_access_denied" - """ - Team does not have access to the vector store - """ - - org_vector_store_access_denied = "org_vector_store_access_denied" - """ - Organization does not have access to the vector store - """ - - team_member_already_in_team = "team_member_already_in_team" - """ - Team member is already in team - """ - - @classmethod - def get_model_access_error_type_for_object( - cls, object_type: Literal["key", "user", "team", "org"] - ) -> "ProxyErrorTypes": - """ - Get the model access error type for object_type - """ - if object_type == "key": - return cls.key_model_access_denied - elif object_type == "team": - return cls.team_model_access_denied - elif object_type == "user": - return cls.user_model_access_denied - elif object_type == "org": - return cls.org_model_access_denied - - @classmethod - def get_vector_store_access_error_type_for_object( - cls, object_type: Literal["key", "team", "org"] - ) -> "ProxyErrorTypes": - """ - Get the vector store access error type for object_type - """ - if object_type == "key": - return cls.key_vector_store_access_denied - elif object_type == "team": - return cls.team_vector_store_access_denied - elif object_type == "org": - return cls.org_vector_store_access_denied - - -DB_CONNECTION_ERROR_TYPES = ( - httpx.ConnectError, - httpx.ReadError, - httpx.ReadTimeout, -) - - -class SSOUserDefinedValues(TypedDict): - models: List[str] - user_id: str - user_email: Optional[str] - user_role: Optional[str] - max_budget: Optional[float] - budget_duration: Optional[str] - - -class VirtualKeyEvent(LiteLLMPydanticObjectBase): - created_by_user_id: str - created_by_user_role: str - created_by_key_alias: Optional[str] - request_kwargs: dict - - -class CreatePassThroughEndpoint(LiteLLMPydanticObjectBase): - path: str - target: str - headers: dict - - -class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): - user_id: str - team_id: str - budget_id: Optional[str] = None - spend: Optional[float] = 0.0 - litellm_budget_table: Optional[LiteLLM_BudgetTable] - - def safe_get_team_member_rpm_limit(self) -> Optional[int]: - if self.litellm_budget_table is not None: - return self.litellm_budget_table.rpm_limit - return None - - def safe_get_team_member_tpm_limit(self) -> Optional[int]: - if self.litellm_budget_table is not None: - return self.litellm_budget_table.tpm_limit - return None - - -#### Organization / Team Member Requests #### - - -class MemberAddRequest(LiteLLMPydanticObjectBase): - member: Union[List[Member], Member] = Field( - description="Member object or list of member objects to add. Each member must include either user_id or user_email, and a role" - ) - - def __init__(self, **data): - member_data = data.get("member") - if isinstance(member_data, list): - # If member is a list of dictionaries, convert each dictionary to a Member object - members = [ - Member(**item) if isinstance(item, dict) else item - for item in member_data - ] - # Replace member_data with the list of Member objects - data["member"] = members - elif isinstance(member_data, dict): - # If member is a dictionary, convert it to a single Member object - member = Member(**member_data) - # Replace member_data with the single Member object - data["member"] = member - # Call the superclass __init__ method to initialize the object - super().__init__(**data) - - -class OrgMemberAddRequest(LiteLLMPydanticObjectBase): - member: Union[List[OrgMember], OrgMember] - - def __init__(self, **data): - member_data = data.get("member") - if isinstance(member_data, list): - # If member is a list of dictionaries, convert each dictionary to a Member object - if all(isinstance(item, dict) for item in member_data): - members = [OrgMember(**item) for item in member_data] - else: - members = [item for item in member_data] - # Replace member_data with the list of Member objects - data["member"] = members - elif isinstance(member_data, dict): - # If member is a dictionary, convert it to a single Member object - member = OrgMember(**member_data) - # Replace member_data with the single Member object - data["member"] = member - # Call the superclass __init__ method to initialize the object - super().__init__(**data) - - -class TeamAddMemberResponse(LiteLLM_TeamTable): - updated_users: List[LiteLLM_UserTable] - updated_team_memberships: List[LiteLLM_TeamMembership] - - -class OrganizationAddMemberResponse(LiteLLMPydanticObjectBase): - organization_id: str - updated_users: List[LiteLLM_UserTable] - updated_organization_memberships: List[LiteLLM_OrganizationMembershipTable] - - -class MemberDeleteRequest(LiteLLMPydanticObjectBase): - user_id: Optional[str] = None - user_email: Optional[str] = None - - @model_validator(mode="before") - @classmethod - def check_user_info(cls, values): - if values.get("user_id") is None and values.get("user_email") is None: - raise ValueError("Either user id or user email must be provided") - return values - - -class MemberUpdateResponse(LiteLLMPydanticObjectBase): - user_id: str - user_email: Optional[str] = None - - -# Team Member Requests -class TeamMemberAddRequest(MemberAddRequest): - """ - Request body for adding members to a team. - - Example: - ```json - { - "team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849", - "member": { - "role": "user", - "user_id": "user123" - }, - "max_budget_in_team": 100.0 - } - ``` - """ - - team_id: str = Field(description="The ID of the team to add the member to") - max_budget_in_team: Optional[float] = Field( - default=None, - description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", - ) - - -class TeamMemberDeleteRequest(MemberDeleteRequest): - team_id: str - - -class TeamMemberUpdateRequest(TeamMemberDeleteRequest): - max_budget_in_team: Optional[float] = None - role: Optional[Literal["admin", "user"]] = None - tpm_limit: Optional[int] = Field( - default=None, description="Tokens per minute limit for this team member" - ) - rpm_limit: Optional[int] = Field( - default=None, description="Requests per minute limit for this team member" - ) - - -class TeamMemberUpdateResponse(MemberUpdateResponse): - team_id: str - max_budget_in_team: Optional[float] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - - -class TeamModelAddRequest(BaseModel): - """Request to add models to a team""" - - team_id: str - models: List[str] - - -class TeamModelDeleteRequest(BaseModel): - """Request to delete models from a team""" - - team_id: str - models: List[str] - - -# Organization Member Requests -class OrganizationMemberAddRequest(OrgMemberAddRequest): - organization_id: str - max_budget_in_organization: Optional[ - float - ] = None # Users max budget within the organization - - -class OrganizationMemberDeleteRequest(MemberDeleteRequest): - organization_id: str - - -ROLES_WITHIN_ORG = [ - LitellmUserRoles.ORG_ADMIN, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, -] - - -class OrganizationMemberUpdateRequest(OrganizationMemberDeleteRequest): - max_budget_in_organization: Optional[float] = None - role: Optional[LitellmUserRoles] = None - - @field_validator("role") - def validate_role( - cls, value: Optional[LitellmUserRoles] - ) -> Optional[LitellmUserRoles]: - if value is not None and value not in ROLES_WITHIN_ORG: - raise ValueError( - f"Invalid role. Must be one of: {[role.value for role in ROLES_WITHIN_ORG]}" - ) - return value - - -class OrganizationMemberUpdateResponse(MemberUpdateResponse): - organization_id: str - max_budget_in_organization: float - - -########################################## - - -class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): - team_member_budget_table: Optional[LiteLLM_BudgetTable] = None - - -class TeamInfoResponseObject(TypedDict): - team_id: str - team_info: TeamInfoResponseObjectTeamTable - keys: List - team_memberships: List[LiteLLM_TeamMembership] - - -class TeamListResponseObject(LiteLLM_TeamTable): - team_memberships: List[LiteLLM_TeamMembership] - keys: List # list of keys that belong to the team - - -class KeyListResponseObject(TypedDict, total=False): - keys: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] - total_count: Optional[int] - current_page: Optional[int] - total_pages: Optional[int] - - -class CurrentItemRateLimit(TypedDict): - current_requests: int - current_tpm: int - current_rpm: int - - -class LoggingCallbackStatus(TypedDict, total=False): - callbacks: List[str] - status: Literal["healthy", "unhealthy"] - details: Optional[str] - - -class KeyHealthResponse(TypedDict, total=False): - key: Literal["healthy", "unhealthy"] - logging_callbacks: Optional[LoggingCallbackStatus] - - -class SpecialHeaders(enum.Enum): - """Used by user_api_key_auth.py to get litellm key""" - - openai_authorization = "Authorization" - azure_authorization = "API-Key" - anthropic_authorization = "x-api-key" - google_ai_studio_authorization = "x-goog-api-key" - azure_apim_authorization = "Ocp-Apim-Subscription-Key" - custom_litellm_api_key = "x-litellm-api-key" - mcp_auth = "x-mcp-auth" - mcp_servers = "x-mcp-servers" - mcp_access_groups = "x-mcp-access-groups" - - -class LitellmDataForBackendLLMCall(TypedDict, total=False): - headers: dict - organization: str - timeout: Optional[float] - stream_timeout: Optional[float] - user: Optional[str] - num_retries: Optional[int] - - -class LitellmMetadataFromRequestHeaders(TypedDict, total=False): - """ - Headers a user can pass that will get added to litellm metadata for the request - """ - - spend_logs_metadata: Optional[dict] - agent_id: Optional[str] - trace_id: Optional[str] - - -class JWTKeyItem(TypedDict, total=False): - kid: str - - -JWKKeyValue = Union[List[JWTKeyItem], JWTKeyItem] - - -class JWKUrlResponse(TypedDict, total=False): - keys: JWKKeyValue - - -class UserManagementEndpointParamDocStringEnums(str, enum.Enum): - user_id_doc_str = ( - "Optional[str] - Specify a user id. If not set, a unique id will be generated." - ) - user_alias_doc_str = ( - "Optional[str] - A descriptive name for you to know who this user id refers to." - ) - teams_doc_str = "Optional[list] - specify a list of team id's a user belongs to." - user_email_doc_str = "Optional[str] - Specify a user email." - send_invite_email_doc_str = ( - "Optional[bool] - Specify if an invite email should be sent." - ) - user_role_doc_str = """Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`""" - max_budget_doc_str = """Optional[float] - Specify max budget for a given user.""" - budget_duration_doc_str = """Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").""" - models_doc_str = """Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)""" - tpm_limit_doc_str = ( - """Optional[int] - Specify tpm limit for a given user (Tokens per minute)""" - ) - rpm_limit_doc_str = ( - """Optional[int] - Specify rpm limit for a given user (Requests per minute)""" - ) - auto_create_key_doc_str = """bool - Default=True. Flag used for returning a key as part of the /user/new response""" - aliases_doc_str = """Optional[dict] - Model aliases for the user - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)""" - config_doc_str = """Optional[dict] - [DEPRECATED PARAM] User-specific config.""" - allowed_cache_controls_doc_str = """Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request-""" - blocked_doc_str = ( - """Optional[bool] - [Not Implemented Yet] Whether the user is blocked.""" - ) - guardrails_doc_str = """Optional[List[str]] - [Not Implemented Yet] List of active guardrails for the user""" - permissions_doc_str = """Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.""" - metadata_doc_str = """Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }""" - max_parallel_requests_doc_str = """Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.""" - soft_budget_doc_str = """Optional[float] - Get alerts when user crosses given budget, doesn't block requests.""" - model_max_budget_doc_str = """Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)""" - model_rpm_limit_doc_str = """Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" - model_tpm_limit_doc_str = """Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" - spend_doc_str = """Optional[float] - Amount spent by user. Default is 0. Will be updated by proxy whenever user is used.""" - team_id_doc_str = """Optional[str] - [DEPRECATED PARAM] The team id of the user. Default is None.""" - duration_doc_str = """Optional[str] - Duration for the key auto-created on `/user/new`. Default is None.""" - - -PassThroughEndpointLoggingResultValues = Union[ - ModelResponse, - TextCompletionResponse, - ImageResponse, - EmbeddingResponse, - VideoObject, - StandardPassThroughResponseObject, -] - - -class PassThroughEndpointLoggingTypedDict(TypedDict): - result: Optional[PassThroughEndpointLoggingResultValues] - kwargs: dict - - -LiteLLM_ManagementEndpoint_MetadataFields = [ - "model_rpm_limit", - "model_tpm_limit", - "rpm_limit_type", - "tpm_limit_type", - "enforced_params", - "temp_budget_increase", - "temp_budget_expiry", - "allowed_vector_store_indexes", -] - -LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ - "guardrails", - "policies", - "tags", - "team_member_key_duration", - "prompts", - "logging", - "secret_manager_settings", - "allowed_passthrough_routes", -] - - -class ProviderBudgetResponseObject(LiteLLMPydanticObjectBase): - """ - Configuration for a single provider's budget settings - """ - - budget_limit: Optional[float] # Budget limit in USD for the time period - time_period: Optional[str] # Time period for budget (e.g., '1d', '30d', '1mo') - spend: Optional[float] = 0.0 # Current spend for this provider - budget_reset_at: Optional[str] = None # When the current budget period resets - - -class ProviderBudgetResponse(LiteLLMPydanticObjectBase): - """ - Complete provider budget configuration and status. - Maps provider names to their budget configs. - """ - - providers: Dict[ - str, ProviderBudgetResponseObject - ] = {} # Dictionary mapping provider names to their budget configurations - - -class ProxyStateVariables(TypedDict): - """ - TypedDict for Proxy state variables. - """ - - spend_logs_row_count: int - - -UI_TEAM_ID = "litellm-dashboard" - - -class JWTAuthBuilderResult(TypedDict): - is_proxy_admin: bool - team_object: Optional[LiteLLM_TeamTable] - user_object: Optional[LiteLLM_UserTable] - end_user_object: Optional[LiteLLM_EndUserTable] - org_object: Optional[LiteLLM_OrganizationTable] - token: str - team_id: Optional[str] - user_id: Optional[str] - end_user_id: Optional[str] - org_id: Optional[str] - team_membership: Optional[LiteLLM_TeamMembership] - - -class ClientSideFallbackModel(TypedDict, total=False): - """ - Dictionary passed when client configuring input - """ - - model: Required[str] - messages: List[AllMessageValues] - - -ALL_FALLBACK_MODEL_VALUES = Union[str, ClientSideFallbackModel] - - -RBAC_ROLES = Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.TEAM, - LitellmUserRoles.INTERNAL_USER, -] - - -class OIDCPermissions(LiteLLMPydanticObjectBase): - models: Optional[List[str]] = None - routes: Optional[List[str]] = None - - -class RoleBasedPermissions(OIDCPermissions): - role: RBAC_ROLES - - model_config = { - "extra": "forbid", - } - - -class RoleMapping(BaseModel): - role: str - internal_role: RBAC_ROLES - - -class JWTLiteLLMRoleMap(BaseModel): - jwt_role: str - litellm_role: LitellmUserRoles - - -class ScopeMapping(OIDCPermissions): - scope: str - - model_config = { - "extra": "forbid", - } - - -class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): - """ - A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth. - - Attributes: - - admin_jwt_scope: The JWT scope required for proxy admin roles. - - admin_allowed_routes: list of allowed routes for proxy admin roles. - - team_jwt_scope: The JWT scope required for proxy team roles. - - team_id_jwt_field: The field in the JWT token that stores the team ID. Default - `client_id`. - - team_allowed_routes: list of allowed routes for proxy team roles. - - user_id_jwt_field: The field in the JWT token that stores the user id (maps to `LiteLLMUserTable`). Use this for internal employees. - - user_email_jwt_field: The field in the JWT token that stores the user email (maps to `LiteLLMUserTable`). Use this for internal employees. - - user_allowed_email_subdomain: If specified, only emails from specified subdomain will be allowed to access proxy. - - end_user_id_jwt_field: The field in the JWT token that stores the end-user ID (maps to `LiteLLMEndUserTable`). Turn this off by setting to `None`. Enables end-user cost tracking. Use this for external customers. - - public_key_ttl: Default - 600s. TTL for caching public JWT keys. - - public_allowed_routes: list of allowed routes for authenticated but unknown litellm role jwt tokens. - - enforce_rbac: If true, enforce RBAC for all routes. - - custom_validate: A custom function to validates the JWT token. - - oidc_userinfo_endpoint: OIDC UserInfo endpoint URL. When set along with oidc_userinfo_enabled, LiteLLM will call this endpoint with the access token to retrieve user identity information. - - oidc_userinfo_enabled: Enable fetching user info from OIDC UserInfo endpoint instead of just decoding JWT token. Default: False. - - oidc_userinfo_cache_ttl: TTL (in seconds) for caching UserInfo responses. Default: 300s (5 minutes). - - See `auth_checks.py` for the specific routes - """ - - admin_jwt_scope: str = "litellm_proxy_admin" - admin_allowed_routes: List[str] = [ - "management_routes", - "spend_tracking_routes", - "global_spend_tracking_routes", - "info_routes", - ] - team_id_jwt_field: Optional[str] = None - team_id_upsert: bool = False - team_ids_jwt_field: Optional[str] = None - upsert_sso_user_to_team: bool = False - team_allowed_routes: List[str] = ["openai_routes", "info_routes"] - team_id_default: Optional[str] = Field( - default=None, - description="If no team_id given, default permissions/spend-tracking to this team.s", - ) - team_alias_jwt_field: Optional[str] = Field( - default=None, - description="The field in the JWT token that stores the team name/alias. Will be resolved to team_id via database lookup.", - ) - - org_id_jwt_field: Optional[str] = None - org_alias_jwt_field: Optional[str] = Field( - default=None, - description="The field in the JWT token that stores the organization name/alias. Will be resolved to org_id via database lookup.", - ) - user_id_jwt_field: Optional[str] = None - user_email_jwt_field: Optional[str] = None - user_allowed_email_domain: Optional[str] = None - user_roles_jwt_field: Optional[str] = None - user_allowed_roles: Optional[List[str]] = None - user_id_upsert: bool = Field( - default=False, description="If user doesn't exist, upsert them into the db." - ) - end_user_id_jwt_field: Optional[str] = None - public_key_ttl: float = 600 - public_allowed_routes: List[str] = ["public_routes"] - enforce_rbac: bool = False - roles_jwt_field: Optional[str] = None # v2 on role mappings - role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[ - str - ] = None # can be either user / team, inferred from the role mapping - scope_mappings: Optional[List[ScopeMapping]] = None - enforce_scope_based_access: bool = False - enforce_team_based_model_access: bool = False - custom_validate: Optional[Callable[..., Literal[True]]] = None - ######################################################### - # Fields for syncing user team membership and roles with IDP provider - jwt_litellm_role_map: Optional[List[JWTLiteLLMRoleMap]] = None - sync_user_role_and_teams: bool = False - ######################################################### - ######################################################### - # OIDC UserInfo Endpoint Configuration - oidc_userinfo_endpoint: Optional[str] = Field( - default=None, - description="OIDC UserInfo endpoint URL. If set, LiteLLM will call this endpoint with the access token to retrieve user identity information.", - ) - oidc_userinfo_enabled: bool = Field( - default=False, - description="Enable fetching user info from OIDC UserInfo endpoint instead of just decoding JWT token.", - ) - oidc_userinfo_cache_ttl: float = Field( - default=300, - description="TTL (in seconds) for caching UserInfo responses. Default: 300s (5 minutes).", - ) - ######################################################### - - def __init__(self, **kwargs: Any) -> None: - # get the attribute names for this Pydantic model - allowed_keys = self.__annotations__.keys() - - invalid_keys = set(kwargs.keys()) - allowed_keys - user_roles_jwt_field = kwargs.get("user_roles_jwt_field") - user_allowed_roles = kwargs.get("user_allowed_roles") - object_id_jwt_field = kwargs.get("object_id_jwt_field") - role_mappings = kwargs.get("role_mappings") - scope_mappings = kwargs.get("scope_mappings") - enforce_scope_based_access = kwargs.get("enforce_scope_based_access") - custom_validate = kwargs.get("custom_validate") - - if custom_validate is not None: - fn = get_instance_fn(custom_validate) - validate_custom_validate_return_type(fn) - kwargs["custom_validate"] = fn - - if invalid_keys: - raise ValueError( - f"Invalid arguments provided: {', '.join(invalid_keys)}. Allowed arguments are: {', '.join(allowed_keys)}." - ) - if (user_roles_jwt_field is not None and user_allowed_roles is None) or ( - user_roles_jwt_field is None and user_allowed_roles is not None - ): - raise ValueError( - "user_allowed_roles must be provided if user_roles_jwt_field is set." - ) - - if object_id_jwt_field is not None and role_mappings is None: - raise ValueError( - "if object_id_jwt_field is set, role_mappings must also be set. Needed to infer if the caller is a user or team." - ) - - if scope_mappings is not None and not enforce_scope_based_access: - raise ValueError( - "scope_mappings must be set if enforce_scope_based_access is true." - ) - - super().__init__(**kwargs) - - -class PrismaCompatibleUpdateDBModel(TypedDict, total=False): - model_name: str - litellm_params: str - model_info: str - updated_at: str - updated_by: str - - -class SpecialManagementEndpointEnums(enum.Enum): - DEFAULT_ORGANIZATION = "default_organization" - - -class TransformRequestBody(BaseModel): - call_type: CallTypes - request_body: dict - - -class DefaultInternalUserParams(LiteLLMPydanticObjectBase): - """ - Default parameters to apply when a new user signs in via SSO or is created on the /user/new API endpoint - """ - - user_role: Optional[ - Literal[ - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - ] - ] = Field( - default=LitellmUserRoles.INTERNAL_USER, - description="Default role assigned to new users created", - ) - max_budget: Optional[float] = Field( - default=None, - description="Default maximum budget (in USD) for new users created", - ) - budget_duration: Optional[str] = Field( - default=None, - description="Default budget duration for new users (e.g. 'daily', 'weekly', 'monthly')", - ) - models: Optional[List[str]] = Field( - default=None, description="Default list of models that new users can access" - ) - - teams: Optional[Union[List[str], List[NewUserRequestTeam]]] = Field( - default=None, - description="Default teams for new users created", - ) - - -class BaseDailySpendTransaction(TypedDict): - date: str - api_key: str - model: Optional[str] - model_group: Optional[str] - mcp_namespaced_tool_name: Optional[str] - custom_llm_provider: Optional[str] - endpoint: Optional[str] - - # token count metrics - prompt_tokens: int - completion_tokens: int - cache_read_input_tokens: int - cache_creation_input_tokens: int - - # request level metrics - spend: float - api_requests: int - successful_requests: int - failed_requests: int - - -class DailyTeamSpendTransaction(BaseDailySpendTransaction): - team_id: str - - -class DailyOrganizationSpendTransaction(BaseDailySpendTransaction): - organization_id: str - - -class DailyUserSpendTransaction(BaseDailySpendTransaction): - user_id: str - - -class DailyEndUserSpendTransaction(BaseDailySpendTransaction): - end_user_id: str - - -class DailyTagSpendTransaction(BaseDailySpendTransaction): - request_id: Optional[str] - tag: str - - -class DailyAgentSpendTransaction(BaseDailySpendTransaction): - agent_id: str - - -class DBSpendUpdateTransactions(TypedDict): - """ - Internal Data Structure for buffering spend updates in Redis or in memory before committing them to the database - """ - - user_list_transactions: Optional[Dict[str, float]] - end_user_list_transactions: Optional[Dict[str, float]] - key_list_transactions: Optional[Dict[str, float]] - team_list_transactions: Optional[Dict[str, float]] - team_member_list_transactions: Optional[Dict[str, float]] - org_list_transactions: Optional[Dict[str, float]] - tag_list_transactions: Optional[Dict[str, float]] - - -class SpendUpdateQueueItem(TypedDict, total=False): - entity_type: Litellm_EntityType - entity_id: str - response_cost: Optional[float] - - -class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): - unified_file_id: str - file_object: Optional[OpenAIFileObject] = None - model_mappings: Dict[str, str] - flat_model_file_ids: List[str] - created_by: Optional[str] - updated_by: Optional[str] - storage_backend: Optional[str] = None - storage_url: Optional[str] = None - - -class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): - unified_object_id: str - model_object_id: str - file_purpose: Literal["batch", "fine-tune", "response"] - file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] - - -class EnterpriseLicenseData(TypedDict, total=False): - expiration_date: str - user_id: str - allowed_features: List[str] - max_users: int - max_teams: int - - -class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): - vector_store_id: str - custom_llm_provider: str - vector_store_name: Optional[str] - vector_store_description: Optional[str] - vector_store_metadata: Optional[Dict[str, Any]] - created_at: Optional[datetime] - updated_at: Optional[datetime] - litellm_credential_name: Optional[str] - litellm_params: Optional[Dict[str, Any]] - team_id: Optional[str] - user_id: Optional[str] - - -class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False): - vector_store: LiteLLM_ManagedVectorStoresTable - - -class CostEstimateRequest(LiteLLMPydanticObjectBase): - """Request body for /cost/estimate endpoint.""" - - model: str = Field(description="Model name (from /model_group/info)") - input_tokens: int = Field(description="Expected input tokens per request", ge=0) - output_tokens: int = Field(description="Expected output tokens per request", ge=0) - num_requests_per_day: Optional[int] = Field( - default=None, description="Number of requests per day", ge=0 - ) - num_requests_per_month: Optional[int] = Field( - default=None, description="Number of requests per month", ge=0 - ) - - -class CostEstimateResponse(LiteLLMPydanticObjectBase): - """Response body for /cost/estimate endpoint.""" - - model: str - input_tokens: int - output_tokens: int - num_requests_per_day: Optional[int] = None - num_requests_per_month: Optional[int] = None - # Per-request costs - cost_per_request: float = Field( - description="Total cost per request (includes margin)" - ) - input_cost_per_request: float = Field( - description="Input token cost per request (before margin)" - ) - output_cost_per_request: float = Field( - description="Output token cost per request (before margin)" - ) - margin_cost_per_request: float = Field( - default=0.0, description="Margin/fee added per request" - ) - # Daily costs (if num_requests_per_day provided) - daily_cost: Optional[float] = Field( - default=None, description="Total daily cost (includes margin)" - ) - daily_input_cost: Optional[float] = Field( - default=None, description="Daily input token cost" - ) - daily_output_cost: Optional[float] = Field( - default=None, description="Daily output token cost" - ) - daily_margin_cost: Optional[float] = Field( - default=None, description="Daily margin/fee" - ) - # Monthly costs (if num_requests_per_month provided) - monthly_cost: Optional[float] = Field( - default=None, description="Total monthly cost (includes margin)" - ) - monthly_input_cost: Optional[float] = Field( - default=None, description="Monthly input token cost" - ) - monthly_output_cost: Optional[float] = Field( - default=None, description="Monthly output token cost" - ) - monthly_margin_cost: Optional[float] = Field( - default=None, description="Monthly margin/fee" - ) - # Pricing info - input_cost_per_token: Optional[float] = None - output_cost_per_token: Optional[float] = None - provider: Optional[str] = None +import enum +import json +from datetime import datetime +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union + +import httpx +from pydantic import ( + BaseModel, + ConfigDict, + Field, + Json, + field_validator, + model_validator, +) +from typing_extensions import Required, TypedDict + +from litellm._uuid import uuid +from litellm.types.integrations.slack_alerting import AlertType +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIFileObject, + ResponsesAPIResponse, +) +from litellm.types.mcp import ( + MCPAuth, + MCPAuthType, + MCPCredentials, + MCPTransport, + MCPTransportType, +) +from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.router import RouterErrors, UpdateRouterConfig +from litellm.types.secret_managers.main import KeyManagementSystem +from litellm.types.utils import ( + CallTypes, + CostBreakdown, + EmbeddingResponse, + GenericBudgetConfigType, + ImageResponse, + LiteLLMBatch, + LiteLLMFineTuningJob, + LiteLLMPydanticObjectBase, + ModelResponse, + ProviderField, + StandardCallbackDynamicParams, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, + StandardPassThroughResponseObject, + TextCompletionResponse, +) +from litellm.types.videos.main import VideoObject + +from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span = Union[_Span, Any] +else: + Span = Any + + +class SupportedDBObjectType(str, enum.Enum): + """ + Supported database object types for fine-grained DB storage control. + Use in general_settings.supported_db_objects to specify which objects to load from DB. + """ + + MODELS = "models" + MCP = "mcp" + GUARDRAILS = "guardrails" + POLICIES = "policies" + VECTOR_STORES = "vector_stores" + PASS_THROUGH_ENDPOINTS = "pass_through_endpoints" + PROMPTS = "prompts" + MODEL_COST_MAP = "model_cost_map" + + def __str__(self): + return str(self.value) + + +class LiteLLMTeamRoles(enum.Enum): + # team admin + TEAM_ADMIN = "admin" + # team member + TEAM_MEMBER = "user" + + +class LitellmUserRoles(str, enum.Enum): + """ + Admin Roles: + PROXY_ADMIN: admin over the platform + PROXY_ADMIN_VIEW_ONLY: can login, view all own keys, view all spend + ORG_ADMIN: admin over a specific organization, can create teams, users only within their organization + + Internal User Roles: + INTERNAL_USER: can login, view/create/delete their own keys, view their spend + INTERNAL_USER_VIEW_ONLY: can login, view their own keys, view their own spend + + + Team Roles: + TEAM: used for JWT auth + + + Customer Roles: + CUSTOMER: External users -> these are customers + + """ + + # Admin Roles + PROXY_ADMIN = "proxy_admin" + PROXY_ADMIN_VIEW_ONLY = "proxy_admin_viewer" + + # Organization admins + ORG_ADMIN = "org_admin" + + # Internal User Roles + INTERNAL_USER = "internal_user" + INTERNAL_USER_VIEW_ONLY = "internal_user_viewer" + + # Team Roles + TEAM = "team" + + # Customer Roles - External users of proxy + CUSTOMER = "customer" + + def __str__(self): + return str(self.value) + + def values(self) -> List[str]: + return list(self.__annotations__.keys()) + + @property + def description(self): + """ + Descriptions for the enum values + """ + descriptions = { + "proxy_admin": "admin over litellm proxy, has all permissions", + "proxy_admin_viewer": "view all keys, view all spend", + "internal_user": "view/create/delete their own keys, view their own spend", + "internal_user_viewer": "view their own keys, view their own spend", + "team": "team scope used for JWT auth", + "customer": "customer", + } + return descriptions.get(self.value, "") + + @property + def ui_label(self): + """ + UI labels for the enum values + """ + ui_labels = { + "proxy_admin": "Admin (All Permissions)", + "proxy_admin_viewer": "Admin (View Only)", + "internal_user": "Internal User (Create/Delete/View)", + "internal_user_viewer": "Internal User (View Only)", + "team": "Team", + "customer": "Customer", + } + return ui_labels.get(self.value, "") + + @property + def is_internal_user_role(self) -> bool: + """returns true if this role is an `internal_user` or `internal_user_viewer` role""" + return self.value in [ + self.INTERNAL_USER, + self.INTERNAL_USER_VIEW_ONLY, + ] + + +class LitellmTableNames(str, enum.Enum): + """ + Enum for Table Names used by LiteLLM + """ + + TEAM_TABLE_NAME = "LiteLLM_TeamTable" + USER_TABLE_NAME = "LiteLLM_UserTable" + KEY_TABLE_NAME = "LiteLLM_VerificationToken" + PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" + MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable" + + +class Litellm_EntityType(enum.Enum): + """ + Enum for types of entities on litellm + + This enum allows specifying the type of entity that is being tracked in the database. + """ + + KEY = "key" + USER = "user" + END_USER = "end_user" + TEAM = "team" + TEAM_MEMBER = "team_member" + ORGANIZATION = "organization" + TAG = "tag" + + # global proxy level entity + PROXY = "proxy" + + +def hash_token(token: str): + import hashlib + + # Hash the string using SHA-256 + hashed_token = hashlib.sha256(token.encode()).hexdigest() + + return hashed_token + + +class KeyManagementRoutes(str, enum.Enum): + """ + Enum for key management routes + """ + + # write routes + KEY_GENERATE = "/key/generate" + KEY_UPDATE = "/key/update" + KEY_DELETE = "/key/delete" + KEY_REGENERATE = "/key/regenerate" + KEY_GENERATE_SERVICE_ACCOUNT = "/key/service-account/generate" + KEY_REGENERATE_WITH_PATH_PARAM = "/key/{key_id}/regenerate" + KEY_BLOCK = "/key/block" + KEY_UNBLOCK = "/key/unblock" + KEY_BULK_UPDATE = "/key/bulk_update" + + # info and health routes + KEY_INFO = "/key/info" + KEY_HEALTH = "/key/health" + + # list routes + KEY_LIST = "/key/list" + + +class LiteLLMRoutes(enum.Enum): + openai_route_names = [ + "chat_completion", + "completion", + "embeddings", + "image_generation", + "video_generation", + "audio_transcriptions", + "moderations", + "model_list", # OpenAI /v1/models route + ] + openai_routes = [ + # chat completions + "/engines/{model}/chat/completions", + "/openai/deployments/{model}/chat/completions", + "/chat/completions", + "/v1/chat/completions", + "/cursor/chat/completions", + # completions + "/engines/{model}/completions", + "/openai/deployments/{model}/completions", + "/completions", + "/v1/completions", + # embeddings + "/engines/{model}/embeddings", + "/openai/deployments/{model}/embeddings", + "/embeddings", + "/v1/embeddings", + # image generation + "/images/generations", + "/v1/images/generations", + # image edit + "/images/edits", + "/v1/images/edits", + # video generation + "/videos", + "/v1/videos", + "/videos/{video_id}", + "/v1/videos/{video_id}", + "/videos/{video_id}/content", + "/v1/videos/{video_id}/content", + "/videos/{video_id}/remix", + "/v1/videos/{video_id}/remix", + # audio transcription + "/audio/transcriptions", + "/v1/audio/transcriptions", + # audio Speech + "/audio/speech", + "/v1/audio/speech", + # moderations + "/moderations", + "/v1/moderations", + # batches + "/v1/batches", + "/batches", + "/v1/batches/{batch_id}", + "/batches/{batch_id}", + "/v1/batches/{batch_id}/cancel", + "/batches/{batch_id}/cancel", + # files + "/v1/files", + "/files", + "/v1/files/{file_id}", + "/files/{file_id}", + "/v1/files/{file_id}/content", + "/files/{file_id}/content", + # fine_tuning + "/fine_tuning/jobs", + "/v1/fine_tuning/jobs", + "/fine_tuning/jobs/{fine_tuning_job_id}/cancel", + "/v1/fine_tuning/jobs/{fine_tuning_job_id}/cancel", + # assistants-related routes + "/assistants", + "/v1/assistants", + "/v1/assistants/{assistant_id}", + "/assistants/{assistant_id}", + "/threads", + "/v1/threads", + "/threads/{thread_id}", + "/v1/threads/{thread_id}", + "/threads/{thread_id}/messages", + "/v1/threads/{thread_id}/messages", + "/threads/{thread_id}/runs", + "/v1/threads/{thread_id}/runs", + # models + "/models", + "/v1/models", + # token counter + "/utils/token_counter", + "/utils/transform_request", + # rerank + "/rerank", + "/v1/rerank", + "/v2/rerank", + # realtime + "/realtime", + "/v1/realtime", + "/realtime?{model}", + "/v1/realtime?{model}", + # responses API + "/responses", + "/v1/responses", + "/responses/{response_id}", + "/v1/responses/{response_id}", + "/responses/{response_id}/input_items", + "/v1/responses/{response_id}/input_items", + "/responses/{response_id}/cancel", + "/v1/responses/{response_id}/cancel", + # vector stores + "/vector_stores", + "/v1/vector_stores", + "/vector_stores/{vector_store_id}/search", + "/v1/vector_stores/{vector_store_id}/search", + "/vector_stores/{vector_store_id}/files", + "/v1/vector_stores/{vector_store_id}/files", + "/vector_stores/{vector_store_id}/files/{file_id}", + "/v1/vector_stores/{vector_store_id}/files/{file_id}", + "/vector_stores/{vector_store_id}/files/{file_id}/content", + "/v1/vector_stores/{vector_store_id}/files/{file_id}/content", + "/vector_store/list", + "/v1/vector_store/list", + # search + "/search", + "/v1/search", + "/search/{search_tool_name}", + "/v1/search/{search_tool_name}", + # OCR + "/ocr", + "/v1/ocr", + # containers API + "/containers", + "/v1/containers", + "/containers/*", + "/v1/containers/*", + ] + + mapped_pass_through_routes = [ + "/bedrock", + "/vertex-ai", + "/vertex_ai", + "/cohere", + "/gemini", + "/anthropic", + "/langfuse", + "/azure", + "/azure_ai", + "/openai", + "/openai_passthrough", + "/assemblyai", + "/eu.assemblyai", + "/vllm", + "/mistral", + "/milvus", + ] + + ######################################################### + # e.g /vllm/*, anthropic/*, etc. + # allows using /anthropic/v1/messages, /vllm/v1/chat/completions, etc. + ######################################################### + passthrough_routes_wildcard = [f"{route}/*" for route in mapped_pass_through_routes] + + litellm_native_routes = [ + "/rag/ingest", + "/v1/rag/ingest", + "/rag/query", + "/v1/rag/query", + ] + + anthropic_routes = [ + "/v1/messages", + "/v1/messages/count_tokens", + "/v1/skills", + "/v1/skills/{skill_id}", + ] + + mcp_routes = [ + "/mcp", + "/mcp/", + "/mcp/{subpath}", + "/mcp/tools", + "/mcp/tools/list", + "/mcp/tools/call", + ] + + agent_routes = [ + "/v1/agents", + "/agents", + "/a2a/{agent_id}", + "/a2a/{agent_id}/message/send", + "/a2a/{agent_id}/message/stream", + "/a2a/{agent_id}/.well-known/agent-card.json", + ] + + google_routes = [ + "/v1beta/models/{model_name:path}:countTokens", + "/v1beta/models/{model_name:path}:generateContent", + "/v1beta/models/{model_name:path}:streamGenerateContent", + "/models/{model_name:path}:countTokens", + "/models/{model_name:path}:generateContent", + "/models/{model_name:path}:streamGenerateContent", + # Google Interactions API + "/interactions", + "/v1beta/interactions", + "/interactions/{interaction_id}", + "/v1beta/interactions/{interaction_id}", + "/interactions/{interaction_id}/cancel", + "/v1beta/interactions/{interaction_id}/cancel", + ] + + apply_guardrail_routes = [ + "/guardrails/apply_guardrail", + ] + + llm_api_routes = ( + openai_routes + + anthropic_routes + + google_routes + + mapped_pass_through_routes + + passthrough_routes_wildcard + + apply_guardrail_routes + + mcp_routes + + litellm_native_routes + + agent_routes + ) + info_routes = [ + "/key/info", + "/key/health", + "/team/info", + "/team/list", + "/v2/team/list", + "/organization/list", + "/team/available", + "/user/info", + "/model/info", + "/v1/model/info", + "/v2/model/info", + "/v2/key/info", + "/model_group/info", + "/health", + "/key/list", + "/user/filter/ui", + "/models", + "/v1/models", + ] + + # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend + master_key_only_routes = [ + "/global/spend/reset", + "/memory-usage-in-mem-cache", + "/memory-usage-in-mem-cache-items", + ] + + key_management_routes = [ + KeyManagementRoutes.KEY_GENERATE.value, + KeyManagementRoutes.KEY_UPDATE.value, + KeyManagementRoutes.KEY_DELETE.value, + KeyManagementRoutes.KEY_INFO.value, + KeyManagementRoutes.KEY_REGENERATE.value, + KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT.value, + KeyManagementRoutes.KEY_REGENERATE_WITH_PATH_PARAM.value, + KeyManagementRoutes.KEY_LIST.value, + KeyManagementRoutes.KEY_BLOCK.value, + KeyManagementRoutes.KEY_UNBLOCK.value, + KeyManagementRoutes.KEY_BULK_UPDATE.value, + ] + + management_routes = [ + # user + "/user/new", + "/user/update", + "/user/delete", + "/user/info", + "/user/list", + # team + "/team/new", + "/team/update", + "/team/delete", + "/team/list", + "/v2/team/list", + "/team/info", + "/team/block", + "/team/unblock", + "/team/available", + "/team/permissions_list", + "/team/permissions_update", + # model + "/model/new", + "/model/update", + "/model/delete", + "/model/info", + ] + key_management_routes + + spend_tracking_routes = [ + # spend + "/spend/keys", + "/spend/users", + "/spend/tags", + "/spend/calculate", + "/spend/logs", + "/cost/estimate", + ] + + global_spend_tracking_routes = [ + # global spend + "/global/spend/logs", + "/global/spend", + "/global/spend/keys", + "/global/spend/teams", + "/global/spend/end_users", + "/global/spend/models", + "/global/predict/spend/logs", + "/global/spend/report", + "/global/spend/provider", + "/global/spend/tags", + ] + + public_routes = set( + [ + "/routes", + "/", + "/health/liveliness", + "/health/liveness", + "/health/readiness", + "/test", + "/config/yaml", + "/metrics", + "/litellm/.well-known/litellm-ui-config", + "/.well-known/litellm-ui-config", + "/public/model_hub", + "/public/agent_hub", + "/public/mcp_hub", + "/public/litellm_model_cost_map", + ] + ) + + ui_routes = [ + "/sso", + "/sso/get/ui_settings", + "/get/ui_settings", + "/login", + "/key/info", + "/config", + "/spend", + "/model/info", + "/v2/model/info", + "/v2/key/info", + "/models", + "/v1/models", + "/global/spend", + "/global/spend/logs", + "/global/spend/keys", + "/global/spend/models", + "/global/spend/tags", + "/global/predict/spend/logs", + "/global/activity", + "/health/services", + ] + info_routes + + internal_user_routes = ( + [ + "/global/spend/tags", + "/global/spend/keys", + "/global/spend/models", + "/global/spend/provider", + "/global/spend/end_users", + "/global/activity", + "/global/activity/model", + "/v1/models/{model_id}", + "/models/{model_id}", + ] + + spend_tracking_routes + + key_management_routes + ) + + internal_user_view_only_routes = ( + spend_tracking_routes + global_spend_tracking_routes + ) + + self_managed_routes = [ + "/team/member_add", + "/team/member_delete", + "/team/member_update", + "/team/permissions_list", + "/team/permissions_update", + "/team/daily/activity", + "/model/new", + "/model/update", + "/model/delete", + "/user/daily/activity", + "/model/{model_id}/update", + "/prompt/list", + "/prompt/info", + ] # routes that manage their own allowed/disallowed logic + + ## Org Admin Routes ## + + # Routes only an Org Admin Can Access + org_admin_only_routes = [ + "/organization/info", + "/organization/delete", + "/organization/member_add", + "/organization/member_update", + ] + + # Routes accessible by Admin Viewer (read-only admin access) + admin_viewer_routes = [ + "/user/list", + "/user/available_users", + "/user/available_roles", + "/user/daily/activity", + "/team/daily/activity", + "/tag/daily/activity", + "/tag/list", + ] + info_routes + + # All routes accesible by an Org Admin + org_admin_allowed_routes = ( + org_admin_only_routes + + management_routes + + self_managed_routes + + admin_viewer_routes + ) + + +class LiteLLMPromptInjectionParams(LiteLLMPydanticObjectBase): + heuristics_check: bool = False + vector_db_check: bool = False + llm_api_check: bool = False + llm_api_name: Optional[str] = None + llm_api_system_prompt: Optional[str] = None + llm_api_fail_call_string: Optional[str] = None + reject_as_response: Optional[bool] = Field( + default=False, + description="Return rejected request error message as a string to the user. Default behaviour is to raise an exception.", + ) + + @model_validator(mode="before") + @classmethod + def check_llm_api_params(cls, values): + llm_api_check = values.get("llm_api_check") + if llm_api_check is True: + if "llm_api_name" not in values or not values["llm_api_name"]: + raise ValueError( + "If llm_api_check is set to True, llm_api_name must be provided" + ) + if ( + "llm_api_system_prompt" not in values + or not values["llm_api_system_prompt"] + ): + raise ValueError( + "If llm_api_check is set to True, llm_api_system_prompt must be provided" + ) + if ( + "llm_api_fail_call_string" not in values + or not values["llm_api_fail_call_string"] + ): + raise ValueError( + "If llm_api_check is set to True, llm_api_fail_call_string must be provided" + ) + return values + + +######### Request Class Definition ###### +class ProxyChatCompletionRequest(LiteLLMPydanticObjectBase): + """ + Pydantic model for chat completion requests that includes both OpenAI standard fields + and LiteLLM-specific parameters. This replaces the previous TypedDict version. + """ + + # Required fields (from ChatCompletionRequest) + model: str + messages: List[AllMessageValues] + + # Standard OpenAI completion parameters (all optional) + frequency_penalty: Optional[float] = None + logit_bias: Optional[Dict[str, float]] = None + logprobs: Optional[bool] = None + top_logprobs: Optional[int] = None + max_tokens: Optional[int] = None + n: Optional[int] = None + presence_penalty: Optional[float] = None + response_format: Optional[Dict[str, Any]] = None + seed: Optional[int] = None + service_tier: Optional[str] = None + stop: Optional[Union[str, List[str]]] = None + stream_options: Optional[Dict[str, Any]] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + tools: Optional[List[Dict[str, Any]]] = None + tool_choice: Optional[Union[str, Dict[str, Any]]] = None + parallel_tool_calls: Optional[bool] = None + function_call: Optional[Union[str, Dict[str, Any]]] = None + functions: Optional[List[Dict[str, Any]]] = None + user: Optional[str] = None + stream: Optional[bool] = None + + # LiteLLM-specific metadata param (from original ChatCompletionRequest) + metadata: Optional[Dict[str, Any]] = None + + # Optional LiteLLM params + guardrails: Optional[List[str]] = None + caching: Optional[bool] = None + num_retries: Optional[int] = None + context_window_fallback_dict: Optional[Dict[str, str]] = None + fallbacks: Optional[List[str]] = None + + +class ModelInfoDelete(LiteLLMPydanticObjectBase): + id: str + + +class ModelInfo(LiteLLMPydanticObjectBase): + id: Optional[str] + mode: Optional[Literal["embedding", "chat", "completion"]] + input_cost_per_token: Optional[float] = 0.0 + output_cost_per_token: Optional[float] = 0.0 + max_tokens: Optional[int] = 2048 # assume 2048 if not set + + # for azure models we need users to specify the base model, one azure you can call deployments - azure/my-random-model + # we look up the base model in model_prices_and_context_window.json + base_model: Optional[ + Literal[ + "gpt-4-1106-preview", + "gpt-4-32k", + "gpt-4", + "gpt-3.5-turbo-16k", + "gpt-3.5-turbo", + "text-embedding-ada-002", + ] + ] + + model_config = ConfigDict(protected_namespaces=(), extra="allow") + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("id") is None: + values.update({"id": str(uuid.uuid4())}) + if values.get("mode") is None: + values.update({"mode": None}) + if values.get("input_cost_per_token") is None: + values.update({"input_cost_per_token": None}) + if values.get("output_cost_per_token") is None: + values.update({"output_cost_per_token": None}) + if values.get("max_tokens") is None: + values.update({"max_tokens": None}) + if values.get("base_model") is None: + values.update({"base_model": None}) + return values + + +class ProviderInfo(LiteLLMPydanticObjectBase): + name: str + fields: List[ProviderField] + + +class BlockUsers(LiteLLMPydanticObjectBase): + user_ids: List[str] # required + + +class ModelParams(LiteLLMPydanticObjectBase): + model_name: str + litellm_params: dict + model_info: ModelInfo + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("model_info") is None: + values.update( + {"model_info": ModelInfo(id=None, mode="chat", base_model=None)} + ) + return values + + +class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): + mcp_servers: Optional[List[str]] = None + mcp_access_groups: Optional[List[str]] = None + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None + vector_stores: Optional[List[str]] = None + agents: Optional[List[str]] = None + agent_access_groups: Optional[List[str]] = None + + +class GenerateRequestBase(LiteLLMPydanticObjectBase): + """ + Overlapping schema between key and user generate/update requests + """ + + key_alias: Optional[str] = None + duration: Optional[str] = None + models: Optional[list] = [] + spend: Optional[float] = 0 + max_budget: Optional[float] = None + user_id: Optional[str] = None + team_id: Optional[str] = None + max_parallel_requests: Optional[int] = None + metadata: Optional[dict] = {} + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + + budget_duration: Optional[str] = None + allowed_cache_controls: Optional[list] = [] + config: Optional[dict] = {} + permissions: Optional[dict] = {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + + model_config = ConfigDict(protected_namespaces=()) + model_rpm_limit: Optional[dict] = None + model_tpm_limit: Optional[dict] = None + guardrails: Optional[List[str]] = None + policies: Optional[List[str]] = None + prompts: Optional[List[str]] = None + blocked: Optional[bool] = None + aliases: Optional[dict] = {} + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + + @field_validator("max_budget", mode="before") + @classmethod + def check_max_budget(cls, v): + if v == "": + return None + return v + + +class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): + index_name: str + index_permissions: List[Literal["read", "write"]] + + +class KeyRequestBase(GenerateRequestBase): + key: Optional[str] = None + budget_id: Optional[str] = None + tags: Optional[List[str]] = None + enforced_params: Optional[List[str]] = None + allowed_routes: Optional[list] = [] + allowed_passthrough_routes: Optional[list] = None + allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None + rpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] + ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm + tpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] + ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm + router_settings: Optional[UpdateRouterConfig] = None + + +class LiteLLMKeyType(str, enum.Enum): + """ + Enum for key types that determine what routes a key can access + """ + + LLM_API = "llm_api" # Can call LLM API routes (chat/completions, embeddings, etc.) + MANAGEMENT = "management" # Can call management routes (user/team/key management) + READ_ONLY = "read_only" # Can only call info/read routes + DEFAULT = "default" # Uses default allowed routes + + +class GenerateKeyRequest(KeyRequestBase): + soft_budget: Optional[float] = None + send_invite_email: Optional[bool] = None + key_type: Optional[LiteLLMKeyType] = Field( + default=LiteLLMKeyType.DEFAULT, + description="Type of key that determines default allowed routes.", + ) + auto_rotate: Optional[bool] = Field( + default=False, description="Whether this key should be automatically rotated" + ) + rotation_interval: Optional[str] = Field( + default=None, + description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True", + ) + organization_id: Optional[str] = None + + +class GenerateKeyResponse(KeyRequestBase): + key: str # type: ignore + key_name: Optional[str] = None + expires: Optional[datetime] = None + user_id: Optional[str] = None + token_id: Optional[str] = None + organization_id: Optional[str] = None + litellm_budget_table: Optional[Any] = None + token: Optional[str] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("token") is not None: + values.update({"key": values.get("token")}) + dict_fields = [ + "metadata", + "aliases", + "config", + "permissions", + "model_max_budget", + "router_settings", + ] + for field in dict_fields: + value = values.get(field) + if value is not None and isinstance(value, str): + try: + values[field] = json.loads(value) + except json.JSONDecodeError: + raise ValueError(f"Field {field} should be a valid dictionary") + + return values + + +class UpdateKeyRequest(KeyRequestBase): + # Note: the defaults of all Params here MUST BE NONE + # else they will get overwritten + key: str # type: ignore + duration: Optional[str] = None + spend: Optional[float] = None + metadata: Optional[dict] = None + temp_budget_increase: Optional[float] = None + temp_budget_expiry: Optional[datetime] = None + auto_rotate: Optional[bool] = None + rotation_interval: Optional[str] = None + + @model_validator(mode="after") + def validate_temp_budget(self) -> "UpdateKeyRequest": + if self.temp_budget_increase is not None or self.temp_budget_expiry is not None: + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + raise ValueError( + "temp_budget_increase and temp_budget_expiry must be set together" + ) + return self + + +class RegenerateKeyRequest(GenerateKeyRequest): + # This needs to be different from UpdateKeyRequest, because "key" is optional for this + key: Optional[str] = None + new_key: Optional[str] = None + duration: Optional[str] = None + spend: Optional[float] = None + metadata: Optional[dict] = None + new_master_key: Optional[str] = None + + +class KeyRequest(LiteLLMPydanticObjectBase): + keys: Optional[List[str]] = None + key_aliases: Optional[List[str]] = None + + @model_validator(mode="before") + @classmethod + def validate_at_least_one(cls, values): + if not values.get("keys") and not values.get("key_aliases"): + raise ValueError( + "At least one of 'keys' or 'key_aliases' must be provided." + ) + return values + + +class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): + id: Optional[int] = None + model_aliases: Optional[Union[str, dict]] = None # json dump the dict + created_by: str + updated_by: str + team: Optional["LiteLLM_TeamTable"] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): + model_id: str + model_name: str + litellm_params: dict + model_info: dict + created_at: Optional[datetime] = None + created_by: str + updated_at: Optional[datetime] = None + updated_by: str + + @model_validator(mode="before") + @classmethod + def check_potential_json_str(cls, values): + if isinstance(values.get("litellm_params"), str): + try: + values["litellm_params"] = json.loads(values["litellm_params"]) + except json.JSONDecodeError: + pass + if isinstance(values.get("model_info"), str): + try: + values["model_info"] = json.loads(values["model_info"]) + except json.JSONDecodeError: + pass + return values + + +# MCP Types +class SpecialMCPServerName(str, enum.Enum): + all_team_servers = "all-team-mcpservers" + all_proxy_servers = "all-proxy-mcpservers" + + +# MCP Proxy Request Types +class NewMCPServerRequest(LiteLLMPydanticObjectBase): + server_id: Optional[str] = None + server_name: Optional[str] = None + alias: Optional[str] = None + description: Optional[str] = None + transport: MCPTransportType = MCPTransport.sse + auth_type: Optional[MCPAuthType] = None + credentials: Optional[MCPCredentials] = None + url: Optional[str] = None + mcp_info: Optional[MCPInfo] = None + mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: Optional[List[str]] = None + extra_headers: Optional[List[str]] = None + static_headers: Optional[Dict[str, str]] = None + # Stdio-specific fields + command: Optional[str] = None + args: List[str] = Field(default_factory=list) + env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None + allow_all_keys: bool = False + + @model_validator(mode="before") + @classmethod + def validate_transport_fields(cls, values): + if isinstance(values, dict): + transport = values.get("transport") + if transport == MCPTransport.stdio: + if not values.get("command"): + raise ValueError("command is required for stdio transport") + if not values.get("args"): + raise ValueError("args is required for stdio transport") + elif transport in [MCPTransport.http, MCPTransport.sse]: + if not values.get("url"): + raise ValueError("url is required for HTTP/SSE transport") + return values + + @model_validator(mode="before") + @classmethod + def validate_credentials_requirements(cls, values): + if not isinstance(values, dict): + return values + + auth_type = values.get("auth_type") + if auth_type in {MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic}: + credentials = values.get("credentials") + auth_value = None + if isinstance(credentials, dict): + auth_value = credentials.get("auth_value") + elif hasattr(credentials, "get"): + auth_value = credentials.get("auth_value") # type: ignore[attr-defined] + + if not auth_value: + raise ValueError( + "auth_value is required when auth_type is api_key, bearer_token, or basic" + ) + + return values + + +class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): + server_id: str + server_name: Optional[str] = None + alias: Optional[str] = None + description: Optional[str] = None + transport: MCPTransportType = MCPTransport.sse + auth_type: Optional[MCPAuthType] = None + credentials: Optional[MCPCredentials] = None + url: Optional[str] = None + mcp_info: Optional[MCPInfo] = None + mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: Optional[List[str]] = None + extra_headers: Optional[List[str]] = None + static_headers: Optional[Dict[str, str]] = None + # Stdio-specific fields + command: Optional[str] = None + args: List[str] = Field(default_factory=list) + env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None + allow_all_keys: bool = False + + @model_validator(mode="before") + @classmethod + def validate_transport_fields(cls, values): + if isinstance(values, dict): + transport = values.get("transport") + if transport == MCPTransport.stdio: + if not values.get("command"): + raise ValueError("command is required for stdio transport") + if not values.get("args"): + raise ValueError("args is required for stdio transport") + elif transport in [MCPTransport.http, MCPTransport.sse]: + if not values.get("url"): + raise ValueError("url is required for HTTP/SSE transport") + return values + + +class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_MCPServerTable record""" + + server_id: str + server_name: Optional[str] = None + alias: Optional[str] = None + description: Optional[str] = None + url: Optional[str] = None + transport: MCPTransportType + auth_type: Optional[MCPAuthType] = None + credentials: Optional[MCPCredentials] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) + mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: List[str] = Field(default_factory=list) + extra_headers: List[str] = Field(default_factory=list) + mcp_info: Optional[MCPInfo] = None + static_headers: Optional[Dict[str, str]] = None + # Health check status + status: Optional[Literal["healthy", "unhealthy", "unknown"]] = Field( + default="unknown", + description="Health status: 'healthy', 'unhealthy', 'unknown'", + ) + last_health_check: Optional[datetime] = None + health_check_error: Optional[str] = None + # Stdio-specific fields + command: Optional[str] = None + args: List[str] = Field(default_factory=list) + env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None + allow_all_keys: bool = False + + +class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): + mcp_server_ids: List[str] + + +######## Skills API Types ######## + + +class NewSkillRequest(LiteLLMPydanticObjectBase): + """Request to create a new skill in LiteLLM database""" + + display_title: Optional[str] = None + description: Optional[str] = None + instructions: Optional[str] = None + file_content: Optional[bytes] = None # Binary content of skill files (zip) + file_name: Optional[str] = None # Original filename + file_type: Optional[str] = None # MIME type (e.g., "application/zip") + metadata: Optional[Dict[str, Any]] = None + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None + + +class UpdateSkillRequest(LiteLLMPydanticObjectBase): + """Request to update an existing skill""" + + skill_id: str + display_title: Optional[str] = None + description: Optional[str] = None + instructions: Optional[str] = None + file_content: Optional[bytes] = None # Binary content of skill files (zip) + file_name: Optional[str] = None # Original filename + file_type: Optional[str] = None # MIME type + metadata: Optional[Dict[str, Any]] = None + + +class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_SkillsTable record""" + + skill_id: str + display_title: Optional[str] = None + description: Optional[str] = None + instructions: Optional[str] = None + source: str = "custom" + latest_version: Optional[str] = None + file_content: Optional[bytes] = None # Binary content of skill files (zip) + file_name: Optional[str] = None # Original filename + file_type: Optional[str] = None # MIME type + metadata: Optional[Dict[str, Any]] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + + +class ListSkillsRequest(LiteLLMPydanticObjectBase): + """Request to list skills from LiteLLM database""" + + limit: Optional[int] = 20 + offset: Optional[int] = 0 + + +class NewUserRequestTeam(LiteLLMPydanticObjectBase): + team_id: str + max_budget_in_team: Optional[float] = None + user_role: Literal["user", "admin"] = "user" + + +class NewUserRequest(GenerateRequestBase): + max_budget: Optional[float] = None + user_email: Optional[str] = None + user_alias: Optional[str] = None + user_role: Optional[ + Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + ] = None + teams: Optional[Union[List[str], List[NewUserRequestTeam]]] = None + auto_create_key: bool = ( + True # flag used for returning a key as part of the /user/new response + ) + send_invite_email: Optional[bool] = None + sso_user_id: Optional[str] = None + organizations: Optional[List[str]] = None + + +class NewUserResponse(GenerateKeyResponse): + max_budget: Optional[float] = None + user_email: Optional[str] = None + user_role: Optional[ + Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + ] = None + teams: Optional[list] = None + user_alias: Optional[str] = None + model_max_budget: Optional[dict] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + + +class UpdateUserRequestNoUserIDorEmail( + GenerateRequestBase +): # shared with BulkUpdateUserRequest + password: Optional[str] = None + spend: Optional[float] = None + metadata: Optional[dict] = None + user_alias: Optional[str] = None + user_role: Optional[ + Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + ] = None + max_budget: Optional[float] = None + + +class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): + # Note: the defaults of all Params here MUST BE NONE + # else they will get overwritten + user_id: Optional[str] = None + user_email: Optional[str] = None + + @model_validator(mode="before") + @classmethod + def check_user_info(cls, values): + if values.get("user_id") is None and values.get("user_email") is None: + raise ValueError("Either user id or user email must be provided") + return values + + +class DeleteUserRequest(LiteLLMPydanticObjectBase): + user_ids: List[str] # required + + +AllowedModelRegion = Literal["eu", "us"] + + +class BudgetNewRequest(LiteLLMPydanticObjectBase): + budget_id: Optional[str] = Field(default=None, description="The unique budget id.") + max_budget: Optional[float] = Field( + default=None, + description="Requests will fail if this budget (in USD) is exceeded.", + ) + soft_budget: Optional[float] = Field( + default=None, + description="Requests will NOT fail if this is exceeded. Will fire alerting though.", + ) + max_parallel_requests: Optional[int] = Field( + default=None, description="Max concurrent requests allowed for this budget id." + ) + tpm_limit: Optional[int] = Field( + default=None, description="Max tokens per minute, allowed for this budget id." + ) + rpm_limit: Optional[int] = Field( + default=None, description="Max requests per minute, allowed for this budget id." + ) + budget_duration: Optional[str] = Field( + default=None, + description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')", + ) + model_max_budget: Optional[GenericBudgetConfigType] = Field( + default=None, + description="Max budget for each model (e.g. {'gpt-4o': {'max_budget': '0.0000001', 'budget_duration': '1d', 'tpm_limit': 1000, 'rpm_limit': 1000}})", + ) + budget_reset_at: Optional[datetime] = Field( + default=None, + description="Datetime when the budget is reset", + ) + + +class BudgetRequest(LiteLLMPydanticObjectBase): + budgets: List[str] + + +class BudgetDeleteRequest(LiteLLMPydanticObjectBase): + id: str + + +class CustomerBase(LiteLLMPydanticObjectBase): + user_id: str + alias: Optional[str] = None + spend: float = 0.0 + allowed_model_region: Optional[AllowedModelRegion] = None + default_model: Optional[str] = None + budget_id: Optional[str] = None + litellm_budget_table: Optional[BudgetNewRequest] = None + blocked: bool = False + + +class NewCustomerRequest(BudgetNewRequest): + """ + Create a new customer, allocate a budget to them + """ + + user_id: str + alias: Optional[str] = None # human-friendly alias + blocked: bool = False # allow/disallow requests for this end-user + budget_id: Optional[str] = None # give either a budget_id or max_budget + spend: Optional[float] = None + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model + + @model_validator(mode="before") + @classmethod + def check_user_info(cls, values): + if values.get("max_budget") is not None and values.get("budget_id") is not None: + raise ValueError("Set either 'max_budget' or 'budget_id', not both.") + + return values + + +class UpdateCustomerRequest(LiteLLMPydanticObjectBase): + """ + Update a Customer, use this to update customer budgets etc + + """ + + user_id: str + alias: Optional[str] = None # human-friendly alias + blocked: bool = False # allow/disallow requests for this end-user + max_budget: Optional[float] = None + budget_id: Optional[str] = None # give either a budget_id or max_budget + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model + + +class DeleteCustomerRequest(LiteLLMPydanticObjectBase): + """ + Delete multiple Customers + """ + + user_ids: List[str] + + +class MemberBase(LiteLLMPydanticObjectBase): + user_id: Optional[str] = Field( + default=None, + description="The unique ID of the user to add. Either user_id or user_email must be provided", + ) + user_email: Optional[str] = Field( + default=None, + description="The email address of the user to add. Either user_id or user_email must be provided", + ) + + @model_validator(mode="before") + @classmethod + def check_user_info(cls, values): + if not isinstance(values, dict): + raise ValueError("input needs to be a dictionary") + if values.get("user_id") is None and values.get("user_email") is None: + raise ValueError("Either user id or user email must be provided") + return values + + +class Member(MemberBase): + role: Literal[ + "admin", + "user", + ] = Field( + description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" + ) + + +class OrgMember(MemberBase): + role: Literal[ + LitellmUserRoles.ORG_ADMIN, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + + +class TeamBase(LiteLLMPydanticObjectBase): + team_alias: Optional[str] = None + team_id: Optional[str] = None + organization_id: Optional[str] = None + admins: list = [] + members: list = [] + members_with_roles: List[Member] = [] + team_member_permissions: Optional[List[str]] = None + metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + + # Budget fields + max_budget: Optional[float] = None + budget_duration: Optional[str] = None + + models: list = [] + blocked: bool = False + router_settings: Optional[dict] = None + + +class NewTeamRequest(TeamBase): + model_aliases: Optional[dict] = None + tags: Optional[list] = None + guardrails: Optional[List[str]] = None + policies: Optional[List[str]] = None + prompts: Optional[List[str]] = None + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + allowed_passthrough_routes: Optional[list] = None + secret_manager_settings: Optional[dict] = None + model_rpm_limit: Optional[Dict[str, int]] = None + rpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput"] + ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm + tpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput"] + ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm + + model_tpm_limit: Optional[Dict[str, int]] = None + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members + team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" + allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class GlobalEndUsersSpend(LiteLLMPydanticObjectBase): + api_key: Optional[str] = None + startTime: Optional[datetime] = None + endTime: Optional[datetime] = None + + +class UpdateTeamRequest(LiteLLMPydanticObjectBase): + """ + UpdateTeamRequest, used by /team/update when you need to update a team + + team_id: str + team_alias: Optional[str] = None + organization_id: Optional[str] = None + metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + models: Optional[list] = None + blocked: Optional[bool] = None + budget_duration: Optional[str] = None + guardrails: Optional[List[str]] = None + policies: Optional[List[str]] = None + """ + + team_id: str # required + team_alias: Optional[str] = None + organization_id: Optional[str] = None + metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + models: Optional[list] = None + blocked: Optional[bool] = None + budget_duration: Optional[str] = None + tags: Optional[list] = None + model_aliases: Optional[dict] = None + guardrails: Optional[List[str]] = None + policies: Optional[List[str]] = None + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + team_member_budget: Optional[float] = None + team_member_budget_duration: Optional[str] = None + team_member_rpm_limit: Optional[int] = None + team_member_tpm_limit: Optional[int] = None + team_member_key_duration: Optional[str] = None + allowed_passthrough_routes: Optional[list] = None + secret_manager_settings: Optional[dict] = None + prompts: Optional[List[str]] = None + model_rpm_limit: Optional[Dict[str, int]] = None + model_tpm_limit: Optional[Dict[str, int]] = None + allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None + router_settings: Optional[dict] = None + + +class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): + """ + internal type used to reset the budget on a team + used by reset_budget() + + team_id: str + spend: float + budget_reset_at: datetime + """ + + team_id: str + spend: float + budget_reset_at: datetime + updated_at: datetime + + +class DeleteTeamRequest(LiteLLMPydanticObjectBase): + team_ids: List[str] # required + + +class BlockTeamRequest(LiteLLMPydanticObjectBase): + team_id: str # required + + +class BlockKeyRequest(LiteLLMPydanticObjectBase): + key: str # required + + +class AddTeamCallback(LiteLLMPydanticObjectBase): + callback_name: str + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" + callback_vars: Dict[str, str] + + @model_validator(mode="before") + @classmethod + def validate_callback_vars(cls, values): + callback_vars = values.get("callback_vars", {}) + valid_keys = set(StandardCallbackDynamicParams.__annotations__.keys()) + for key, value in callback_vars.items(): + if key not in valid_keys: + raise ValueError( + f"Invalid callback variable: {key}. Must be one of {valid_keys}" + ) + if not isinstance(value, str): + callback_vars[key] = str(value) + return values + + +class TeamCallbackMetadata(LiteLLMPydanticObjectBase): + success_callback: Optional[List[str]] = [] + failure_callback: Optional[List[str]] = [] + callbacks: Optional[List[str]] = [] + # for now - only supported for langfuse + callback_vars: Optional[Dict[str, str]] = {} + + @model_validator(mode="before") + @classmethod + def validate_callback_vars(cls, values): + success_callback = values.get("success_callback", []) + if success_callback is None: + values.pop("success_callback", None) + failure_callback = values.get("failure_callback", []) + if failure_callback is None: + values.pop("failure_callback", None) + callbacks = values.get("callbacks", []) + if callbacks is None: + values.pop("callbacks", None) + + callback_vars = values.get("callback_vars", {}) + if callback_vars is None: + values.pop("callback_vars", None) + if all(val is None for val in values.values()): + return { + "success_callback": [], + "failure_callback": [], + "callbacks": [], + "callback_vars": {}, + } + valid_keys = set(StandardCallbackDynamicParams.__annotations__.keys()) + if callback_vars is not None: + for key in callback_vars: + if key not in valid_keys: + raise ValueError( + f"Invalid callback variable: {key}. Must be one of {valid_keys}" + ) + return values + + +class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_ObjectPermissionTable record""" + + object_permission_id: str + mcp_servers: Optional[List[str]] = [] + mcp_access_groups: Optional[List[str]] = [] + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None + """ + Mapping - server_id -> list of tools + + Enforces allowed tools for a specific key/team/organization + { + "1234567890": ["tool_name_1", "tool_name_2"] + } + """ + + vector_stores: Optional[List[str]] = [] + agents: Optional[List[str]] = [] + agent_access_groups: Optional[List[str]] = [] + + +class LiteLLM_TeamTable(TeamBase): + team_id: str # type: ignore + spend: Optional[float] = None + max_parallel_requests: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + model_id: Optional[int] = None + litellm_model_table: Optional[LiteLLM_ModelTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + updated_at: Optional[datetime] = None + created_at: Optional[datetime] = None + + ######################################################### + # Object Permission - MCP, Vector Stores etc. + ######################################################### + object_permission_id: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + dict_fields = [ + "metadata", + "aliases", + "config", + "permissions", + "model_max_budget", + "model_aliases", + "router_settings", + ] + + if isinstance(values, BaseModel): + values = values.model_dump() + + if ( + isinstance(values.get("members_with_roles"), dict) + and not values["members_with_roles"] + ): + values["members_with_roles"] = [] + + for field in dict_fields: + value = values.get(field) + if value is not None and isinstance(value, str): + try: + values[field] = json.loads(value) + except json.JSONDecodeError: + raise ValueError(f"Field {field} should be a valid dictionary") + + return values + + +class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): + last_refreshed_at: Optional[float] = None + + +class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): + """ + Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable + plus metadata captured at deletion time. + """ + + id: Optional[str] = None + deleted_at: Optional[datetime] = None + deleted_by: Optional[str] = None + deleted_by_api_key: Optional[str] = None + litellm_changed_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class TeamRequest(LiteLLMPydanticObjectBase): + teams: List[str] + + +class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_BudgetTable record""" + + budget_id: Optional[str] = None + soft_budget: Optional[float] = None + max_budget: Optional[float] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + model_max_budget: Optional[dict] = None + budget_duration: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): + """Represents all params for a LiteLLM_BudgetTable record""" + + budget_reset_at: Optional[datetime] = None + created_at: datetime + + +class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): + """ + Used to track spend of a user_id within a team_id + """ + + spend: Optional[float] = None + user_id: Optional[str] = None + team_id: Optional[str] = None + budget_id: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class NewOrganizationRequest(LiteLLM_BudgetTable): + organization_id: Optional[str] = None + organization_alias: str + models: List = [] + budget_id: Optional[str] = None + metadata: Optional[dict] = None + model_rpm_limit: Optional[Dict[str, int]] = None + model_tpm_limit: Optional[Dict[str, int]] = None + + ######################################################### + # Object Permission - MCP, Vector Stores etc. + ######################################################### + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + + +class OrganizationRequest(LiteLLMPydanticObjectBase): + organizations: List[str] + + +class DeleteOrganizationRequest(LiteLLMPydanticObjectBase): + organization_ids: List[str] # required + + +class TeamDefaultSettings(LiteLLMPydanticObjectBase): + team_id: str + + model_config = ConfigDict( + extra="allow" + ) # allow params not defined here, these fall in litellm.completion(**kwargs) + + +class DynamoDBArgs(LiteLLMPydanticObjectBase): + billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"] + read_capacity_units: Optional[int] = None + write_capacity_units: Optional[int] = None + ssl_verify: Optional[bool] = None + region_name: str + user_table_name: str = "LiteLLM_UserTable" + key_table_name: str = "LiteLLM_VerificationToken" + config_table_name: str = "LiteLLM_Config" + spend_table_name: str = "LiteLLM_SpendLogs" + aws_role_name: Optional[str] = None + aws_session_name: Optional[str] = None + aws_web_identity_token: Optional[str] = None + aws_provider_id: Optional[str] = None + aws_policy_arns: Optional[List[str]] = None + aws_policy: Optional[str] = None + aws_duration_seconds: Optional[int] = None + assume_role_aws_role_name: Optional[str] = None + assume_role_aws_session_name: Optional[str] = None + + +class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): + """ + Settings for a specific guardrail on a passthrough endpoint. + + Allows field-level targeting for guardrail execution. + """ + + request_fields: Optional[List[str]] = Field( + default=None, + description="JSONPath expressions for input field targeting (pre_call). Examples: 'query', 'documents[*].text', 'messages[*].content'. If not specified, guardrail runs on entire request payload.", + ) + response_fields: Optional[List[str]] = Field( + default=None, + description="JSONPath expressions for output field targeting (post_call). Examples: 'results[*].text', 'output'. If not specified, guardrail runs on entire response payload.", + ) + + +# Type alias for the guardrails dict: guardrail_name -> settings (or None for defaults) +PassThroughGuardrailsConfig = Dict[str, Optional[PassThroughGuardrailSettings]] + + +class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): + id: Optional[str] = Field( + default=None, + description="Optional unique identifier for the pass-through endpoint. If not provided, endpoints will be identified by path for backwards compatibility.", + ) + path: str = Field(description="The route to be added to the LiteLLM Proxy Server.") + target: str = Field( + description="The URL to which requests for this path should be forwarded." + ) + headers: dict = Field( + default={}, + description="Key-value pairs of headers to be forwarded with the request. You can set any key value pair here and it will be forwarded to your target endpoint", + ) + include_subpath: bool = Field( + default=False, + description="If True, requests to subpaths of the path will be forwarded to the target endpoint. For example, if the path is /bria and include_subpath is True, requests to /bria/v1/text-to-image/base/2.3 will be forwarded to the target endpoint.", + ) + cost_per_request: float = Field( + default=0.0, + description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.", + ) + auth: bool = Field( + default=False, + description="Whether authentication is required for the pass-through endpoint. If True, requests to the endpoint will require a valid LiteLLM API key.", + ) + guardrails: Optional[PassThroughGuardrailsConfig] = Field( + default=None, + description="Guardrails configuration for this passthrough endpoint. Dict keys are guardrail names, values are optional settings for field targeting. When set, all org/team/key level guardrails will also execute. Defaults to None (no guardrails execute).", + ) + + +class PassThroughEndpointResponse(LiteLLMPydanticObjectBase): + endpoints: List[PassThroughGenericEndpoint] + + +class ConfigFieldUpdate(LiteLLMPydanticObjectBase): + field_name: str + field_value: Any + config_type: Literal["general_settings"] + + +class ConfigFieldDelete(LiteLLMPydanticObjectBase): + config_type: Literal["general_settings"] + field_name: str + + +class CallbackDelete(LiteLLMPydanticObjectBase): + callback_name: str + + +class FieldDetail(BaseModel): + field_name: str + field_type: str + field_description: str + field_default_value: Any = None + stored_in_db: Optional[bool] + + +class ConfigList(LiteLLMPydanticObjectBase): + field_name: str + field_type: str + field_description: str + field_value: Any + stored_in_db: Optional[bool] + field_default_value: Any + premium_field: bool = False + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields + + +class UserHeaderMapping(LiteLLMPydanticObjectBase): + """ + Map an incoming HTTP header to a LiteLLM user role. + """ + + header_name: str + litellm_user_role: Literal[ + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.CUSTOMER, + ] + + model_config = { + "extra": "forbid", + } + + +UserMCPManagementMode = Literal["restricted", "view_all"] + + +class ConfigGeneralSettings(LiteLLMPydanticObjectBase): + """ + Documents all the fields supported by `general_settings` in config.yaml + """ + + completion_model: Optional[str] = Field( + None, description="proxy level default model for all chat completion calls" + ) + key_management_system: Optional[KeyManagementSystem] = Field( + None, description="key manager to load keys from / decrypt keys with" + ) + use_google_kms: Optional[bool] = Field( + None, description="decrypt keys with google kms" + ) + use_azure_key_vault: Optional[bool] = Field( + None, description="load keys from azure key vault" + ) + master_key: Optional[str] = Field( + None, description="require a key for all calls to proxy" + ) + database_url: Optional[str] = Field( + None, + description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", + ) + database_connection_pool_limit: Optional[int] = Field( + 10, + description="default connection pool for prisma client connecting to postgres db", + ) + database_connection_timeout: Optional[float] = Field( + 60, description="default timeout for a connection to the database" + ) + database_type: Optional[Literal["dynamo_db"]] = Field( + None, description="to use dynamodb instead of postgres db" + ) + database_args: Optional[DynamoDBArgs] = Field( + None, + description="custom args for instantiating dynamodb client - e.g. billing provision", + ) + otel: Optional[bool] = Field( + None, + description="[BETA] OpenTelemetry support - this might change, use with caution.", + ) + custom_auth: Optional[str] = Field( + None, + description="override user_api_key_auth with your own auth script - https://docs.litellm.ai/docs/proxy/virtual_keys#custom-auth", + ) + max_parallel_requests: Optional[int] = Field( + None, + description="maximum parallel requests for each api key", + ) + global_max_parallel_requests: Optional[int] = Field( + None, description="global max parallel requests to allow for a proxy instance." + ) + max_request_size_mb: Optional[int] = Field( + None, + description="max request size in MB, if a request is larger than this size it will be rejected", + ) + max_response_size_mb: Optional[int] = Field( + None, + description="max response size in MB, if a response is larger than this size it will be rejected", + ) + infer_model_from_keys: Optional[bool] = Field( + None, + description="for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY)", + ) + background_health_checks: Optional[bool] = Field( + None, description="run health checks in background" + ) + health_check_interval: int = Field( + 300, description="background health check interval in seconds" + ) + alerting: Optional[List] = Field( + None, + description="List of alerting integrations. Today, just slack - `alerting: ['slack']`", + ) + alert_types: Optional[List[AlertType]] = Field( + None, + description="List of alerting types. By default it is all alerts", + ) + alert_to_webhook_url: Optional[Dict] = Field( + None, + description="Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`", + ) + alerting_args: Optional[Dict] = Field( + None, description="Controllable params for slack alerting - e.g. ttl in cache." + ) + alerting_threshold: Optional[int] = Field( + None, + description="sends alerts if requests hang for 5min+", + ) + ui_access_mode: Optional[Literal["admin_only", "all"]] = Field( + "all", description="Control access to the Proxy UI" + ) + allowed_routes: Optional[List] = Field( + None, description="Proxy API Endpoints you want users to be able to access" + ) + reject_clientside_metadata_tags: Optional[bool] = Field( + None, + description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", + ) + enable_public_model_hub: bool = Field( + default=False, + description="Public model hub for users to see what models they have access to, supported openai params, etc.", + ) + pass_through_endpoints: Optional[List[PassThroughGenericEndpoint]] = Field( + default=None, + description="Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through", + ) + user_header_name: Optional[str] = Field( + None, + description="[DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings.", + ) + user_header_mappings: Optional[List[UserHeaderMapping]] = None + supported_db_objects: Optional[List[SupportedDBObjectType]] = Field( + None, + description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map'. If not set, all objects are loaded (default behavior).", + ) + user_mcp_management_mode: Optional[UserMCPManagementMode] = Field( + None, + description="Controls how non-admin users interact with MCP servers in the dashboard. 'restricted' shows only accessible servers, 'view_all' lists every server in read-only mode.", + ) + store_prompts_in_spend_logs: Optional[bool] = Field( + None, + description="If True, stores request messages and responses in spend logs. Default is False.", + ) + maximum_spend_logs_retention_period: Optional[str] = Field( + None, + description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", + ) + + +class ConfigYAML(LiteLLMPydanticObjectBase): + """ + Documents all the fields supported by the config.yaml + """ + + environment_variables: Optional[dict] = Field( + None, + description="Object to pass in additional environment variables via POST request", + ) + model_list: Optional[List[ModelParams]] = Field( + None, + description="List of supported models on the server, with model-specific configs", + ) + litellm_settings: Optional[dict] = Field( + None, + description="litellm Module settings. See __init__.py for all, example litellm.drop_params=True, litellm.set_verbose=True, litellm.api_base, litellm.cache", + ) + general_settings: Optional[ConfigGeneralSettings] = None + router_settings: Optional[UpdateRouterConfig] = Field( + None, + description="litellm router object settings. See router.py __init__ for all, example router.num_retries=5, router.timeout=5, router.max_retries=5, router.retry_after=5", + ) + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): + token: Optional[str] = None + key_name: Optional[str] = None + key_alias: Optional[str] = None + spend: float = 0.0 + max_budget: Optional[float] = None + expires: Optional[Union[str, datetime]] = None + models: List = [] + aliases: Dict = {} + config: Dict = {} + user_id: Optional[str] = None + team_id: Optional[str] = None + max_parallel_requests: Optional[int] = None + metadata: Dict = {} + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + allowed_cache_controls: Optional[list] = [] + allowed_routes: Optional[list] = [] + permissions: Dict = {} + model_spend: Dict = {} + model_max_budget: Dict = {} + soft_budget_cooldown: bool = False + blocked: Optional[bool] = None + litellm_budget_table: Optional[dict] = None + org_id: Optional[str] = None # org id for a given key + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + object_permission_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + rotation_count: Optional[int] = 0 # Number of times key has been rotated + auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated + rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") + last_rotation_at: Optional[datetime] = None # When this key was last rotated + key_rotation_at: Optional[datetime] = None # When this key should next be rotated + router_settings: Optional[ + Dict + ] = None # Router settings for this key (Key > Team > Global precedence) + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): + """ + Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken + plus metadata captured at deletion time. + """ + + id: Optional[str] = None + deleted_at: Optional[datetime] = None + deleted_by: Optional[str] = None + deleted_by_api_key: Optional[str] = None + litellm_changed_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): + """ + Combined view of litellm verification token + litellm team table (select values) + """ + + team_spend: Optional[float] = None + team_alias: Optional[str] = None + team_tpm_limit: Optional[int] = None + team_rpm_limit: Optional[int] = None + team_max_budget: Optional[float] = None + team_models: List = [] + team_blocked: bool = False + soft_budget: Optional[float] = None + team_model_aliases: Optional[Dict] = None + team_member: Optional[Member] = None + team_metadata: Optional[Dict] = None + team_object_permission_id: Optional[str] = None + + # Team Member Specific Params + team_member_spend: Optional[float] = None + team_member_tpm_limit: Optional[int] = None + team_member_rpm_limit: Optional[int] = None + + # End User Params + end_user_id: Optional[str] = None + end_user_tpm_limit: Optional[int] = None + end_user_rpm_limit: Optional[int] = None + end_user_max_budget: Optional[float] = None + + # Organization Params + organization_max_budget: Optional[float] = None + organization_tpm_limit: Optional[int] = None + organization_rpm_limit: Optional[int] = None + organization_metadata: Optional[dict] = None + + # Time stamps + last_refreshed_at: Optional[float] = None # last time joint view was pulled from db + + def __init__(self, **kwargs): + # Handle litellm_budget_table_* keys (budget table overrides when key value is None or empty) + for key, value in list(kwargs.items()): + if key.startswith("litellm_budget_table_") and value is not None: + # Extract the corresponding attribute name + attr_name = key.replace("litellm_budget_table_", "") + # Use key's value from kwargs (from DB view), not class default + current = kwargs.get(attr_name) + if current is None: + current = getattr(self, attr_name, None) + # Apply budget value when key has no value, or for model_max_budget when key has empty dict + should_apply = current is None or ( + attr_name == "model_max_budget" + and isinstance(current, dict) + and len(current) == 0 + ) + if should_apply: + kwargs[attr_name] = value + if key == "end_user_id" and value is not None and isinstance(value, int): + kwargs[key] = str(value) + + if kwargs.get("organization_id") is not None: + kwargs["org_id"] = kwargs.pop("organization_id") + # Initialize the superclass + super().__init__(**kwargs) + + +class UserAPIKeyAuth( + LiteLLM_VerificationTokenView +): # the expected response object for user api key auth + """ + Return the row in the db + """ + + api_key: Optional[str] = None + user_role: Optional[LitellmUserRoles] = None + allowed_model_region: Optional[AllowedModelRegion] = None + parent_otel_span: Optional[Span] = None + rpm_limit_per_model: Optional[Dict[str, int]] = None + tpm_limit_per_model: Optional[Dict[str, int]] = None + user_tpm_limit: Optional[int] = None + user_rpm_limit: Optional[int] = None + user_email: Optional[str] = None + user_spend: Optional[float] = None + user_max_budget: Optional[float] = None + request_route: Optional[str] = None + user: Optional[Any] = None # Expanded user object when expand=user is used + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @model_validator(mode="before") + @classmethod + def check_api_key(cls, values): + # If values is already an instance (not a dict), return it as-is + if not isinstance(values, dict): + return values + if values.get("api_key") is not None: + values.update( + {"token": cls._safe_hash_litellm_api_key(values.get("api_key"))} + ) + if isinstance(values.get("api_key"), str): + values.update( + {"api_key": cls._safe_hash_litellm_api_key(values.get("api_key"))} + ) + return values + + @classmethod + def _safe_hash_litellm_api_key(cls, api_key: str) -> str: + """ + Helper to ensure all logged keys are hashed + Covers: + 1. Regular API keys from LiteLLM DB + 2. JWT tokens used for connecting to LiteLLM API + """ + if api_key.startswith("sk-"): + return hash_token(api_key) + from litellm.proxy.auth.handle_jwt import JWTHandler + + if JWTHandler.is_jwt(token=api_key): + return f"hashed-jwt-{hash_token(token=api_key)}" + return api_key + + @classmethod + def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth": + """ + Returns a `UserAPIKeyAuth` object for the litellm internal health check service account. + + This is used to track number of requests/spend for health check calls. + """ + from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + + return cls( + api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + team_id=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + ) + + @classmethod + def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth": + """ + Returns a `UserAPIKeyAuth` object for the litellm internal health check service account. + + This is used to track number of requests/spend for health check calls. + """ + from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME + + return cls( + api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + ) + + @classmethod + def get_litellm_internal_jobs_user_api_key_auth(cls) -> "UserAPIKeyAuth": + """ + Returns a `UserAPIKeyAuth` object for internal LiteLLM jobs like key rotation. + + This is used to track actions performed by automated system jobs. + """ + from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + + return cls( + api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + team_id="system", + key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + team_alias="system", + user_id="system", + ) + + +class UserInfoResponse(LiteLLMPydanticObjectBase): + user_id: Optional[str] + user_info: Optional[Union[dict, BaseModel]] + keys: List + teams: List + + +class LiteLLM_Config(LiteLLMPydanticObjectBase): + param_name: str + param_value: Dict + + +class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): + """ + This is the table that track what organizations a user belongs to and users spend within the organization + """ + + user_id: str + organization_id: str + user_role: Optional[str] = None + spend: float = 0.0 + budget_id: Optional[str] = None + created_at: datetime + updated_at: datetime + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable): + """Represents user-controllable params for a LiteLLM_OrganizationTable record""" + + organization_id: Optional[str] = None + organization_alias: Optional[str] = None + budget_id: Optional[str] = None + spend: Optional[float] = None + metadata: Optional[dict] = None + models: Optional[List[str]] = None + updated_by: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + model_tpm_limit: Optional[Dict[str, int]] = None + model_rpm_limit: Optional[Dict[str, int]] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if values.get(field) is not None: + # add to metadata + if values.get("metadata") is None: + values.update({"metadata": {}}) + values["metadata"][field] = values.get(field) + values.pop(field) + return values + + +class LiteLLM_UserTable(LiteLLMPydanticObjectBase): + user_id: str + max_budget: Optional[float] = None + spend: float = 0.0 + model_max_budget: Optional[Dict] = {} + model_spend: Optional[Dict] = {} + user_email: Optional[str] = None + user_alias: Optional[str] = None + models: list = [] + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + user_role: Optional[str] = None + organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None + teams: List[str] = [] + sso_user_id: Optional[str] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + metadata: Optional[dict] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + if values.get("models") is None: + values.update({"models": []}) + if values.get("teams") is None: + values.update({"teams": []}) + return values + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_OrganizationTable record""" + + organization_id: Optional[str] = None + organization_alias: Optional[str] = None + budget_id: str + spend: float = 0.0 + metadata: Optional[dict] = None + models: List[str] + created_by: str + updated_by: str + users: Optional[List[LiteLLM_UserTable]] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + + ######################################################### + # Object Permission - MCP, Vector Stores etc. + ######################################################### + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + object_permission_id: Optional[str] = None + + +class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable): + """Returned by the /organization/info endpoint and /organization/list endpoint""" + + members: List[LiteLLM_OrganizationMembershipTable] = [] + teams: List[LiteLLM_TeamTable] = [] + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + created_at: datetime + updated_at: datetime + + +class NewOrganizationResponse(LiteLLM_OrganizationTable): + organization_id: str # type: ignore + created_at: datetime + updated_at: datetime + + +class LiteLLM_UserTableFiltered(BaseModel): # done to avoid exposing sensitive data + user_id: str + user_email: Optional[str] = None + + +class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable): + key_count: int = 0 + + +class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): + user_id: str + blocked: bool + alias: Optional[str] = None + spend: float = 0.0 + allowed_model_region: Optional[AllowedModelRegion] = None + default_model: Optional[str] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + return values + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_TagTable(LiteLLMPydanticObjectBase): + tag_name: str + description: Optional[str] = None + models: List[str] = [] + model_info: Optional[dict] = None + spend: float = 0.0 + budget_id: Optional[str] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + if values.get("models") is None: + values.update({"models": []}) + return values + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_SpendLogs(LiteLLMPydanticObjectBase): + request_id: str + api_key: str + model: Optional[str] = "" + api_base: Optional[str] = "" + call_type: str + spend: Optional[float] = 0.0 + total_tokens: Optional[int] = 0 + prompt_tokens: Optional[int] = 0 + completion_tokens: Optional[int] = 0 + startTime: Union[str, datetime, None] + endTime: Union[str, datetime, None] + user: Optional[str] = "" + metadata: Optional[Json] = {} + cache_hit: Optional[str] = "False" + cache_key: Optional[str] = None + request_tags: Optional[Json] = None + requester_ip_address: Optional[str] = None + messages: Optional[Union[str, list, dict]] + response: Optional[Union[str, list, dict]] + + +class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): + request_id: Optional[str] = str(uuid.uuid4()) + api_base: Optional[str] = "" + model_group: Optional[str] = "" + litellm_model_name: Optional[str] = "" + model_id: Optional[str] = "" + request_kwargs: Optional[dict] = {} + exception_type: Optional[str] = "" + status_code: Optional[str] = "" + exception_string: Optional[str] = "" + startTime: Union[str, datetime, None] + endTime: Union[str, datetime, None] + + +AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "rotated"] + + +class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): + id: str + updated_at: datetime + changed_by: Optional[Any] = None + changed_by_api_key: Optional[str] = None + action: AUDIT_ACTIONS + table_name: LitellmTableNames + object_id: str + before_value: Optional[Json] = None + updated_values: Optional[Json] = None + + @model_validator(mode="before") + @classmethod + def cast_changed_by_to_str(cls, values): + if values.get("changed_by") is not None: + values["changed_by"] = str(values["changed_by"]) + return values + + @model_validator(mode="after") + def mask_api_keys(self): + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + + masker = SensitiveDataMasker(sensitive_patterns={"key"}) + + if self.before_value is not None: + json_before_value: Optional[dict] = None + if isinstance(self.before_value, str): + json_before_value = json.loads(self.before_value) + elif isinstance(self.before_value, dict): + json_before_value = self.before_value + + if json_before_value is not None: + json_before_value = masker.mask_dict(json_before_value) + self.before_value = json.dumps(json_before_value, default=str) + + if self.updated_values is not None: + json_updated_values: Optional[dict] = None + if isinstance(self.updated_values, str): + json_updated_values = json.loads(self.updated_values) + elif isinstance(self.updated_values, dict): + json_updated_values = self.updated_values + + if json_updated_values is not None: + json_updated_values = masker.mask_dict(json_updated_values) + self.updated_values = json.dumps(json_updated_values, default=str) + + return self + + +class LiteLLM_SpendLogs_ResponseObject(LiteLLMPydanticObjectBase): + response: Optional[List[Union[LiteLLM_SpendLogs, Any]]] = None + + +class TokenCountRequest(LiteLLMPydanticObjectBase): + model: str + prompt: Optional[str] = None + messages: Optional[List[dict]] = None + """ + Anthropic token counting endpoint uses /messages + """ + + contents: Optional[List[dict]] = None + """ + Google /countTokens endpoint expects contents to be a list of dicts with the following structure: + """ + + +class CallInfo(LiteLLMPydanticObjectBase): + """Used for slack budget alerting""" + + spend: float + max_budget: Optional[float] = None + soft_budget: Optional[float] = None + token: Optional[str] = Field(default=None, description="Hashed value of that key") + customer_id: Optional[str] = None + user_id: Optional[str] = None + team_id: Optional[str] = None + team_alias: Optional[str] = None + organization_id: Optional[str] = None + user_email: Optional[str] = None + key_alias: Optional[str] = None + projected_exceeded_date: Optional[str] = None + projected_spend: Optional[float] = None + event_group: Litellm_EntityType + + +class WebhookEvent(CallInfo): + event: Literal[ + "budget_crossed", + "max_budget_alert", + "soft_budget_crossed", + "threshold_crossed", + "projected_limit_exceeded", + "key_created", + "key_rotated", + "internal_user_created", + "spend_tracked", + ] + event_message: str # human-readable description of event + event_group: Litellm_EntityType + + +class SpecialModelNames(enum.Enum): + all_team_models = "all-team-models" + all_proxy_models = "all-proxy-models" + no_default_models = "no-default-models" + + +class SpecialProxyStrings(enum.Enum): + default_user_id = "default_user_id" # global proxy admin + + +class InvitationNew(LiteLLMPydanticObjectBase): + user_id: str + + +class InvitationUpdate(LiteLLMPydanticObjectBase): + invitation_id: str + is_accepted: bool + + +class InvitationDelete(LiteLLMPydanticObjectBase): + invitation_id: str + + +class InvitationModel(LiteLLMPydanticObjectBase): + id: str + user_id: str + is_accepted: bool + accepted_at: Optional[datetime] + expires_at: datetime + created_at: datetime + created_by: str + updated_at: datetime + updated_by: str + + +class InvitationClaim(LiteLLMPydanticObjectBase): + invitation_link: str + user_id: str + password: str + + +class ConfigFieldInfo(LiteLLMPydanticObjectBase): + field_name: str + field_value: Any + + +class CallbackOnUI(LiteLLMPydanticObjectBase): + litellm_callback_name: str + litellm_callback_params: Optional[list] + ui_callback_name: str + + +class AllCallbacks(LiteLLMPydanticObjectBase): + langfuse: CallbackOnUI = CallbackOnUI( + litellm_callback_name="langfuse", + ui_callback_name="Langfuse", + litellm_callback_params=[ + "LANGFUSE_PUBLIC_KEY", + "LANGFUSE_SECRET_KEY", + "LANGFUSE_HOST", + ], + ) + + otel: CallbackOnUI = CallbackOnUI( + litellm_callback_name="otel", + ui_callback_name="OpenTelemetry", + litellm_callback_params=[ + "OTEL_EXPORTER", + "OTEL_ENDPOINT", + "OTEL_HEADERS", + ], + ) + + s3: CallbackOnUI = CallbackOnUI( + litellm_callback_name="s3", + ui_callback_name="s3 Bucket (AWS)", + litellm_callback_params=[ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION_NAME", + ], + ) + + openmeter: CallbackOnUI = CallbackOnUI( + litellm_callback_name="openmeter", + ui_callback_name="OpenMeter", + litellm_callback_params=[ + "OPENMETER_API_ENDPOINT", + "OPENMETER_API_KEY", + ], + ) + + custom_callback_api: CallbackOnUI = CallbackOnUI( + litellm_callback_name="custom_callback_api", + litellm_callback_params=["GENERIC_LOGGER_ENDPOINT", "GENERIC_LOGGER_HEADERS"], + ui_callback_name="Custom Callback API", + ) + + generic_api: CallbackOnUI = CallbackOnUI( + litellm_callback_name="generic_api", + litellm_callback_params=["GENERIC_LOGGER_ENDPOINT", "GENERIC_LOGGER_HEADERS"], + ui_callback_name="Custom Callback API", + ) + + datadog: CallbackOnUI = CallbackOnUI( + litellm_callback_name="datadog", + litellm_callback_params=["DD_API_KEY", "DD_SITE"], + ui_callback_name="Datadog", + ) + + braintrust: CallbackOnUI = CallbackOnUI( + litellm_callback_name="braintrust", + litellm_callback_params=["BRAINTRUST_API_KEY", "BRAINTRUST_API_BASE"], + ui_callback_name="Braintrust", + ) + + langsmith: CallbackOnUI = CallbackOnUI( + litellm_callback_name="langsmith", + litellm_callback_params=[ + "LANGSMITH_API_KEY", + "LANGSMITH_PROJECT", + "LANGSMITH_DEFAULT_RUN_NAME", + ], + ui_callback_name="Langsmith", + ) + + lago: CallbackOnUI = CallbackOnUI( + litellm_callback_name="lago", + litellm_callback_params=[ + "LAGO_API_BASE", + "LAGO_API_KEY", + "LAGO_API_EVENT_CODE", + "LAGO_API_CHARGE_BY", + ], + ui_callback_name="Lago Billing", + ) + + traceloop: CallbackOnUI = CallbackOnUI( + litellm_callback_name="traceloop", + litellm_callback_params=[ + "TRACELOOP_API_KEY", + ], + ui_callback_name="Traceloop", + ) + + +class SpendLogsMetadata(TypedDict): + """ + Specific metadata k,v pairs logged to spendlogs for easier cost tracking + """ + + additional_usage_values: Optional[ + dict + ] # covers provider-specific usage information - e.g. prompt caching + user_api_key: Optional[str] + user_api_key_alias: Optional[str] + user_api_key_team_id: Optional[str] + user_api_key_org_id: Optional[str] + user_api_key_user_id: Optional[str] + user_api_key_team_alias: Optional[str] + spend_logs_metadata: Optional[ + dict + ] # special param to log k,v pairs to spendlogs for a call + requester_ip_address: Optional[str] + applied_guardrails: Optional[List[str]] + mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] + vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] + guardrail_information: Optional[List[StandardLoggingGuardrailInformation]] + status: StandardLoggingPayloadStatus + proxy_server_request: Optional[str] + batch_models: Optional[List[str]] + error_information: Optional[StandardLoggingPayloadErrorInformation] + usage_object: Optional[dict] + model_map_information: Optional[StandardLoggingModelInformation] + cold_storage_object_key: Optional[ + str + ] # S3/GCS object key for cold storage retrieval + litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds + cost_breakdown: Optional[ + CostBreakdown + ] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) + + +class SpendLogsPayload(TypedDict): + request_id: str + call_type: str + api_key: str + spend: float + total_tokens: int + prompt_tokens: int + completion_tokens: int + startTime: Union[datetime, str] + endTime: Union[datetime, str] + completionStartTime: Optional[Union[datetime, str]] + model: str + model_id: Optional[str] + model_group: Optional[str] + mcp_namespaced_tool_name: Optional[str] + agent_id: Optional[str] + api_base: str + user: str + metadata: str # json str + cache_hit: str + cache_key: str + request_tags: str # json str + team_id: Optional[str] + organization_id: Optional[str] + end_user: Optional[str] + requester_ip_address: Optional[str] + custom_llm_provider: Optional[str] + messages: Optional[Union[str, list, dict]] + response: Optional[Union[str, list, dict]] + proxy_server_request: Optional[str] + session_id: Optional[str] + status: Literal["success", "failure"] + + +class SpanAttributes(str, enum.Enum): + # Note: We've taken this from opentelemetry-semantic-conventions-ai + # I chose to not add a new dependency to litellm for this + + # Semantic Conventions for LLM requests, this needs to be removed after + # OpenTelemetry Semantic Conventions support Gen AI. + # Issue at https://github.com/open-telemetry/opentelemetry-python/issues/3868 + # Refer to https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/llm-spans.md + + LLM_SYSTEM = "gen_ai.system" + LLM_REQUEST_MODEL = "gen_ai.request.model" + LLM_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens" + LLM_REQUEST_TEMPERATURE = "gen_ai.request.temperature" + LLM_REQUEST_TOP_P = "gen_ai.request.top_p" + LLM_PROMPTS = "gen_ai.prompt" + LLM_COMPLETIONS = "gen_ai.completion" + LLM_RESPONSE_MODEL = "gen_ai.response.model" + LLM_USAGE_COMPLETION_TOKENS = "gen_ai.usage.completion_tokens" + LLM_USAGE_PROMPT_TOKENS = "gen_ai.usage.prompt_tokens" + + # OTEL 1.38 attributes + GEN_AI_INPUT_MESSAGES = "gen_ai.input.messages" + GEN_AI_OUTPUT_MESSAGES = "gen_ai.output.messages" + GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens" + GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens" + GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens" + GEN_AI_OPERATION_NAME = "gen_ai.operation.name" + GEN_AI_REQUEST_ID = "gen_ai.request.id" + GEN_AI_SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions" + GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons" + + LLM_TOKEN_TYPE = "gen_ai.token.type" + # To be added + # LLM_RESPONSE_FINISH_REASON = "gen_ai.response.finish_reasons" + # LLM_RESPONSE_ID = "gen_ai.response.id" + + # LLM + LLM_REQUEST_TYPE = "llm.request.type" + LLM_USAGE_TOTAL_TOKENS = "llm.usage.total_tokens" + LLM_USAGE_TOKEN_TYPE = "llm.usage.token_type" + LLM_USER = "llm.user" + LLM_HEADERS = "llm.headers" + LLM_TOP_K = "llm.top_k" + LLM_IS_STREAMING = "llm.is_streaming" + LLM_FREQUENCY_PENALTY = "llm.frequency_penalty" + LLM_PRESENCE_PENALTY = "llm.presence_penalty" + LLM_CHAT_STOP_SEQUENCES = "llm.chat.stop_sequences" + LLM_REQUEST_FUNCTIONS = "llm.request.functions" + LLM_REQUEST_REPETITION_PENALTY = "llm.request.repetition_penalty" + LLM_RESPONSE_FINISH_REASON = "llm.response.finish_reason" + LLM_RESPONSE_STOP_REASON = "llm.response.stop_reason" + LLM_CONTENT_COMPLETION_CHUNK = "llm.content.completion.chunk" + + # OpenAI + LLM_OPENAI_RESPONSE_SYSTEM_FINGERPRINT = "gen_ai.openai.system_fingerprint" + LLM_OPENAI_API_BASE = "gen_ai.openai.api_base" + LLM_OPENAI_API_VERSION = "gen_ai.openai.api_version" + LLM_OPENAI_API_TYPE = "gen_ai.openai.api_type" + + +class ManagementEndpointLoggingPayload(LiteLLMPydanticObjectBase): + route: str + request_data: dict + response: Optional[dict] = None + exception: Optional[Any] = None + start_time: Optional[datetime] = None + end_time: Optional[datetime] = None + + +class ProxyException(Exception): + # NOTE: DO NOT MODIFY THIS + # This is used to map exactly to OPENAI Exceptions + def __init__( + self, + message: str, + type: str, + param: Optional[str], + code: Optional[Union[int, str]] = None, # maps to status code + headers: Optional[Dict[str, str]] = None, + openai_code: Optional[str] = None, # maps to 'code' in openai + provider_specific_fields: Optional[dict] = None, + ): + self.message = str(message) + self.type = type + self.param = param + self.openai_code = openai_code or code + # If we look on official python OpenAI lib, the code should be a string: + # https://github.com/openai/openai-python/blob/195c05a64d39c87b2dfdf1eca2d339597f1fce03/src/openai/types/shared/error_object.py#L11 + # Related LiteLLM issue: https://github.com/BerriAI/litellm/discussions/4834 + self.code = str(code) + if headers is not None: + for k, v in headers.items(): + if not isinstance(v, str): + headers[k] = str(v) + self.headers = headers or {} + self.provider_specific_fields = provider_specific_fields + # rules for proxyExceptions + # Litellm router.py returns "No healthy deployment available" when there are no deployments available + # Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487 + if ( + "No healthy deployment available" in self.message + or "No deployments available" in self.message + ): + self.code = "429" + elif RouterErrors.no_deployments_with_tag_routing.value in self.message: + self.code = "401" + + def to_dict(self) -> dict: + """Converts the ProxyException instance to a dictionary.""" + error_dict: Dict[str, Optional[Union[str, Dict]]] = { + "message": self.message, + "type": self.type, + "param": self.param, + "code": self.code, + } + if self.provider_specific_fields: + error_dict["provider_specific_fields"] = self.provider_specific_fields + return error_dict + + +class CommonProxyErrors(str, enum.Enum): + db_not_connected_error = ( + "DB not connected. See https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + no_llm_router = "No models configured on proxy" + not_allowed_access = "Admin-only endpoint. Not allowed to access this." + not_premium_user = "You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/enterprise#trial. \nPricing: https://www.litellm.ai/#pricing" + max_parallel_request_limit_reached = ( + "Crossed TPM / RPM / Max Parallel Request Limit" + ) + missing_enterprise_package = "Missing litellm-enterprise package. Please install it to use this feature. Run `pip install litellm-enterprise`" + missing_enterprise_package_docker = ( + "This uses the enterprise folder - only available on the Docker image." + ) + + +class SpendCalculateRequest(LiteLLMPydanticObjectBase): + model: Optional[str] = None + messages: Optional[List] = None + completion_response: Optional[dict] = None + + +class ProxyErrorTypes(str, enum.Enum): + budget_exceeded = "budget_exceeded" + """ + Object was over budget + """ + no_db_connection = "no_db_connection" + """ + No database connection + """ + + token_not_found_in_db = "token_not_found_in_db" + """ + Requested token was not found in the database + """ + + key_model_access_denied = "key_model_access_denied" + """ + Key does not have access to the model + """ + + team_model_access_denied = "team_model_access_denied" + """ + Team does not have access to the model + """ + + user_model_access_denied = "user_model_access_denied" + """ + User does not have access to the model + """ + + org_model_access_denied = "org_model_access_denied" + """ + Organization does not have access to the model + """ + + expired_key = "expired_key" + """ + Key has expired + """ + + auth_error = "auth_error" + """ + General authentication error + """ + + internal_server_error = "internal_server_error" + """ + Internal server error + """ + + bad_request_error = "bad_request_error" + """ + Bad request error + """ + + not_found_error = "not_found_error" + """ + Not found error + """ + + validation_error = "validation_error" + """ + Validation error + """ + + cache_ping_error = "cache_ping_error" + """ + Cache ping error + """ + + team_member_permission_error = "team_member_permission_error" + """ + Team member permission error + """ + + key_vector_store_access_denied = "key_vector_store_access_denied" + """ + Key does not have access to the vector store + """ + + team_vector_store_access_denied = "team_vector_store_access_denied" + """ + Team does not have access to the vector store + """ + + org_vector_store_access_denied = "org_vector_store_access_denied" + """ + Organization does not have access to the vector store + """ + + team_member_already_in_team = "team_member_already_in_team" + """ + Team member is already in team + """ + + @classmethod + def get_model_access_error_type_for_object( + cls, object_type: Literal["key", "user", "team", "org"] + ) -> "ProxyErrorTypes": + """ + Get the model access error type for object_type + """ + if object_type == "key": + return cls.key_model_access_denied + elif object_type == "team": + return cls.team_model_access_denied + elif object_type == "user": + return cls.user_model_access_denied + elif object_type == "org": + return cls.org_model_access_denied + + @classmethod + def get_vector_store_access_error_type_for_object( + cls, object_type: Literal["key", "team", "org"] + ) -> "ProxyErrorTypes": + """ + Get the vector store access error type for object_type + """ + if object_type == "key": + return cls.key_vector_store_access_denied + elif object_type == "team": + return cls.team_vector_store_access_denied + elif object_type == "org": + return cls.org_vector_store_access_denied + + +DB_CONNECTION_ERROR_TYPES = ( + httpx.ConnectError, + httpx.ReadError, + httpx.ReadTimeout, +) + + +class SSOUserDefinedValues(TypedDict): + models: List[str] + user_id: str + user_email: Optional[str] + user_role: Optional[str] + max_budget: Optional[float] + budget_duration: Optional[str] + + +class VirtualKeyEvent(LiteLLMPydanticObjectBase): + created_by_user_id: str + created_by_user_role: str + created_by_key_alias: Optional[str] + request_kwargs: dict + + +class CreatePassThroughEndpoint(LiteLLMPydanticObjectBase): + path: str + target: str + headers: dict + + +class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): + user_id: str + team_id: str + budget_id: Optional[str] = None + spend: Optional[float] = 0.0 + litellm_budget_table: Optional[LiteLLM_BudgetTable] + + def safe_get_team_member_rpm_limit(self) -> Optional[int]: + if self.litellm_budget_table is not None: + return self.litellm_budget_table.rpm_limit + return None + + def safe_get_team_member_tpm_limit(self) -> Optional[int]: + if self.litellm_budget_table is not None: + return self.litellm_budget_table.tpm_limit + return None + + +#### Organization / Team Member Requests #### + + +class MemberAddRequest(LiteLLMPydanticObjectBase): + member: Union[List[Member], Member] = Field( + description="Member object or list of member objects to add. Each member must include either user_id or user_email, and a role" + ) + + def __init__(self, **data): + member_data = data.get("member") + if isinstance(member_data, list): + # If member is a list of dictionaries, convert each dictionary to a Member object + members = [ + Member(**item) if isinstance(item, dict) else item + for item in member_data + ] + # Replace member_data with the list of Member objects + data["member"] = members + elif isinstance(member_data, dict): + # If member is a dictionary, convert it to a single Member object + member = Member(**member_data) + # Replace member_data with the single Member object + data["member"] = member + # Call the superclass __init__ method to initialize the object + super().__init__(**data) + + +class OrgMemberAddRequest(LiteLLMPydanticObjectBase): + member: Union[List[OrgMember], OrgMember] + + def __init__(self, **data): + member_data = data.get("member") + if isinstance(member_data, list): + # If member is a list of dictionaries, convert each dictionary to a Member object + if all(isinstance(item, dict) for item in member_data): + members = [OrgMember(**item) for item in member_data] + else: + members = [item for item in member_data] + # Replace member_data with the list of Member objects + data["member"] = members + elif isinstance(member_data, dict): + # If member is a dictionary, convert it to a single Member object + member = OrgMember(**member_data) + # Replace member_data with the single Member object + data["member"] = member + # Call the superclass __init__ method to initialize the object + super().__init__(**data) + + +class TeamAddMemberResponse(LiteLLM_TeamTable): + updated_users: List[LiteLLM_UserTable] + updated_team_memberships: List[LiteLLM_TeamMembership] + + +class OrganizationAddMemberResponse(LiteLLMPydanticObjectBase): + organization_id: str + updated_users: List[LiteLLM_UserTable] + updated_organization_memberships: List[LiteLLM_OrganizationMembershipTable] + + +class MemberDeleteRequest(LiteLLMPydanticObjectBase): + user_id: Optional[str] = None + user_email: Optional[str] = None + + @model_validator(mode="before") + @classmethod + def check_user_info(cls, values): + if values.get("user_id") is None and values.get("user_email") is None: + raise ValueError("Either user id or user email must be provided") + return values + + +class MemberUpdateResponse(LiteLLMPydanticObjectBase): + user_id: str + user_email: Optional[str] = None + + +# Team Member Requests +class TeamMemberAddRequest(MemberAddRequest): + """ + Request body for adding members to a team. + + Example: + ```json + { + "team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849", + "member": { + "role": "user", + "user_id": "user123" + }, + "max_budget_in_team": 100.0 + } + ``` + """ + + team_id: str = Field(description="The ID of the team to add the member to") + max_budget_in_team: Optional[float] = Field( + default=None, + description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", + ) + + +class TeamMemberDeleteRequest(MemberDeleteRequest): + team_id: str + + +class TeamMemberUpdateRequest(TeamMemberDeleteRequest): + max_budget_in_team: Optional[float] = None + role: Optional[Literal["admin", "user"]] = None + tpm_limit: Optional[int] = Field( + default=None, description="Tokens per minute limit for this team member" + ) + rpm_limit: Optional[int] = Field( + default=None, description="Requests per minute limit for this team member" + ) + + +class TeamMemberUpdateResponse(MemberUpdateResponse): + team_id: str + max_budget_in_team: Optional[float] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + + +class TeamModelAddRequest(BaseModel): + """Request to add models to a team""" + + team_id: str + models: List[str] + + +class TeamModelDeleteRequest(BaseModel): + """Request to delete models from a team""" + + team_id: str + models: List[str] + + +# Organization Member Requests +class OrganizationMemberAddRequest(OrgMemberAddRequest): + organization_id: str + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization + + +class OrganizationMemberDeleteRequest(MemberDeleteRequest): + organization_id: str + + +ROLES_WITHIN_ORG = [ + LitellmUserRoles.ORG_ADMIN, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, +] + + +class OrganizationMemberUpdateRequest(OrganizationMemberDeleteRequest): + max_budget_in_organization: Optional[float] = None + role: Optional[LitellmUserRoles] = None + + @field_validator("role") + def validate_role( + cls, value: Optional[LitellmUserRoles] + ) -> Optional[LitellmUserRoles]: + if value is not None and value not in ROLES_WITHIN_ORG: + raise ValueError( + f"Invalid role. Must be one of: {[role.value for role in ROLES_WITHIN_ORG]}" + ) + return value + + +class OrganizationMemberUpdateResponse(MemberUpdateResponse): + organization_id: str + max_budget_in_organization: float + + +########################################## + + +class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): + team_member_budget_table: Optional[LiteLLM_BudgetTable] = None + + +class TeamInfoResponseObject(TypedDict): + team_id: str + team_info: TeamInfoResponseObjectTeamTable + keys: List + team_memberships: List[LiteLLM_TeamMembership] + + +class TeamListResponseObject(LiteLLM_TeamTable): + team_memberships: List[LiteLLM_TeamMembership] + keys: List # list of keys that belong to the team + + +class KeyListResponseObject(TypedDict, total=False): + keys: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] + total_count: Optional[int] + current_page: Optional[int] + total_pages: Optional[int] + + +class CurrentItemRateLimit(TypedDict): + current_requests: int + current_tpm: int + current_rpm: int + + +class LoggingCallbackStatus(TypedDict, total=False): + callbacks: List[str] + status: Literal["healthy", "unhealthy"] + details: Optional[str] + + +class KeyHealthResponse(TypedDict, total=False): + key: Literal["healthy", "unhealthy"] + logging_callbacks: Optional[LoggingCallbackStatus] + + +class SpecialHeaders(enum.Enum): + """Used by user_api_key_auth.py to get litellm key""" + + openai_authorization = "Authorization" + azure_authorization = "API-Key" + anthropic_authorization = "x-api-key" + google_ai_studio_authorization = "x-goog-api-key" + azure_apim_authorization = "Ocp-Apim-Subscription-Key" + custom_litellm_api_key = "x-litellm-api-key" + mcp_auth = "x-mcp-auth" + mcp_servers = "x-mcp-servers" + mcp_access_groups = "x-mcp-access-groups" + + +class LitellmDataForBackendLLMCall(TypedDict, total=False): + headers: dict + organization: str + timeout: Optional[float] + stream_timeout: Optional[float] + user: Optional[str] + num_retries: Optional[int] + + +class LitellmMetadataFromRequestHeaders(TypedDict, total=False): + """ + Headers a user can pass that will get added to litellm metadata for the request + """ + + spend_logs_metadata: Optional[dict] + agent_id: Optional[str] + trace_id: Optional[str] + + +class JWTKeyItem(TypedDict, total=False): + kid: str + + +JWKKeyValue = Union[List[JWTKeyItem], JWTKeyItem] + + +class JWKUrlResponse(TypedDict, total=False): + keys: JWKKeyValue + + +class UserManagementEndpointParamDocStringEnums(str, enum.Enum): + user_id_doc_str = ( + "Optional[str] - Specify a user id. If not set, a unique id will be generated." + ) + user_alias_doc_str = ( + "Optional[str] - A descriptive name for you to know who this user id refers to." + ) + teams_doc_str = "Optional[list] - specify a list of team id's a user belongs to." + user_email_doc_str = "Optional[str] - Specify a user email." + send_invite_email_doc_str = ( + "Optional[bool] - Specify if an invite email should be sent." + ) + user_role_doc_str = """Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`""" + max_budget_doc_str = """Optional[float] - Specify max budget for a given user.""" + budget_duration_doc_str = """Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").""" + models_doc_str = """Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)""" + tpm_limit_doc_str = ( + """Optional[int] - Specify tpm limit for a given user (Tokens per minute)""" + ) + rpm_limit_doc_str = ( + """Optional[int] - Specify rpm limit for a given user (Requests per minute)""" + ) + auto_create_key_doc_str = """bool - Default=True. Flag used for returning a key as part of the /user/new response""" + aliases_doc_str = """Optional[dict] - Model aliases for the user - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)""" + config_doc_str = """Optional[dict] - [DEPRECATED PARAM] User-specific config.""" + allowed_cache_controls_doc_str = """Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request-""" + blocked_doc_str = ( + """Optional[bool] - [Not Implemented Yet] Whether the user is blocked.""" + ) + guardrails_doc_str = """Optional[List[str]] - [Not Implemented Yet] List of active guardrails for the user""" + permissions_doc_str = """Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.""" + metadata_doc_str = """Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }""" + max_parallel_requests_doc_str = """Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.""" + soft_budget_doc_str = """Optional[float] - Get alerts when user crosses given budget, doesn't block requests.""" + model_max_budget_doc_str = """Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)""" + model_rpm_limit_doc_str = """Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" + model_tpm_limit_doc_str = """Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" + spend_doc_str = """Optional[float] - Amount spent by user. Default is 0. Will be updated by proxy whenever user is used.""" + team_id_doc_str = """Optional[str] - [DEPRECATED PARAM] The team id of the user. Default is None.""" + duration_doc_str = """Optional[str] - Duration for the key auto-created on `/user/new`. Default is None.""" + + +PassThroughEndpointLoggingResultValues = Union[ + ModelResponse, + TextCompletionResponse, + ImageResponse, + EmbeddingResponse, + VideoObject, + StandardPassThroughResponseObject, +] + + +class PassThroughEndpointLoggingTypedDict(TypedDict): + result: Optional[PassThroughEndpointLoggingResultValues] + kwargs: dict + + +LiteLLM_ManagementEndpoint_MetadataFields = [ + "model_rpm_limit", + "model_tpm_limit", + "rpm_limit_type", + "tpm_limit_type", + "enforced_params", + "temp_budget_increase", + "temp_budget_expiry", + "allowed_vector_store_indexes", +] + +LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ + "guardrails", + "policies", + "tags", + "team_member_key_duration", + "prompts", + "logging", + "secret_manager_settings", + "allowed_passthrough_routes", +] + + +class ProviderBudgetResponseObject(LiteLLMPydanticObjectBase): + """ + Configuration for a single provider's budget settings + """ + + budget_limit: Optional[float] # Budget limit in USD for the time period + time_period: Optional[str] # Time period for budget (e.g., '1d', '30d', '1mo') + spend: Optional[float] = 0.0 # Current spend for this provider + budget_reset_at: Optional[str] = None # When the current budget period resets + + +class ProviderBudgetResponse(LiteLLMPydanticObjectBase): + """ + Complete provider budget configuration and status. + Maps provider names to their budget configs. + """ + + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations + + +class ProxyStateVariables(TypedDict): + """ + TypedDict for Proxy state variables. + """ + + spend_logs_row_count: int + + +UI_TEAM_ID = "litellm-dashboard" + + +class JWTAuthBuilderResult(TypedDict): + is_proxy_admin: bool + team_object: Optional[LiteLLM_TeamTable] + user_object: Optional[LiteLLM_UserTable] + end_user_object: Optional[LiteLLM_EndUserTable] + org_object: Optional[LiteLLM_OrganizationTable] + token: str + team_id: Optional[str] + user_id: Optional[str] + end_user_id: Optional[str] + org_id: Optional[str] + team_membership: Optional[LiteLLM_TeamMembership] + + +class ClientSideFallbackModel(TypedDict, total=False): + """ + Dictionary passed when client configuring input + """ + + model: Required[str] + messages: List[AllMessageValues] + + +ALL_FALLBACK_MODEL_VALUES = Union[str, ClientSideFallbackModel] + + +RBAC_ROLES = Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.TEAM, + LitellmUserRoles.INTERNAL_USER, +] + + +class OIDCPermissions(LiteLLMPydanticObjectBase): + models: Optional[List[str]] = None + routes: Optional[List[str]] = None + + +class RoleBasedPermissions(OIDCPermissions): + role: RBAC_ROLES + + model_config = { + "extra": "forbid", + } + + +class RoleMapping(BaseModel): + role: str + internal_role: RBAC_ROLES + + +class JWTLiteLLMRoleMap(BaseModel): + jwt_role: str + litellm_role: LitellmUserRoles + + +class ScopeMapping(OIDCPermissions): + scope: str + + model_config = { + "extra": "forbid", + } + + +class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): + """ + A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth. + + Attributes: + - admin_jwt_scope: The JWT scope required for proxy admin roles. + - admin_allowed_routes: list of allowed routes for proxy admin roles. + - team_jwt_scope: The JWT scope required for proxy team roles. + - team_id_jwt_field: The field in the JWT token that stores the team ID. Default - `client_id`. + - team_allowed_routes: list of allowed routes for proxy team roles. + - user_id_jwt_field: The field in the JWT token that stores the user id (maps to `LiteLLMUserTable`). Use this for internal employees. + - user_email_jwt_field: The field in the JWT token that stores the user email (maps to `LiteLLMUserTable`). Use this for internal employees. + - user_allowed_email_subdomain: If specified, only emails from specified subdomain will be allowed to access proxy. + - end_user_id_jwt_field: The field in the JWT token that stores the end-user ID (maps to `LiteLLMEndUserTable`). Turn this off by setting to `None`. Enables end-user cost tracking. Use this for external customers. + - public_key_ttl: Default - 600s. TTL for caching public JWT keys. + - public_allowed_routes: list of allowed routes for authenticated but unknown litellm role jwt tokens. + - enforce_rbac: If true, enforce RBAC for all routes. + - custom_validate: A custom function to validates the JWT token. + - oidc_userinfo_endpoint: OIDC UserInfo endpoint URL. When set along with oidc_userinfo_enabled, LiteLLM will call this endpoint with the access token to retrieve user identity information. + - oidc_userinfo_enabled: Enable fetching user info from OIDC UserInfo endpoint instead of just decoding JWT token. Default: False. + - oidc_userinfo_cache_ttl: TTL (in seconds) for caching UserInfo responses. Default: 300s (5 minutes). + + See `auth_checks.py` for the specific routes + """ + + admin_jwt_scope: str = "litellm_proxy_admin" + admin_allowed_routes: List[str] = [ + "management_routes", + "spend_tracking_routes", + "global_spend_tracking_routes", + "info_routes", + ] + team_id_jwt_field: Optional[str] = None + team_id_upsert: bool = False + team_ids_jwt_field: Optional[str] = None + upsert_sso_user_to_team: bool = False + team_allowed_routes: List[str] = ["openai_routes", "info_routes"] + team_id_default: Optional[str] = Field( + default=None, + description="If no team_id given, default permissions/spend-tracking to this team.s", + ) + team_alias_jwt_field: Optional[str] = Field( + default=None, + description="The field in the JWT token that stores the team name/alias. Will be resolved to team_id via database lookup.", + ) + + org_id_jwt_field: Optional[str] = None + org_alias_jwt_field: Optional[str] = Field( + default=None, + description="The field in the JWT token that stores the organization name/alias. Will be resolved to org_id via database lookup.", + ) + user_id_jwt_field: Optional[str] = None + user_email_jwt_field: Optional[str] = None + user_allowed_email_domain: Optional[str] = None + user_roles_jwt_field: Optional[str] = None + user_allowed_roles: Optional[List[str]] = None + user_id_upsert: bool = Field( + default=False, description="If user doesn't exist, upsert them into the db." + ) + end_user_id_jwt_field: Optional[str] = None + public_key_ttl: float = 600 + public_allowed_routes: List[str] = ["public_routes"] + enforce_rbac: bool = False + roles_jwt_field: Optional[str] = None # v2 on role mappings + role_mappings: Optional[List[RoleMapping]] = None + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping + scope_mappings: Optional[List[ScopeMapping]] = None + enforce_scope_based_access: bool = False + enforce_team_based_model_access: bool = False + custom_validate: Optional[Callable[..., Literal[True]]] = None + ######################################################### + # Fields for syncing user team membership and roles with IDP provider + jwt_litellm_role_map: Optional[List[JWTLiteLLMRoleMap]] = None + sync_user_role_and_teams: bool = False + ######################################################### + ######################################################### + # OIDC UserInfo Endpoint Configuration + oidc_userinfo_endpoint: Optional[str] = Field( + default=None, + description="OIDC UserInfo endpoint URL. If set, LiteLLM will call this endpoint with the access token to retrieve user identity information.", + ) + oidc_userinfo_enabled: bool = Field( + default=False, + description="Enable fetching user info from OIDC UserInfo endpoint instead of just decoding JWT token.", + ) + oidc_userinfo_cache_ttl: float = Field( + default=300, + description="TTL (in seconds) for caching UserInfo responses. Default: 300s (5 minutes).", + ) + ######################################################### + + def __init__(self, **kwargs: Any) -> None: + # get the attribute names for this Pydantic model + allowed_keys = self.__annotations__.keys() + + invalid_keys = set(kwargs.keys()) - allowed_keys + user_roles_jwt_field = kwargs.get("user_roles_jwt_field") + user_allowed_roles = kwargs.get("user_allowed_roles") + object_id_jwt_field = kwargs.get("object_id_jwt_field") + role_mappings = kwargs.get("role_mappings") + scope_mappings = kwargs.get("scope_mappings") + enforce_scope_based_access = kwargs.get("enforce_scope_based_access") + custom_validate = kwargs.get("custom_validate") + + if custom_validate is not None: + fn = get_instance_fn(custom_validate) + validate_custom_validate_return_type(fn) + kwargs["custom_validate"] = fn + + if invalid_keys: + raise ValueError( + f"Invalid arguments provided: {', '.join(invalid_keys)}. Allowed arguments are: {', '.join(allowed_keys)}." + ) + if (user_roles_jwt_field is not None and user_allowed_roles is None) or ( + user_roles_jwt_field is None and user_allowed_roles is not None + ): + raise ValueError( + "user_allowed_roles must be provided if user_roles_jwt_field is set." + ) + + if object_id_jwt_field is not None and role_mappings is None: + raise ValueError( + "if object_id_jwt_field is set, role_mappings must also be set. Needed to infer if the caller is a user or team." + ) + + if scope_mappings is not None and not enforce_scope_based_access: + raise ValueError( + "scope_mappings must be set if enforce_scope_based_access is true." + ) + + super().__init__(**kwargs) + + +class PrismaCompatibleUpdateDBModel(TypedDict, total=False): + model_name: str + litellm_params: str + model_info: str + updated_at: str + updated_by: str + + +class SpecialManagementEndpointEnums(enum.Enum): + DEFAULT_ORGANIZATION = "default_organization" + + +class TransformRequestBody(BaseModel): + call_type: CallTypes + request_body: dict + + +class DefaultInternalUserParams(LiteLLMPydanticObjectBase): + """ + Default parameters to apply when a new user signs in via SSO or is created on the /user/new API endpoint + """ + + user_role: Optional[ + Literal[ + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ] + ] = Field( + default=LitellmUserRoles.INTERNAL_USER, + description="Default role assigned to new users created", + ) + max_budget: Optional[float] = Field( + default=None, + description="Default maximum budget (in USD) for new users created", + ) + budget_duration: Optional[str] = Field( + default=None, + description="Default budget duration for new users (e.g. 'daily', 'weekly', 'monthly')", + ) + models: Optional[List[str]] = Field( + default=None, description="Default list of models that new users can access" + ) + + teams: Optional[Union[List[str], List[NewUserRequestTeam]]] = Field( + default=None, + description="Default teams for new users created", + ) + + +class BaseDailySpendTransaction(TypedDict): + date: str + api_key: str + model: Optional[str] + model_group: Optional[str] + mcp_namespaced_tool_name: Optional[str] + custom_llm_provider: Optional[str] + endpoint: Optional[str] + + # token count metrics + prompt_tokens: int + completion_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + + # request level metrics + spend: float + api_requests: int + successful_requests: int + failed_requests: int + + +class DailyTeamSpendTransaction(BaseDailySpendTransaction): + team_id: str + + +class DailyOrganizationSpendTransaction(BaseDailySpendTransaction): + organization_id: str + + +class DailyUserSpendTransaction(BaseDailySpendTransaction): + user_id: str + + +class DailyEndUserSpendTransaction(BaseDailySpendTransaction): + end_user_id: str + + +class DailyTagSpendTransaction(BaseDailySpendTransaction): + request_id: Optional[str] + tag: str + + +class DailyAgentSpendTransaction(BaseDailySpendTransaction): + agent_id: str + + +class DBSpendUpdateTransactions(TypedDict): + """ + Internal Data Structure for buffering spend updates in Redis or in memory before committing them to the database + """ + + user_list_transactions: Optional[Dict[str, float]] + end_user_list_transactions: Optional[Dict[str, float]] + key_list_transactions: Optional[Dict[str, float]] + team_list_transactions: Optional[Dict[str, float]] + team_member_list_transactions: Optional[Dict[str, float]] + org_list_transactions: Optional[Dict[str, float]] + tag_list_transactions: Optional[Dict[str, float]] + + +class SpendUpdateQueueItem(TypedDict, total=False): + entity_type: Litellm_EntityType + entity_id: str + response_cost: Optional[float] + + +class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): + unified_file_id: str + file_object: Optional[OpenAIFileObject] = None + model_mappings: Dict[str, str] + flat_model_file_ids: List[str] + created_by: Optional[str] + updated_by: Optional[str] + storage_backend: Optional[str] = None + storage_url: Optional[str] = None + + +class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): + unified_object_id: str + model_object_id: str + file_purpose: Literal["batch", "fine-tune", "response"] + file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] + + +class EnterpriseLicenseData(TypedDict, total=False): + expiration_date: str + user_id: str + allowed_features: List[str] + max_users: int + max_teams: int + + +class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): + vector_store_id: str + custom_llm_provider: str + vector_store_name: Optional[str] + vector_store_description: Optional[str] + vector_store_metadata: Optional[Dict[str, Any]] + created_at: Optional[datetime] + updated_at: Optional[datetime] + litellm_credential_name: Optional[str] + litellm_params: Optional[Dict[str, Any]] + team_id: Optional[str] + user_id: Optional[str] + + +class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False): + vector_store: LiteLLM_ManagedVectorStoresTable + + +class CostEstimateRequest(LiteLLMPydanticObjectBase): + """Request body for /cost/estimate endpoint.""" + + model: str = Field(description="Model name (from /model_group/info)") + input_tokens: int = Field(description="Expected input tokens per request", ge=0) + output_tokens: int = Field(description="Expected output tokens per request", ge=0) + num_requests_per_day: Optional[int] = Field( + default=None, description="Number of requests per day", ge=0 + ) + num_requests_per_month: Optional[int] = Field( + default=None, description="Number of requests per month", ge=0 + ) + + +class CostEstimateResponse(LiteLLMPydanticObjectBase): + """Response body for /cost/estimate endpoint.""" + + model: str + input_tokens: int + output_tokens: int + num_requests_per_day: Optional[int] = None + num_requests_per_month: Optional[int] = None + # Per-request costs + cost_per_request: float = Field( + description="Total cost per request (includes margin)" + ) + input_cost_per_request: float = Field( + description="Input token cost per request (before margin)" + ) + output_cost_per_request: float = Field( + description="Output token cost per request (before margin)" + ) + margin_cost_per_request: float = Field( + default=0.0, description="Margin/fee added per request" + ) + # Daily costs (if num_requests_per_day provided) + daily_cost: Optional[float] = Field( + default=None, description="Total daily cost (includes margin)" + ) + daily_input_cost: Optional[float] = Field( + default=None, description="Daily input token cost" + ) + daily_output_cost: Optional[float] = Field( + default=None, description="Daily output token cost" + ) + daily_margin_cost: Optional[float] = Field( + default=None, description="Daily margin/fee" + ) + # Monthly costs (if num_requests_per_month provided) + monthly_cost: Optional[float] = Field( + default=None, description="Total monthly cost (includes margin)" + ) + monthly_input_cost: Optional[float] = Field( + default=None, description="Monthly input token cost" + ) + monthly_output_cost: Optional[float] = Field( + default=None, description="Monthly output token cost" + ) + monthly_margin_cost: Optional[float] = Field( + default=None, description="Monthly margin/fee" + ) + # Pricing info + input_cost_per_token: Optional[float] = None + output_cost_per_token: Optional[float] = None + provider: Optional[str] = None diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index ac02c915366..1950f442813 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,192 +1,208 @@ -import json -from typing import List, Optional - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import DualCache -from litellm.integrations.custom_logger import Span -from litellm.proxy._types import UserAPIKeyAuth -from litellm.router_strategy.budget_limiter import RouterBudgetLimiting -from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ( - BudgetConfig, - GenericBudgetConfigType, - StandardLoggingPayload, -) - -VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX = "virtual_key_spend" - - -class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): - """ - Handles budgets for model + virtual key - - Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d - """ - - def __init__(self, dual_cache: DualCache): - self.dual_cache = dual_cache - self.redis_increment_operation_queue = [] - - async def is_key_within_model_budget( - self, - user_api_key_dict: UserAPIKeyAuth, - model: str, - ) -> bool: - """ - Check if the user_api_key_dict is within the model budget - - Raises: - BudgetExceededError: If the user_api_key_dict has exceeded the model budget - """ - _model_max_budget = user_api_key_dict.model_max_budget - internal_model_max_budget: GenericBudgetConfigType = {} - - for _model, _budget_info in _model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - - verbose_proxy_logger.debug( - "internal_model_max_budget %s", - json.dumps(internal_model_max_budget, indent=4, default=str), - ) - - # check if current model is in internal_model_max_budget - _current_model_budget_info = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if _current_model_budget_info is None: - verbose_proxy_logger.debug( - f"Model {model} not found in internal_model_max_budget" - ) - return True - - # check if current model is within budget - if ( - _current_model_budget_info.max_budget - and _current_model_budget_info.max_budget > 0 - ): - _current_spend = await self._get_virtual_key_spend_for_model( - user_api_key_hash=user_api_key_dict.token, - model=model, - key_budget_config=_current_model_budget_info, - ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): - raise litellm.BudgetExceededError( - message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", - current_cost=_current_spend, - max_budget=_current_model_budget_info.max_budget, - ) - - return True - - async def _get_virtual_key_spend_for_model( - self, - user_api_key_hash: Optional[str], - model: str, - key_budget_config: BudgetConfig, - ) -> Optional[float]: - """ - Get the current spend for a virtual key for a model - - Lookup model in this order: - 1. model: directly look up `model` - 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - """ - - # 1. model: directly look up `model` - virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}" - _current_spend = await self.dual_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - - if _current_spend is None: - # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - # if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview - virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" - _current_spend = await self.dual_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - return _current_spend - - def _get_request_model_budget_config( - self, model: str, internal_model_max_budget: GenericBudgetConfigType - ) -> Optional[BudgetConfig]: - """ - Get the budget config for the request model - - 1. Check if `model` is in `internal_model_max_budget` - 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` - """ - return internal_model_max_budget.get( - model, None - ) or internal_model_max_budget.get( - self._get_model_without_custom_llm_provider(model), None - ) - - def _get_model_without_custom_llm_provider(self, model: str) -> str: - if "/" in model: - return model.split("/")[-1] - return model - - async def async_filter_deployments( - self, - model: str, - healthy_deployments: List, - messages: Optional[List[AllMessageValues]], - request_kwargs: Optional[dict] = None, - parent_otel_span: Optional[Span] = None, # type: ignore - ) -> List[dict]: - return healthy_deployments - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - """ - Track spend for virtual key + model in DualCache - - Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d - """ - verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event") - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) - if standard_logging_payload is None: - raise ValueError("standard_logging_payload is required") - - _litellm_params: dict = kwargs.get("litellm_params", {}) or {} - _metadata: dict = _litellm_params.get("metadata", {}) or {} - user_api_key_model_max_budget: Optional[dict] = _metadata.get( - "user_api_key_model_max_budget", None - ) - if ( - user_api_key_model_max_budget is None - or len(user_api_key_model_max_budget) == 0 - ): - verbose_proxy_logger.debug( - "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget is None or empty. `user_api_key_model_max_budget`=%s", - user_api_key_model_max_budget, - ) - return - response_cost: float = standard_logging_payload.get("response_cost", 0) - model = standard_logging_payload.get("model") - - virtual_key = standard_logging_payload.get("metadata").get("user_api_key_hash") - model = standard_logging_payload.get("model") - if virtual_key is not None: - budget_config = BudgetConfig(time_period="1d", budget_limit=0.1) - virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_config.budget_duration}" - virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}" - await self._increment_spend_for_key( - budget_config=budget_config, - spend_key=virtual_spend_key, - start_time_key=virtual_start_time_key, - response_cost=response_cost, - ) - verbose_proxy_logger.debug( - "current state of in memory cache %s", - json.dumps( - self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str - ), - ) +import json +from typing import List, Optional + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import Span +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ( + BudgetConfig, + GenericBudgetConfigType, + StandardLoggingPayload, +) + +VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX = "virtual_key_spend" + + +class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): + """ + Handles budgets for model + virtual key + + Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d + """ + + def __init__(self, dual_cache: DualCache): + self.dual_cache = dual_cache + self.redis_increment_operation_queue = [] + + async def is_key_within_model_budget( + self, + user_api_key_dict: UserAPIKeyAuth, + model: str, + ) -> bool: + """ + Check if the user_api_key_dict is within the model budget + + Raises: + BudgetExceededError: If the user_api_key_dict has exceeded the model budget + """ + _model_max_budget = user_api_key_dict.model_max_budget + internal_model_max_budget: GenericBudgetConfigType = {} + + for _model, _budget_info in _model_max_budget.items(): + internal_model_max_budget[_model] = BudgetConfig(**_budget_info) + + verbose_proxy_logger.debug( + "internal_model_max_budget %s", + json.dumps(internal_model_max_budget, indent=4, default=str), + ) + + # check if current model is in internal_model_max_budget + _current_model_budget_info = self._get_request_model_budget_config( + model=model, internal_model_max_budget=internal_model_max_budget + ) + if _current_model_budget_info is None: + verbose_proxy_logger.debug( + f"Model {model} not found in internal_model_max_budget" + ) + return True + + # check if current model is within budget + if ( + _current_model_budget_info.max_budget + and _current_model_budget_info.max_budget > 0 + ): + _current_spend = await self._get_virtual_key_spend_for_model( + user_api_key_hash=user_api_key_dict.token, + model=model, + key_budget_config=_current_model_budget_info, + ) + if ( + _current_spend is not None + and _current_model_budget_info.max_budget is not None + and _current_spend > _current_model_budget_info.max_budget + ): + raise litellm.BudgetExceededError( + message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", + current_cost=_current_spend, + max_budget=_current_model_budget_info.max_budget, + ) + + return True + + async def _get_virtual_key_spend_for_model( + self, + user_api_key_hash: Optional[str], + model: str, + key_budget_config: BudgetConfig, + ) -> Optional[float]: + """ + Get the current spend for a virtual key for a model + + Lookup model in this order: + 1. model: directly look up `model` + 2. If 1, does not exist, check if passed as {custom_llm_provider}/model + """ + + # 1. model: directly look up `model` + virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}" + _current_spend = await self.dual_cache.async_get_cache( + key=virtual_key_model_spend_cache_key, + ) + + if _current_spend is None: + # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model + # if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview + virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" + _current_spend = await self.dual_cache.async_get_cache( + key=virtual_key_model_spend_cache_key, + ) + return _current_spend + + def _get_request_model_budget_config( + self, model: str, internal_model_max_budget: GenericBudgetConfigType + ) -> Optional[BudgetConfig]: + """ + Get the budget config for the request model + + 1. Check if `model` is in `internal_model_max_budget` + 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` + """ + return internal_model_max_budget.get( + model, None + ) or internal_model_max_budget.get( + self._get_model_without_custom_llm_provider(model), None + ) + + def _get_model_without_custom_llm_provider(self, model: str) -> str: + if "/" in model: + return model.split("/")[-1] + return model + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: List, + messages: Optional[List[AllMessageValues]], + request_kwargs: Optional[dict] = None, + parent_otel_span: Optional[Span] = None, # type: ignore + ) -> List[dict]: + return healthy_deployments + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Track spend for virtual key + model in DualCache + + Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d + """ + verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event") + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object", None + ) + if standard_logging_payload is None: + raise ValueError("standard_logging_payload is required") + + _litellm_params: dict = kwargs.get("litellm_params", {}) or {} + _metadata: dict = _litellm_params.get("metadata", {}) or {} + user_api_key_model_max_budget: Optional[dict] = _metadata.get( + "user_api_key_model_max_budget", None + ) + if ( + user_api_key_model_max_budget is None + or len(user_api_key_model_max_budget) == 0 + ): + verbose_proxy_logger.debug( + "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget is None or empty. `user_api_key_model_max_budget`=%s", + user_api_key_model_max_budget, + ) + return + response_cost: float = standard_logging_payload.get("response_cost", 0) + model = standard_logging_payload.get("model") + virtual_key = standard_logging_payload.get("metadata", {}).get( + "user_api_key_hash" + ) + + if virtual_key is None or model is None: + return + + # Resolve per-model budget config (same logic as is_key_within_model_budget) + internal_model_max_budget: GenericBudgetConfigType = {} + for _model, _budget_info in user_api_key_model_max_budget.items(): + internal_model_max_budget[_model] = BudgetConfig(**_budget_info) + key_budget_config = self._get_request_model_budget_config( + model=model, internal_model_max_budget=internal_model_max_budget + ) + if key_budget_config is None or not key_budget_config.budget_duration: + verbose_proxy_logger.debug( + "Not incrementing model spend: no budget config or budget_duration for model=%s", + model, + ) + return + + virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" + virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}" + await self._increment_spend_for_key( + budget_config=key_budget_config, + spend_key=virtual_spend_key, + start_time_key=virtual_start_time_key, + response_cost=response_cost, + ) + verbose_proxy_logger.debug( + "current state of in memory cache %s", + json.dumps( + self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str + ), + ) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index e43da32565a..ad52d7afbd3 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -1,321 +1,352 @@ -""" -BUDGET MANAGEMENT - -All /budget management endpoints - -/budget/new -/budget/info -/budget/update -/budget/delete -/budget/settings -/budget/list -""" - -#### BUDGET TABLE MANAGEMENT #### -from datetime import timedelta - -from fastapi import APIRouter, Depends, HTTPException - -from litellm.litellm_core_utils.duration_parser import duration_in_seconds -from litellm.proxy._types import * -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.utils import jsonify_object - -router = APIRouter() - - -@router.post( - "/budget/new", - tags=["budget management"], - dependencies=[Depends(user_api_key_auth)], -) -async def new_budget( - budget_obj: BudgetNewRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Create a new budget object. Can apply this to teams, orgs, end-users, keys. - - Parameters: - - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. - - max_budget: Optional[float] - The max budget for the budget. - - soft_budget: Optional[float] - The soft budget for the budget. - - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. - - tpm_limit: Optional[int] - The tokens per minute limit for the budget. - - rpm_limit: Optional[int] - The requests per minute limit for the budget. - - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. - """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - # Validate budget values are not negative - if budget_obj.max_budget is not None and budget_obj.max_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"} - ) - if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"} - ) - - # if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future - if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None: - budget_obj.budget_reset_at = datetime.utcnow() + timedelta( - seconds=duration_in_seconds(duration=budget_obj.budget_duration) - ) - - budget_obj_json = budget_obj.model_dump(exclude_none=True) - budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries - response = await prisma_client.db.litellm_budgettable.create( - data={ - **budget_obj_jsonified, # type: ignore - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } # type: ignore - ) - - return response - - -@router.post( - "/budget/update", - tags=["budget management"], - dependencies=[Depends(user_api_key_auth)], -) -async def update_budget( - budget_obj: BudgetNewRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Update an existing budget object. - - Parameters: - - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. - - max_budget: Optional[float] - The max budget for the budget. - - soft_budget: Optional[float] - The soft budget for the budget. - - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. - - tpm_limit: Optional[int] - The tokens per minute limit for the budget. - - rpm_limit: Optional[int] - The requests per minute limit for the budget. - - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. - """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - if budget_obj.budget_id is None: - raise HTTPException(status_code=400, detail={"error": "budget_id is required"}) - - # Validate budget values are not negative - if budget_obj.max_budget is not None and budget_obj.max_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"} - ) - if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"} - ) - - response = await prisma_client.db.litellm_budgettable.update( - where={"budget_id": budget_obj.budget_id}, - data={ - **budget_obj.model_dump(exclude_unset=True), # type: ignore - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - }, # type: ignore - ) - - return response - - -@router.post( - "/budget/info", - tags=["budget management"], - dependencies=[Depends(user_api_key_auth)], -) -async def info_budget(data: BudgetRequest): - """ - Get the budget id specific information - - Parameters: - - budgets: List[str] - The list of budget ids to get information for - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException(status_code=500, detail={"error": "No db connected"}) - - if len(data.budgets) == 0: - raise HTTPException( - status_code=400, - detail={ - "error": f"Specify list of budget id's to query. Passed in={data.budgets}" - }, - ) - response = await prisma_client.db.litellm_budgettable.find_many( - where={"budget_id": {"in": data.budgets}}, - ) - - return response - - -@router.get( - "/budget/settings", - tags=["budget management"], - dependencies=[Depends(user_api_key_auth)], -) -async def budget_settings( - budget_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Get list of configurable params + current value for a budget item + description of each field - - Used on Admin UI. - - Query Parameters: - - budget_id: str - The budget id to get information for - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=400, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=400, - detail={ - "error": "{}, your role={}".format( - CommonProxyErrors.not_allowed_access.value, - user_api_key_dict.user_role, - ) - }, - ) - - ## get budget item from db - db_budget_row = await prisma_client.db.litellm_budgettable.find_first( - where={"budget_id": budget_id} - ) - - if db_budget_row is not None: - db_budget_row_dict = db_budget_row.model_dump(exclude_none=True) - else: - db_budget_row_dict = {} - - allowed_args = { - "max_parallel_requests": {"type": "Integer"}, - "tpm_limit": {"type": "Integer"}, - "rpm_limit": {"type": "Integer"}, - "budget_duration": {"type": "String"}, - "max_budget": {"type": "Float"}, - "soft_budget": {"type": "Float"}, - } - - return_val = [] - - for field_name, field_info in BudgetNewRequest.model_fields.items(): - if field_name in allowed_args: - _stored_in_db = True - - _response_obj = ConfigList( - field_name=field_name, - field_type=allowed_args[field_name]["type"], - field_description=field_info.description or "", - field_value=db_budget_row_dict.get(field_name, None), - stored_in_db=_stored_in_db, - field_default_value=field_info.default, - ) - return_val.append(_response_obj) - - return return_val - - -@router.get( - "/budget/list", - tags=["budget management"], - dependencies=[Depends(user_api_key_auth)], -) -async def list_budget( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """List all the created budgets in proxy db. Used on Admin UI.""" - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=400, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=400, - detail={ - "error": "{}, your role={}".format( - CommonProxyErrors.not_allowed_access.value, - user_api_key_dict.user_role, - ) - }, - ) - - response = await prisma_client.db.litellm_budgettable.find_many() - - return response - - -@router.post( - "/budget/delete", - tags=["budget management"], - dependencies=[Depends(user_api_key_auth)], -) -async def delete_budget( - data: BudgetDeleteRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Delete budget - - Parameters: - - id: str - The budget id to delete - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=400, - detail={ - "error": "{}, your role={}".format( - CommonProxyErrors.not_allowed_access.value, - user_api_key_dict.user_role, - ) - }, - ) - - response = await prisma_client.db.litellm_budgettable.delete( - where={"budget_id": data.id} - ) - - return response +""" +BUDGET MANAGEMENT + +All /budget management endpoints + +/budget/new +/budget/info +/budget/update +/budget/delete +/budget/settings +/budget/list +""" + +#### BUDGET TABLE MANAGEMENT #### +from datetime import timedelta + +from fastapi import APIRouter, Depends, HTTPException + +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import * +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import jsonify_object + +router = APIRouter() + + +@router.post( + "/budget/new", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def new_budget( + budget_obj: BudgetNewRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create a new budget object. Can apply this to teams, orgs, end-users, keys. + + Parameters: + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. + - max_budget: Optional[float] - The max budget for the budget. + - soft_budget: Optional[float] - The soft budget for the budget. + - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. + - tpm_limit: Optional[int] - The tokens per minute limit for the budget. + - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} + - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Validate budget values are not negative + if budget_obj.max_budget is not None and budget_obj.max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}" + }, + ) + if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}" + }, + ) + + # Validate model_max_budget if present + if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_model_max_budget, + ) + + try: + validate_model_max_budget(budget_obj.model_max_budget) + except ValueError as e: + raise HTTPException(status_code=400, detail={"error": str(e)}) + + # if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future + if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None: + budget_obj.budget_reset_at = datetime.utcnow() + timedelta( + seconds=duration_in_seconds(duration=budget_obj.budget_duration) + ) + + budget_obj_json = budget_obj.model_dump(exclude_none=True) + budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries + response = await prisma_client.db.litellm_budgettable.create( + data={ + **budget_obj_jsonified, # type: ignore + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } # type: ignore + ) + + return response + + +@router.post( + "/budget/update", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_budget( + budget_obj: BudgetNewRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update an existing budget object. + + Parameters: + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. + - max_budget: Optional[float] - The max budget for the budget. + - soft_budget: Optional[float] - The soft budget for the budget. + - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. + - tpm_limit: Optional[int] - The tokens per minute limit for the budget. + - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} + - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + if budget_obj.budget_id is None: + raise HTTPException(status_code=400, detail={"error": "budget_id is required"}) + + # Validate budget values are not negative + if budget_obj.max_budget is not None and budget_obj.max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}" + }, + ) + if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}" + }, + ) + + # Validate model_max_budget if present in update + if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_model_max_budget, + ) + + try: + validate_model_max_budget(budget_obj.model_max_budget) + except ValueError as e: + raise HTTPException(status_code=400, detail={"error": str(e)}) + + response = await prisma_client.db.litellm_budgettable.update( + where={"budget_id": budget_obj.budget_id}, + data={ + **budget_obj.model_dump(exclude_unset=True), # type: ignore + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + }, # type: ignore + ) + + return response + + +@router.post( + "/budget/info", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def info_budget(data: BudgetRequest): + """ + Get the budget id specific information + + Parameters: + - budgets: List[str] - The list of budget ids to get information for + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + if len(data.budgets) == 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"Specify list of budget id's to query. Passed in={data.budgets}" + }, + ) + response = await prisma_client.db.litellm_budgettable.find_many( + where={"budget_id": {"in": data.budgets}}, + ) + + return response + + +@router.get( + "/budget/settings", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def budget_settings( + budget_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get list of configurable params + current value for a budget item + description of each field + + Used on Admin UI. + + Query Parameters: + - budget_id: str - The budget id to get information for + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=400, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + ## get budget item from db + db_budget_row = await prisma_client.db.litellm_budgettable.find_first( + where={"budget_id": budget_id} + ) + + if db_budget_row is not None: + db_budget_row_dict = db_budget_row.model_dump(exclude_none=True) + else: + db_budget_row_dict = {} + + allowed_args = { + "max_parallel_requests": {"type": "Integer"}, + "tpm_limit": {"type": "Integer"}, + "rpm_limit": {"type": "Integer"}, + "budget_duration": {"type": "String"}, + "max_budget": {"type": "Float"}, + "soft_budget": {"type": "Float"}, + "model_max_budget": {"type": "Object"}, + } + + return_val = [] + + for field_name, field_info in BudgetNewRequest.model_fields.items(): + if field_name in allowed_args: + _stored_in_db = True + + _response_obj = ConfigList( + field_name=field_name, + field_type=allowed_args[field_name]["type"], + field_description=field_info.description or "", + field_value=db_budget_row_dict.get(field_name, None), + stored_in_db=_stored_in_db, + field_default_value=field_info.default, + ) + return_val.append(_response_obj) + + return return_val + + +@router.get( + "/budget/list", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_budget( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """List all the created budgets in proxy db. Used on Admin UI.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=400, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + response = await prisma_client.db.litellm_budgettable.find_many() + + return response + + +@router.post( + "/budget/delete", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_budget( + data: BudgetDeleteRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete budget + + Parameters: + - id: str - The budget id to delete + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=400, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + response = await prisma_client.db.litellm_budgettable.delete( + where={"budget_id": data.id} + ) + + return response diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 380e8bddc99..2c4d20a84bd 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1,4465 +1,4496 @@ -""" -KEY MANAGEMENT - -All /key management endpoints - -/key/generate -/key/info -/key/update -/key/delete -""" - -import asyncio -import copy -import json -import secrets -import traceback -from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Literal, Optional, Tuple, cast - -import fastapi -import yaml -from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm._uuid import uuid -from litellm.caching import DualCache -from litellm.constants import ( - LENGTH_OF_LITELLM_GENERATED_KEY, - LITELLM_PROXY_ADMIN_NAME, - UI_SESSION_TOKEN_TEAM_ID, -) -from litellm.litellm_core_utils.duration_parser import duration_in_seconds -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._experimental.mcp_server.db import ( - rotate_mcp_server_credentials_master_key, -) -from litellm.proxy._types import * -from litellm.proxy._types import LiteLLM_VerificationToken -from litellm.types.proxy.management_endpoints.key_management_endpoints import ( - BulkUpdateKeyRequest, - BulkUpdateKeyRequestItem, - BulkUpdateKeyResponse, - FailedKeyUpdate, - SuccessfulKeyUpdate, -) -from litellm.proxy.auth.auth_checks import ( - _cache_key_object, - _delete_cache_key_object, - can_team_access_model, - get_key_object, - get_org_object, - get_team_object, -) -from litellm.proxy.auth.auth_utils import abbreviate_api_key -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks -from litellm.proxy.management_endpoints.common_utils import ( - _is_user_team_admin, - _set_object_metadata_field, -) -from litellm.proxy.management_endpoints.model_management_endpoints import ( - _add_model_to_db, -) -from litellm.proxy.management_helpers.object_permission_utils import ( - _set_object_permission, - attach_object_permission_to_dict, - handle_update_object_permission_common, -) -from litellm.proxy.management_helpers.team_member_permission_checks import ( - TeamMemberPermissionChecks, -) -from litellm.proxy.management_helpers.utils import management_endpoint_wrapper -from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key -from litellm.proxy.utils import ( - PrismaClient, - _hash_token_if_needed, - handle_exception_on_proxy, - is_valid_api_key, - jsonify_object, -) -from litellm.router import Router -from litellm.secret_managers.main import get_secret -from litellm.types.router import Deployment -from litellm.types.utils import ( - BudgetConfig, - PersonalUIKeyGenerationConfig, - TeamUIKeyGenerationConfig, -) - - -def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]): - return data.team_id is not None - - -def _get_user_in_team( - team_table: LiteLLM_TeamTableCachedObj, user_id: Optional[str] -) -> Optional[Member]: - if user_id is None: - return None - for member in team_table.members_with_roles: - if member.user_id is not None and member.user_id == user_id: - return member - - return None - - -def _calculate_key_rotation_time(rotation_interval: str) -> datetime: - """ - Helper function to calculate the next rotation time for a key based on the rotation interval. - - Args: - rotation_interval: String representing the rotation interval (e.g., '30d', '90d', '1h') - - Returns: - datetime: The calculated next rotation time in UTC - """ - now = datetime.now(timezone.utc) - interval_seconds = duration_in_seconds(rotation_interval) - return now + timedelta(seconds=interval_seconds) - - -def _set_key_rotation_fields( - data: dict, auto_rotate: bool, rotation_interval: Optional[str] -) -> None: - """ - Helper function to set rotation fields in key data if auto_rotate is enabled. - - Args: - data: Dictionary to update with rotation fields - auto_rotate: Whether auto rotation is enabled - rotation_interval: The rotation interval string (required if auto_rotate is True) - """ - if auto_rotate and rotation_interval: - data.update( - { - "auto_rotate": auto_rotate, - "rotation_interval": rotation_interval, - "key_rotation_at": _calculate_key_rotation_time(rotation_interval), - } - ) - - -def _is_allowed_to_make_key_request( - user_api_key_dict: UserAPIKeyAuth, - user_id: Optional[str], - team_id: Optional[str], -) -> bool: - """ - Assert user only creates/updates keys for themselves - - Relevant issue: https://github.com/BerriAI/litellm/issues/7336 - """ - ## BASE CASE - PROXY ADMIN - if ( - user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - ): - return True - - if user_id is not None: - assert ( - user_id == user_api_key_dict.user_id - ), "User can only create keys for themselves. Got user_id={}, Your ID={}".format( - user_id, user_api_key_dict.user_id - ) - - if team_id is not None: - if ( - user_api_key_dict.team_id is not None - and user_api_key_dict.team_id == UI_TEAM_ID - ): - return True # handle https://github.com/BerriAI/litellm/issues/7482 - - return True - - -def _team_key_operation_team_member_check( - assigned_user_id: Optional[str], - team_table: LiteLLM_TeamTableCachedObj, - user_api_key_dict: UserAPIKeyAuth, - team_key_generation: TeamUIKeyGenerationConfig, - route: KeyManagementRoutes, -): - if assigned_user_id is not None: - key_assigned_user_in_team = _get_user_in_team( - team_table=team_table, user_id=assigned_user_id - ) - - if key_assigned_user_in_team is None: - raise HTTPException( - status_code=400, - detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}", - ) - - team_member_object = _get_user_in_team( - team_table=team_table, user_id=user_api_key_dict.user_id - ) - - is_admin = ( - user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - ) - - if is_admin: - return True - elif team_member_object is None: - raise HTTPException( - status_code=400, - detail=f"User={user_api_key_dict.user_id} not assigned to team={team_table.team_id}", - ) - elif ( - "allowed_team_member_roles" in team_key_generation - and team_member_object.role - not in team_key_generation["allowed_team_member_roles"] - ): - raise HTTPException( - status_code=400, - detail=f"Team member role {team_member_object.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}", - ) - - TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( - team_member_object=team_member_object, - team_table=team_table, - route=route, - ) - return True - - -def _key_generation_required_param_check( - data: GenerateKeyRequest, required_params: Optional[List[str]] -): - if required_params is None: - return True - - data_dict = data.model_dump(exclude_unset=True) - for param in required_params: - if param not in data_dict: - raise HTTPException( - status_code=400, - detail=f"Required param {param} not in data", - ) - return True - - -def _team_key_generation_check( - team_table: LiteLLM_TeamTableCachedObj, - user_api_key_dict: UserAPIKeyAuth, - data: GenerateKeyRequest, - route: KeyManagementRoutes, -): - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return True - if ( - litellm.key_generation_settings is not None - and "team_key_generation" in litellm.key_generation_settings - ): - _team_key_generation = litellm.key_generation_settings["team_key_generation"] - else: - _team_key_generation = TeamUIKeyGenerationConfig( - allowed_team_member_roles=["admin", "user"], - ) - - _team_key_operation_team_member_check( - assigned_user_id=data.user_id, - team_table=team_table, - user_api_key_dict=user_api_key_dict, - team_key_generation=_team_key_generation, - route=route, - ) - _key_generation_required_param_check( - data, - _team_key_generation.get("required_params"), - ) - - return True - - -def _personal_key_membership_check( - user_api_key_dict: UserAPIKeyAuth, - personal_key_generation: Optional[PersonalUIKeyGenerationConfig], -): - if ( - personal_key_generation is None - or "allowed_user_roles" not in personal_key_generation - ): - return True - - if user_api_key_dict.user_role not in personal_key_generation["allowed_user_roles"]: - raise HTTPException( - status_code=400, - detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}", # type: ignore - ) - - return True - - -def _personal_key_generation_check( - user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest -): - if ( - litellm.key_generation_settings is None - or litellm.key_generation_settings.get("personal_key_generation") is None - ): - return True - - _personal_key_generation = litellm.key_generation_settings["personal_key_generation"] # type: ignore - - _personal_key_membership_check( - user_api_key_dict, - personal_key_generation=_personal_key_generation, - ) - - _key_generation_required_param_check( - data, - _personal_key_generation.get("required_params"), - ) - - return True - - -def key_generation_check( - team_table: Optional[LiteLLM_TeamTableCachedObj], - user_api_key_dict: UserAPIKeyAuth, - data: GenerateKeyRequest, - route: KeyManagementRoutes, -) -> bool: - """ - Check if admin has restricted key creation to certain roles for teams or individuals - """ - - ## check if key is for team or individual - is_team_key = _is_team_key(data=data) - if is_team_key: - if team_table is None and litellm.key_generation_settings is not None: - raise HTTPException( - status_code=400, - detail=f"Unable to find team object in database. Team ID: {data.team_id}", - ) - elif team_table is None: - return True # assume user is assigning team_id without using the team table - return _team_key_generation_check( - team_table=team_table, - user_api_key_dict=user_api_key_dict, - data=data, - route=route, - ) - else: - return _personal_key_generation_check( - user_api_key_dict=user_api_key_dict, data=data - ) - - -def common_key_access_checks( - user_api_key_dict: UserAPIKeyAuth, - data: Union[GenerateKeyRequest, UpdateKeyRequest], - llm_router: Optional[Router], - premium_user: bool, - user_id: Optional[str] = None, -) -> Literal[True]: - """ - Check if user is allowed to make a key request, for this key - """ - try: - _is_allowed_to_make_key_request( - user_api_key_dict=user_api_key_dict, - user_id=user_id or data.user_id, - team_id=data.team_id, - ) - except AssertionError as e: - raise HTTPException( - status_code=403, - detail=str(e), - ) - except Exception as e: - raise HTTPException( - status_code=500, - detail=str(e), - ) - - _check_model_access_group( - models=data.models, - llm_router=llm_router, - premium_user=premium_user, - ) - return True - - -router = APIRouter() - - -def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: - """ - Handle the key type. - """ - key_type = data.key_type - data_json.pop("key_type", None) - if key_type == LiteLLMKeyType.LLM_API: - data_json["allowed_routes"] = ["llm_api_routes"] - elif key_type == LiteLLMKeyType.MANAGEMENT: - data_json["allowed_routes"] = ["management_routes"] - elif key_type == LiteLLMKeyType.READ_ONLY: - data_json["allowed_routes"] = ["info_routes"] - return data_json - - -async def validate_team_id_used_in_service_account_request( - team_id: Optional[str], - prisma_client: Optional[PrismaClient], -): - """ - Validate team_id is used in the request body for generating a service account key - """ - if team_id is None: - raise HTTPException( - status_code=400, - detail="team_id is required for service account keys. Please specify `team_id` in the request body.", - ) - - if prisma_client is None: - raise HTTPException( - status_code=400, - detail="prisma_client is required for service account keys. Please specify `prisma_client` in the request body.", - ) - - # check if team_id exists in the database - team = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - ) - if team is None: - raise HTTPException( - status_code=400, - detail="team_id does not exist in the database. Please specify a valid `team_id` in the request body.", - ) - return True - - -async def _common_key_generation_helper( # noqa: PLR0915 - data: GenerateKeyRequest, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str], - team_table: Optional[LiteLLM_TeamTableCachedObj], -) -> GenerateKeyResponse: - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - llm_router, - premium_user, - prisma_client, - ) - - common_key_access_checks( - user_api_key_dict=user_api_key_dict, - data=data, - llm_router=llm_router, - premium_user=premium_user, - ) - - if ( - data.metadata is not None - and data.metadata.get("service_account_id") is not None - and data.team_id is None - ): - await validate_team_id_used_in_service_account_request( - team_id=data.team_id, - prisma_client=prisma_client, - ) - - # check if user set default key/generate params on config.yaml - if litellm.default_key_generate_params is not None: - for elem in data: - key, value = elem - if value is None and key in [ - "max_budget", - "user_id", - "team_id", - "max_parallel_requests", - "tpm_limit", - "rpm_limit", - "budget_duration", - ]: - setattr(data, key, litellm.default_key_generate_params.get(key, None)) - elif key == "models" and value == []: - setattr(data, key, litellm.default_key_generate_params.get(key, [])) - elif key == "metadata" and value == {}: - setattr(data, key, litellm.default_key_generate_params.get(key, {})) - - # check if user set default key/generate params on config.yaml - if litellm.upperbound_key_generate_params is not None: - for elem in data: - key, value = elem - upperbound_value = getattr( - litellm.upperbound_key_generate_params, key, None - ) - if upperbound_value is not None: - if value is None: - # Use the upperbound value if user didn't provide a value - setattr(data, key, upperbound_value) - else: - # Compare with upperbound for numeric fields - if key in [ - "max_budget", - "max_parallel_requests", - "tpm_limit", - "rpm_limit", - ]: - if value > upperbound_value: - raise HTTPException( - status_code=400, - detail={ - "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" - }, - ) - # Compare durations - elif key in ["budget_duration", "duration"]: - upperbound_duration = duration_in_seconds( - duration=upperbound_value - ) - # Handle special case where duration is "-1" (never expires) - if value == "-1": - user_duration = float('inf') # Infinite duration - else: - user_duration = duration_in_seconds(duration=value) - if user_duration > upperbound_duration: - raise HTTPException( - status_code=400, - detail={ - "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" - }, - ) - - # APPLY ENTERPRISE KEY MANAGEMENT PARAMS - try: - from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import ( - apply_enterprise_key_management_params, - ) - - data = apply_enterprise_key_management_params(data, team_table) - except Exception as e: - verbose_proxy_logger.debug( - "litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - {}".format( - str(e) - ) - ) - - # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable - _budget_id = data.budget_id - if prisma_client is not None and data.soft_budget is not None: - # create the Budget Row for the LiteLLM Verification Token - budget_row = LiteLLM_BudgetTable( - soft_budget=data.soft_budget, - model_max_budget=data.model_max_budget or {}, - ) - new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - - _budget = await prisma_client.db.litellm_budgettable.create( - data={ - **new_budget, # type: ignore - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } - ) - _budget_id = getattr(_budget, "budget_id", None) - - # ADD METADATA FIELDS - # Set Management Endpoint Metadata Fields - for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: - if getattr(data, field, None) is not None: - _set_object_metadata_field( - object_data=data, - field_name=field, - value=getattr(data, field), - ) - delattr(data, field) - - for field in LiteLLM_ManagementEndpoint_MetadataFields: - if getattr(data, field, None) is not None: - _set_object_metadata_field( - object_data=data, - field_name=field, - value=getattr(data, field), - ) - delattr(data, field) - - data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore - - data_json = handle_key_type(data, data_json) - - # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users - if "max_budget" in data_json: - data_json["key_max_budget"] = data_json.pop("max_budget", None) - if _budget_id is not None: - data_json["budget_id"] = _budget_id - - if "budget_duration" in data_json: - data_json["key_budget_duration"] = data_json.pop("budget_duration", None) - - if user_api_key_dict.user_id is not None: - data_json["created_by"] = user_api_key_dict.user_id - data_json["updated_by"] = user_api_key_dict.user_id - - # Set tags on the new key - if "tags" in data_json: - from litellm.proxy.proxy_server import premium_user - - if premium_user is not True and data_json["tags"] is not None: - raise ValueError( - f"Only premium users can add tags to keys. {CommonProxyErrors.not_premium_user.value}" - ) - - _metadata = data_json.get("metadata") - if not _metadata: - data_json["metadata"] = {"tags": data_json["tags"]} - else: - data_json["metadata"]["tags"] = data_json["tags"] - - data_json.pop("tags") - - data_json = await _set_object_permission( - data_json=data_json, - prisma_client=prisma_client, - ) - - await _enforce_unique_key_alias( - key_alias=data_json.get("key_alias", None), - prisma_client=prisma_client, - ) - - # Validate user-provided key format - if data.key is not None and not data.key.startswith("sk-"): - raise HTTPException( - status_code=400, - detail={ - "error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {data.key}" - }, - ) - - # check org key limits - done here to handle inheriting org id from team - if data.organization_id is not None: - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - if prisma_client: - org_table = await get_org_object( - org_id=data.organization_id, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) - if org_table is None: - raise HTTPException( - status_code=400, - detail=f"Organization not found for organization_id={data.organization_id}", - ) - await _check_org_key_limits( - org_table=org_table, - data=data, - prisma_client=prisma_client, - ) - - response = await generate_key_helper_fn( - request_type="key", **data_json, table_name="key" - ) - - response["soft_budget"] = ( - data.soft_budget - ) # include the user-input soft budget in the response - - response = GenerateKeyResponse(**response) - - response.token = ( - response.token_id - ) # remap token to use the hash, and leave the key in the `key` field [TODO]: clean up generate_key_helper_fn to do this - - asyncio.create_task( - KeyManagementEventHooks.async_key_generated_hook( - data=data, - response=response, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - ) - - return response - - -def _check_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], - data: Union[GenerateKeyRequest, UpdateKeyRequest], - entity_rpm_limit: Optional[int], - entity_tpm_limit: Optional[int], - entity_model_rpm_limit_dict: Dict[str, int], - entity_model_tpm_limit_dict: Dict[str, int], - entity_type: str, # "team" or "organization" -) -> None: - """ - Generic function to check if a key is allocating model specific limits. - Raises an error if we're overallocating. - """ - model_rpm_limit = getattr(data, "model_rpm_limit", None) or ( - data.metadata.get("model_rpm_limit", None) if data.metadata else None - ) - model_tpm_limit = getattr(data, "model_tpm_limit", None) or ( - data.metadata.get("model_tpm_limit", None) if data.metadata else None - ) - if model_rpm_limit is None and model_tpm_limit is None: - return - - # get total model specific tpm/rpm limit - model_specific_rpm_limit: Dict[str, int] = {} - model_specific_tpm_limit: Dict[str, int] = {} - - for key in keys: - if key.metadata.get("model_rpm_limit", None) is not None: - for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items(): - model_specific_rpm_limit[model] = ( - model_specific_rpm_limit.get(model, 0) + rpm_limit - ) - if key.metadata.get("model_tpm_limit", None) is not None: - for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items(): - model_specific_tpm_limit[model] = ( - model_specific_tpm_limit.get(model, 0) + tpm_limit - ) - - if model_rpm_limit is not None: - for model, rpm_limit in model_rpm_limit.items(): - if ( - entity_rpm_limit is not None - and model_specific_rpm_limit.get(model, 0) + rpm_limit - > entity_rpm_limit - ): - raise HTTPException( - status_code=400, - detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", - ) - elif entity_model_rpm_limit_dict: - entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model) - if ( - entity_model_specific_rpm_limit - and model_specific_rpm_limit.get(model, 0) + rpm_limit - > entity_model_specific_rpm_limit - ): - raise HTTPException( - status_code=400, - detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit}", - ) - - if model_tpm_limit is not None: - for model, tpm_limit in model_tpm_limit.items(): - if ( - entity_tpm_limit is not None - and model_specific_tpm_limit.get(model, 0) + tpm_limit - > entity_tpm_limit - ): - raise HTTPException( - status_code=400, - detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", - ) - elif entity_model_tpm_limit_dict: - entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model) - if ( - entity_model_specific_tpm_limit - and model_specific_tpm_limit.get(model, 0) + tpm_limit - > entity_model_specific_tpm_limit - ): - raise HTTPException( - status_code=400, - detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit}", - ) - - -def _check_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], - data: Union[GenerateKeyRequest, UpdateKeyRequest], - entity_rpm_limit: Optional[int], - entity_tpm_limit: Optional[int], - entity_type: str, # "team" or "organization" -) -> None: - """ - Generic function to check if a key is allocating rpm/tpm limits. - Raises an error if we're overallocating. - """ - if keys is not None and len(keys) > 0: - allocated_tpm = sum(key.tpm_limit for key in keys if key.tpm_limit is not None) - allocated_rpm = sum(key.rpm_limit for key in keys if key.rpm_limit is not None) - else: - allocated_tpm = 0 - allocated_rpm = 0 - - if ( - data.tpm_limit is not None - and entity_tpm_limit is not None - and data.tpm_limit + allocated_tpm > entity_tpm_limit - ): - raise HTTPException( - status_code=400, - detail=f"Allocated TPM limit={allocated_tpm} + Key TPM limit={data.tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", - ) - if ( - data.rpm_limit is not None - and entity_rpm_limit is not None - and data.rpm_limit + allocated_rpm > entity_rpm_limit - ): - raise HTTPException( - status_code=400, - detail=f"Allocated RPM limit={allocated_rpm} + Key RPM limit={data.rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", - ) - - -def check_team_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], - team_table: LiteLLM_TeamTableCachedObj, - data: Union[GenerateKeyRequest, UpdateKeyRequest], -) -> None: - """ - Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating. - """ - entity_model_rpm_limit_dict = {} - entity_model_tpm_limit_dict = {} - if team_table.metadata: - entity_model_rpm_limit_dict = team_table.metadata.get("model_rpm_limit", {}) - entity_model_tpm_limit_dict = team_table.metadata.get("model_tpm_limit", {}) - - _check_key_model_specific_limits( - keys=keys, - data=data, - entity_rpm_limit=team_table.rpm_limit, - entity_tpm_limit=team_table.tpm_limit, - entity_model_rpm_limit_dict=entity_model_rpm_limit_dict, - entity_model_tpm_limit_dict=entity_model_tpm_limit_dict, - entity_type="team", - ) - - -def check_team_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], - team_table: LiteLLM_TeamTableCachedObj, - data: Union[GenerateKeyRequest, UpdateKeyRequest], -) -> None: - """ - Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. - """ - _check_key_rpm_tpm_limits( - keys=keys, - data=data, - entity_rpm_limit=team_table.rpm_limit, - entity_tpm_limit=team_table.tpm_limit, - entity_type="team", - ) - - -async def _check_team_key_limits( - team_table: LiteLLM_TeamTableCachedObj, - data: Union[GenerateKeyRequest, UpdateKeyRequest], - prisma_client: PrismaClient, -) -> None: - """ - Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. - - Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" - """ - if ( - data.tpm_limit_type != "guaranteed_throughput" - and data.rpm_limit_type != "guaranteed_throughput" - ): - return - # get all team keys - # calculate allocated tpm/rpm limit - # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - - keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": team_table.team_id}, - ) - check_team_key_model_specific_limits( - keys=keys, - team_table=team_table, - data=data, - ) - check_team_key_rpm_tpm_limits( - keys=keys, - team_table=team_table, - data=data, - ) - - -def check_org_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], - org_table: LiteLLM_OrganizationTable, - data: Union[GenerateKeyRequest, UpdateKeyRequest], -) -> None: - """ - Check if the organization key is allocating model specific limits. If so, raise an error if we're overallocating. - """ - # Get org limits from budget table if available - entity_rpm_limit = None - entity_tpm_limit = None - entity_model_rpm_limit_dict = {} - entity_model_tpm_limit_dict = {} - - if org_table.litellm_budget_table is not None: - entity_rpm_limit = org_table.litellm_budget_table.rpm_limit - entity_tpm_limit = org_table.litellm_budget_table.tpm_limit - - if org_table.metadata: - entity_model_rpm_limit_dict = org_table.metadata.get("model_rpm_limit", {}) - entity_model_tpm_limit_dict = org_table.metadata.get("model_tpm_limit", {}) - - _check_key_model_specific_limits( - keys=keys, - data=data, - entity_rpm_limit=entity_rpm_limit, - entity_tpm_limit=entity_tpm_limit, - entity_model_rpm_limit_dict=entity_model_rpm_limit_dict, - entity_model_tpm_limit_dict=entity_model_tpm_limit_dict, - entity_type="organization", - ) - - -def check_org_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], - org_table: LiteLLM_OrganizationTable, - data: Union[GenerateKeyRequest, UpdateKeyRequest], -) -> None: - """ - Check if the organization key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. - """ - # Get org limits from budget table if available - entity_rpm_limit = None - entity_tpm_limit = None - - if org_table.litellm_budget_table is not None: - entity_rpm_limit = org_table.litellm_budget_table.rpm_limit - entity_tpm_limit = org_table.litellm_budget_table.tpm_limit - - _check_key_rpm_tpm_limits( - keys=keys, - data=data, - entity_rpm_limit=entity_rpm_limit, - entity_tpm_limit=entity_tpm_limit, - entity_type="organization", - ) - - -async def _check_org_key_limits( - org_table: LiteLLM_OrganizationTable, - data: Union[GenerateKeyRequest, UpdateKeyRequest], - prisma_client: PrismaClient, -) -> None: - """ - Check if the organization key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. - - Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" - """ - - rpm_limit_type = getattr(data, "rpm_limit_type", None) or ( - data.metadata.get("rpm_limit_type", None) if data.metadata else None - ) - tpm_limit_type = getattr(data, "tpm_limit_type", None) or ( - data.metadata.get("tpm_limit_type", None) if data.metadata else None - ) - - if ( - tpm_limit_type != "guaranteed_throughput" - and rpm_limit_type != "guaranteed_throughput" - ): - return - # get all organization keys - # calculate allocated tpm/rpm limit - # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"organization_id": org_table.organization_id}, - ) - check_org_key_model_specific_limits( - keys=keys, - org_table=org_table, - data=data, - ) - check_org_key_rpm_tpm_limits( - keys=keys, - org_table=org_table, - data=data, - ) - - -@router.post( - "/key/generate", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], - response_model=GenerateKeyResponse, -) -@management_endpoint_wrapper -async def generate_key_fn( - data: GenerateKeyRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -): - """ - Generate an API key based on the provided data. - - Docs: https://docs.litellm.ai/docs/proxy/virtual_keys - - Parameters: - - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - - key_alias: Optional[str] - User defined key alias - - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. - - team_id: Optional[str] - The team id of the key - - user_id: Optional[str] - The user id of the key - - organization_id: Optional[str] - The organization id of the key. If not set, and team_id is set, the organization id will be the same as the team id. If conflict, an error will be raised. - - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) - - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models - - config: Optional[dict] - any key-specific configs, overrides config in config.yaml - - spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend - - send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key - - max_budget: Optional[float] - Specify max budget for a given key. - - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - - guardrails: Optional[List[str]] - List of active guardrails for the key - - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. - - tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". - - rpm_limit_type: Optional[str] - Type of rpm limit. Options: "best_effort_throughput" (no error if we're overallocating rpm), "guaranteed_throughput" (raise an error if we're overallocating rpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". - - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request - - blocked: Optional[bool] - Whether the key is blocked. - - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) - - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) - - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. - - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) - - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] - - allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored. - - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default". - - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. - - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated) - - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True. - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. - - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings. - - Examples: - - 1. Allow users to turn on/off pii masking - - ```bash - curl --location 'http://0.0.0.0:4000/key/generate' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "permissions": {"allow_pii_controls": true} - }' - ``` - - Returns: - - key: (str) The generated api key - - expires: (datetime) Datetime object for when key expires. - - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. - """ - try: - from litellm.proxy._types import CommonProxyErrors - from litellm.proxy.proxy_server import ( - prisma_client, - user_api_key_cache, - user_custom_key_generate, - ) - - if prisma_client is None: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - verbose_proxy_logger.debug("entered /key/generate") - - # Validate budget values are not negative - if data.max_budget is not None and data.max_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} - ) - if data.soft_budget is not None and data.soft_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} - ) - - if user_custom_key_generate is not None: - if asyncio.iscoroutinefunction(user_custom_key_generate): - result = await user_custom_key_generate(data) # type: ignore - else: - raise ValueError("user_custom_key_generate must be a coroutine") - decision = result.get("decision", True) - message = result.get("message", "Authentication Failed - Custom Auth Rule") - if not decision: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=message - ) - team_table: Optional[LiteLLM_TeamTableCachedObj] = None - if data.team_id is not None: - try: - team_table = await get_team_object( - team_id=data.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - check_db_only=True, - ) - except Exception as e: - verbose_proxy_logger.debug( - f"Error getting team object in `/key/generate`: {e}" - ) - - key_generation_check( - team_table=team_table, - user_api_key_dict=user_api_key_dict, - data=data, - route=KeyManagementRoutes.KEY_GENERATE, - ) - - if team_table is not None: - await _check_team_key_limits( - team_table=team_table, - data=data, - prisma_client=prisma_client, - ) - - return await _common_key_generation_helper( - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - team_table=team_table, - ) - - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {}".format( - str(e) - ) - ) - raise handle_exception_on_proxy(e) - - -@router.post( - "/key/service-account/generate", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], -) -@management_endpoint_wrapper -async def generate_service_account_key_fn( - data: GenerateKeyRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -): - """ - Generate a Service Account API key based on the provided data. This key does not belong to any user. It belongs to the team. - - Why use a service account key? - - Prevent key from being deleted when user is deleted. - - Apply team limits, not team member limits to key. - - Docs: https://docs.litellm.ai/docs/proxy/virtual_keys - - Parameters: - - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - - key_alias: Optional[str] - User defined key alias - - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. - - team_id: Optional[str] - The team id of the key - - user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key - - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) - - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models - - config: Optional[dict] - any key-specific configs, overrides config in config.yaml - - spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend - - send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key - - max_budget: Optional[float] - Specify max budget for a given key. - - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - - guardrails: Optional[List[str]] - List of active guardrails for the key - - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. - - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request - - blocked: Optional[bool] - Whether the key is blocked. - - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) - - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) - - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. - - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) - - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] - - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - Examples: - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. - - - 1. Allow users to turn on/off pii masking - - ```bash - curl --location 'http://0.0.0.0:4000/key/generate' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "permissions": {"allow_pii_controls": true} - }' - ``` - - Returns: - - key: (str) The generated api key - - expires: (datetime) Datetime object for when key expires. - - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. - - """ - from litellm.proxy._types import CommonProxyErrors - from litellm.proxy.proxy_server import ( - prisma_client, - user_api_key_cache, - user_custom_key_generate, - ) - - if prisma_client is None: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - await validate_team_id_used_in_service_account_request( - team_id=data.team_id, - prisma_client=prisma_client, - ) - - verbose_proxy_logger.debug("entered /key/generate") - - if user_custom_key_generate is not None: - if asyncio.iscoroutinefunction(user_custom_key_generate): - result = await user_custom_key_generate(data) # type: ignore - else: - raise ValueError("user_custom_key_generate must be a coroutine") - decision = result.get("decision", True) - message = result.get("message", "Authentication Failed - Custom Auth Rule") - if not decision: - raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) - team_table: Optional[LiteLLM_TeamTableCachedObj] = None - if data.team_id is not None: - try: - team_table = await get_team_object( - team_id=data.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - check_db_only=True, - ) - except Exception as e: - verbose_proxy_logger.debug( - f"Error getting team object in `/key/generate`: {e}" - ) - team_table = None - - if team_table is not None: - await _check_team_key_limits( - team_table=team_table, - data=data, - prisma_client=prisma_client, - ) - - key_generation_check( - team_table=team_table, - user_api_key_dict=user_api_key_dict, - data=data, - route=KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT, - ) - - data.user_id = None # do not allow user_id to be set for service account keys - - return await _common_key_generation_helper( - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - team_table=team_table, - ) - - -def prepare_metadata_fields( - data: BaseModel, non_default_values: dict, existing_metadata: dict -) -> dict: - """ - Check LiteLLM_ManagementEndpoint_MetadataFields (proxy/_types.py) for fields that are allowed to be updated - """ - if "metadata" not in non_default_values: # allow user to set metadata to none - non_default_values["metadata"] = existing_metadata.copy() - - casted_metadata = cast(dict, non_default_values["metadata"]) - - data_json = data.model_dump(exclude_unset=True, exclude_none=True) - - try: - for k, v in data_json.items(): - if k in LiteLLM_ManagementEndpoint_MetadataFields: - if isinstance(v, datetime): - casted_metadata[k] = v.isoformat() - else: - casted_metadata[k] = v - if k in LiteLLM_ManagementEndpoint_MetadataFields_Premium: - from litellm.proxy.utils import _premium_user_check - - _premium_user_check(k) - casted_metadata[k] = v - - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - {}".format( - str(e) - ) - ) - - non_default_values["metadata"] = casted_metadata - return non_default_values - - -async def prepare_key_update_data( - data: Union[UpdateKeyRequest, RegenerateKeyRequest], - existing_key_row: LiteLLM_VerificationToken, -): - data_json: dict = data.model_dump(exclude_unset=True) - data_json.pop("key", None) - data_json.pop("new_key", None) - if ( - data.metadata is not None - and data.metadata.get("service_account_id") is not None - and (data.team_id or existing_key_row.team_id) is None - ): - raise HTTPException( - status_code=400, - detail="team_id is required for service account keys. Please specify `team_id` in the request body.", - ) - non_default_values = {} - # ADD METADATA FIELDS - # Set Management Endpoint Metadata Fields - for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: - if getattr(data, field, None) is not None: - _set_object_metadata_field( - object_data=data, - field_name=field, - value=getattr(data, field), - ) - for k, v in data_json.items(): - if ( - k in LiteLLM_ManagementEndpoint_MetadataFields - or k in LiteLLM_ManagementEndpoint_MetadataFields_Premium - ): - continue - non_default_values[k] = v - - if "duration" in non_default_values: - duration = non_default_values.pop("duration") - if duration == "-1": - # Set expires to None to indicate the key never expires - non_default_values["expires"] = None - elif duration and (isinstance(duration, str)) and len(duration) > 0: - duration_s = duration_in_seconds(duration=duration) - expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - non_default_values["expires"] = expires - - if "budget_duration" in non_default_values: - budget_duration = non_default_values.pop("budget_duration") - if ( - budget_duration - and (isinstance(budget_duration, str)) - and len(budget_duration) > 0 - ): - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - key_reset_at = get_budget_reset_time(budget_duration=budget_duration) - non_default_values["budget_reset_at"] = key_reset_at - non_default_values["budget_duration"] = budget_duration - - if "object_permission" in non_default_values: - non_default_values = await _handle_update_object_permission( - data_json=non_default_values, - existing_key_row=existing_key_row, - ) - - _metadata = existing_key_row.metadata or {} - - # validate model_max_budget - if "model_max_budget" in non_default_values: - validate_model_max_budget(non_default_values["model_max_budget"]) - - # Serialize router_settings to JSON if present - if "router_settings" in non_default_values and non_default_values["router_settings"] is not None: - non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"]) - - non_default_values = prepare_metadata_fields( - data=data, non_default_values=non_default_values, existing_metadata=_metadata - ) - - return non_default_values - - -async def _handle_update_object_permission( - data_json: dict, - existing_key_row: LiteLLM_VerificationToken, -) -> dict: - """ - Handle the update of object permission. - """ - from litellm.proxy.proxy_server import prisma_client - - # Use the common helper to handle the object permission update - object_permission_id = await handle_update_object_permission_common( - data_json=data_json, - existing_object_permission_id=existing_key_row.object_permission_id, - prisma_client=prisma_client, - ) - - # Add the object_permission_id to data_json if one was created/updated - if object_permission_id is not None: - data_json["object_permission_id"] = object_permission_id - verbose_proxy_logger.debug( - f"updated object_permission_id: {object_permission_id}" - ) - - return data_json - - -def is_different_team( - data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken -) -> bool: - if data.team_id is None: - return False - if existing_key_row.team_id is None: - return True - return data.team_id != existing_key_row.team_id - - -def _validate_max_budget(max_budget: Optional[float]) -> None: - """ - Validate that max_budget is not negative. - - Args: - max_budget: The max_budget value to validate - - Raises: - HTTPException: If max_budget is negative - """ - if max_budget is not None and max_budget < 0: - raise HTTPException( - status_code=400, - detail={ - "error": f"max_budget cannot be negative. Received: {max_budget}" - }, - ) - - -async def _get_and_validate_existing_key( - token: str, prisma_client: Optional[PrismaClient] -) -> LiteLLM_VerificationToken: - """ - Get existing key from database and validate it exists. - - Args: - token: The key token to look up - prisma_client: Prisma client instance - - Returns: - LiteLLM_VerificationToken: The existing key row - - Raises: - HTTPException: If key is not found - """ - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": "Database not connected"}, - ) - - existing_key_row = await prisma_client.get_data( - token=token, - table_name="key", - query_type="find_unique", - ) - - if existing_key_row is None: - raise HTTPException( - status_code=404, - detail={"error": f"Key not found: {token}"}, - ) - - return existing_key_row - - -async def _process_single_key_update( - key_update_item: BulkUpdateKeyRequestItem, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str], - prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, - proxy_logging_obj: Any, - llm_router: Optional[Router], -) -> Dict[str, Any]: - """ - Process a single key update with all validations and checks. - - This function encapsulates all the logic for updating a single key, - including validation, permission checks, team checks, and database updates. - - Args: - key_update_item: The key update request item - user_api_key_dict: The authenticated user's API key info - litellm_changed_by: Optional header for tracking who made the change - prisma_client: Prisma client instance - user_api_key_cache: User API key cache - proxy_logging_obj: Proxy logging object - llm_router: LLM router instance - - Returns: - Dict containing the updated key information - - Raises: - HTTPException: For various validation and permission errors - """ - # Validate max_budget - _validate_max_budget(key_update_item.max_budget) - - # Get and validate existing key - existing_key_row = await _get_and_validate_existing_key( - token=key_update_item.key, - prisma_client=prisma_client, - ) - - # Check team member permissions - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=existing_key_row, - user_api_key_cache=user_api_key_cache, - ) - - # Create UpdateKeyRequest from BulkUpdateKeyRequestItem - update_key_request = UpdateKeyRequest( - key=key_update_item.key, - budget_id=key_update_item.budget_id, - max_budget=key_update_item.max_budget, - team_id=key_update_item.team_id, - tags=key_update_item.tags, - ) - - # Get team object and check team limits if team_id is provided - team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - if update_key_request.team_id is not None: - team_obj = await get_team_object( - team_id=update_key_request.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - - if team_obj is not None and prisma_client is not None: - await _check_team_key_limits( - team_table=team_obj, - data=update_key_request, - prisma_client=prisma_client, - ) - - # Validate team change if team is being changed - if is_different_team( - data=update_key_request, existing_key_row=existing_key_row - ): - if llm_router is None: - raise HTTPException( - status_code=400, - detail={ - "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." - }, - ) - if team_obj is None: - raise HTTPException( - status_code=500, - detail={ - "error": "Team object not found for team change validation" - }, - ) - validate_key_team_change( - key=existing_key_row, - team=team_obj, - change_initiated_by=user_api_key_dict, - llm_router=llm_router, - ) - - # Prepare update data - non_default_values = await prepare_key_update_data( - data=update_key_request, existing_key_row=existing_key_row - ) - - # Update key in database - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": "Database not connected"}, - ) - - _data = {**non_default_values, "token": key_update_item.key} - response = await prisma_client.update_data( - token=key_update_item.key, data=_data - ) - - # Delete cache - await _delete_cache_key_object( - hashed_token=hash_token(key_update_item.key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - # Trigger async hook - asyncio.create_task( - KeyManagementEventHooks.async_key_updated_hook( - data=update_key_request, - existing_key_row=existing_key_row, - response=response, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - ) - - if response is None: - raise ValueError("Failed to update key got response = None") - - # Extract and format updated key info - updated_key_info = response.get("data", {}) - if hasattr(updated_key_info, "model_dump"): - updated_key_info = updated_key_info.model_dump() - elif hasattr(updated_key_info, "dict"): - updated_key_info = updated_key_info.dict() - - updated_key_info.pop("token", None) - - return updated_key_info - - -@router.post( - "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) -@management_endpoint_wrapper -async def update_key_fn( - request: Request, - data: UpdateKeyRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -): - """ - Update an existing API key's parameters. - - Parameters: - - key: str - The key to update - - key_alias: Optional[str] - User-friendly key alias - - user_id: Optional[str] - User ID associated with key - - team_id: Optional[str] - Team ID associated with key - - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - - models: Optional[list] - Model_name's a user is allowed to call - - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) - - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) - - spend: Optional[float] - Amount spent by key - - max_budget: Optional[float] - Max budget for key - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. - - max_parallel_requests: Optional[int] - Rate limit for parallel requests - - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} - - tpm_limit: Optional[int] - Tokens per minute limit - - rpm_limit: Optional[int] - Requests per minute limit - - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} - - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} - - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - - allowed_cache_controls: Optional[list] - List of allowed cache control values - - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) or "-1" to never expire - - permissions: Optional[dict] - Key-specific permissions - - send_invite_email: Optional[bool] - Send invite email to user_id - - guardrails: Optional[List[str]] - List of active guardrails for the key - - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - - blocked: Optional[bool] - Whether the key is blocked - - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) - - config: Optional[dict] - [DEPRECATED PARAM] Key-specific config. - - temp_budget_increase: Optional[float] - Temporary budget increase for the key (Enterprise only). - - temp_budget_expiry: Optional[str] - Expiry time for the temporary budget increase (Enterprise only). - - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] - - allowed_passthrough_routes: Optional[list] - List of allowed pass through routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through routes the key can access, without specifying the routes. If allowed_routes is specified, allowed_passthrough_routes is ignored. - - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. - - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - - auto_rotate: Optional[bool] - Whether this key should be automatically rotated - - rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. - - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings. - - Example: - ```bash - curl --location 'http://0.0.0.0:4000/key/update' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "key": "sk-1234", - "key_alias": "my-key", - "user_id": "user-1234", - "team_id": "team-1234", - "max_budget": 100, - "metadata": {"any_key": "any-val"}, - }' - ``` - """ - from litellm.proxy.proxy_server import ( - llm_router, - premium_user, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - try: - # Validate budget values are not negative - if data.max_budget is not None and data.max_budget < 0: - raise HTTPException( - status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} - ) - - data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) - key = data_json.pop("key") - - # get the row from db - if prisma_client is None: - raise Exception("Not connected to DB!") - - existing_key_row = await prisma_client.get_data( - token=data.key, table_name="key", query_type="find_unique" - ) - - if existing_key_row is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team not found, passed team_id={data.team_id}"}, - ) - - ## sanity check - prevent non-proxy admin user from updating key to belong to a different user - if ( - data.user_id is not None - and data.user_id != existing_key_row.user_id - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - ): - raise HTTPException( - status_code=403, - detail=f"User={data.user_id} is not allowed to update key={key} to belong to user={existing_key_row.user_id}", - ) - - common_key_access_checks( - user_api_key_dict=user_api_key_dict, - data=data, - user_id=existing_key_row.user_id, - llm_router=llm_router, - premium_user=premium_user, - ) - - # check if user has permission to update key - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=existing_key_row, - user_api_key_cache=user_api_key_cache, - ) - - # Only check team limits if key has a team_id - team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - if data.team_id is not None: - team_obj = await get_team_object( - team_id=data.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - - if team_obj is not None: - await _check_team_key_limits( - team_table=team_obj, - data=data, - prisma_client=prisma_client, - ) - - # if team change - check if this is possible - if is_different_team(data=data, existing_key_row=existing_key_row): - if llm_router is None: - raise HTTPException( - status_code=400, - detail={ - "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." - }, - ) - # team_obj should be set since is_different_team() returns True only when data.team_id is not None - if team_obj is None: - raise HTTPException( - status_code=500, - detail={ - "error": "Team object not found for team change validation" - }, - ) - validate_key_team_change( - key=existing_key_row, - team=team_obj, - change_initiated_by=user_api_key_dict, - llm_router=llm_router, - ) - - # Set Management Endpoint Metadata Fields - - non_default_values = await prepare_key_update_data( - data=data, existing_key_row=existing_key_row - ) - - await _enforce_unique_key_alias( - key_alias=non_default_values.get("key_alias", None), - prisma_client=prisma_client, - existing_key_token=existing_key_row.token, - ) - - # Handle rotation fields if auto_rotate is being enabled - _set_key_rotation_fields( - non_default_values, - non_default_values.get("auto_rotate", False), - non_default_values.get("rotation_interval"), - ) - - _data = {**non_default_values, "token": key} - response = await prisma_client.update_data(token=key, data=_data) - - # Delete - key from cache, since it's been updated! - # key updated - a new model could have been added to this key. it should not block requests after this is done - await _delete_cache_key_object( - hashed_token=hash_token(key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - asyncio.create_task( - KeyManagementEventHooks.async_key_updated_hook( - data=data, - existing_key_row=existing_key_row, - response=response, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - ) - - if response is None: - raise ValueError("Failed to update key got response = None") - - return {"key": key, **response["data"]} - # update based on remaining passed in values - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.update_key_fn(): Exception occured - {}".format( - str(e) - ) - ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) - - -@router.post( - "/key/bulk_update", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], - response_model=BulkUpdateKeyResponse, -) -@management_endpoint_wrapper -async def bulk_update_keys( - data: BulkUpdateKeyRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -): - """ - Bulk update multiple keys at once. - - This endpoint allows updating multiple keys in a single request. Each key update - is processed independently - if some updates fail, others will still succeed. - - Parameters: - - keys: List[BulkUpdateKeyRequestItem] - List of key update requests, each containing: - - key: str - The key identifier (token) to update - - budget_id: Optional[str] - Budget ID associated with the key - - max_budget: Optional[float] - Max budget for key - - team_id: Optional[str] - Team ID associated with key - - tags: Optional[List[str]] - Tags for organizing keys - - Returns: - - total_requested: int - Total number of keys requested for update - - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info - - failed_updates: List[FailedKeyUpdate] - List of failed updates with key_info and failed_reason - - Example request: - ```bash - curl --location 'http://0.0.0.0:4000/key/bulk_update' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "keys": [ - { - "key": "sk-1234", - "max_budget": 100.0, - "team_id": "team-123", - "tags": ["production", "api"] - }, - { - "key": "sk-5678", - "budget_id": "budget-456", - "tags": ["staging"] - } - ] - }' - ``` - """ - from litellm.proxy.proxy_server import ( - llm_router, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: - raise HTTPException( - status_code=403, - detail={ - "error": "Only proxy admins can perform bulk key updates" - }, - ) - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": "Database not connected"}, - ) - - if not data.keys: - raise HTTPException( - status_code=400, - detail={"error": "No keys provided for update"}, - ) - - MAX_BATCH_SIZE = 500 - if len(data.keys) > MAX_BATCH_SIZE: - raise HTTPException( - status_code=400, - detail={ - "error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys." - }, - ) - - successful_updates: List[SuccessfulKeyUpdate] = [] - failed_updates: List[FailedKeyUpdate] = [] - - for key_update_item in data.keys: - try: - # Process single key update using reusable function - updated_key_info = await _process_single_key_update( - key_update_item=key_update_item, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - ) - - successful_updates.append( - SuccessfulKeyUpdate( - key=key_update_item.key, - key_info=updated_key_info, - ) - ) - - except Exception as e: - verbose_proxy_logger.exception( - f"Failed to update key {key_update_item.key}: {e}" - ) - - if isinstance(e, HTTPException): - error_detail = e.detail - if isinstance(error_detail, dict): - error_message = error_detail.get("error", str(e)) - else: - error_message = str(error_detail) - else: - error_message = str(e) - - key_info = None - try: - existing_key_row = await prisma_client.get_data( - token=key_update_item.key, - table_name="key", - query_type="find_unique", - ) - if existing_key_row is not None: - if hasattr(existing_key_row, "model_dump"): - key_info = existing_key_row.model_dump() - elif hasattr(existing_key_row, "dict"): - key_info = existing_key_row.dict() - if key_info: - key_info.pop("token", None) - except Exception: - pass - - failed_updates.append( - FailedKeyUpdate( - key=key_update_item.key, - key_info=key_info, - failed_reason=error_message, - ) - ) - - return BulkUpdateKeyResponse( - total_requested=len(data.keys), - successful_updates=successful_updates, - failed_updates=failed_updates, - ) - - -def validate_key_team_change( - key: LiteLLM_VerificationToken, - team: LiteLLM_TeamTable, - change_initiated_by: UserAPIKeyAuth, - llm_router: Router, -): - """ - Validate that a key can be moved to a new team. - - - The team must have access to the key's models - - The key's user_id must be a member of the team - - The key's tpm/rpm limit must be less than the team's tpm/rpm limit - - The person initiating the change must be either Proxy Admin or Team Admin - """ - # Check if the team has access to the key's models - if len(key.models) > 0: - for model in key.models: - can_team_access_model( - model=model, - team_object=team, - llm_router=llm_router, - ) - - # Check if the key's user_id is a member of the team - member_object = _get_user_in_team( - team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id - ) - if key.user_id is not None: - if not member_object: - raise HTTPException( - status_code=403, - detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", - ) - - # Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit - if key.tpm_limit is not None: - if team.tpm_limit and key.tpm_limit > team.tpm_limit: - raise HTTPException( - status_code=403, - detail=f"Key={key.token} has a tpm_limit={key.tpm_limit} which is greater than the team's tpm_limit={team.tpm_limit}.", - ) - if team.rpm_limit and key.rpm_limit and key.rpm_limit > team.rpm_limit: - raise HTTPException( - status_code=403, - detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.", - ) - - # Check if the person initiating the change is a Proxy Admin or Team Admin - if change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return - elif _is_user_team_admin( - user_api_key_dict=change_initiated_by, - team_obj=team, - ): - return - # this teams member permissions allow updating a - elif TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( - team_member_object=member_object, - team_table=cast(LiteLLM_TeamTableCachedObj, team), - route=KeyManagementRoutes.KEY_UPDATE.value, - ): - return - else: - raise HTTPException( - status_code=403, - detail=f"User={change_initiated_by.user_id} is not a Proxy Admin or Team Admin for team={team.team_id}. Please ask your Proxy Admin to allow this action under 'Member Permissions' for this team.", - ) - - -@router.post( - "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) -@management_endpoint_wrapper -async def delete_key_fn( - data: KeyRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -): - """ - Delete a key from the key management system. - - Parameters:: - - keys (List[str]): A list of keys or hashed keys to delete. Example {"keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]} - - key_aliases (List[str]): A list of key aliases to delete. Can be passed instead of `keys`.Example {"key_aliases": ["alias1", "alias2"]} - - Returns: - - deleted_keys (List[str]): A list of deleted keys. Example {"deleted_keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]} - - Example: - ```bash - curl --location 'http://0.0.0.0:4000/key/delete' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "keys": ["sk-QWrxEynunsNpV1zT48HIrw"] - }' - ``` - - Raises: - HTTPException: If an error occurs during key deletion. - """ - try: - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - if prisma_client is None: - raise Exception("Not connected to DB!") - - # Normalize litellm_changed_by: if it's a Header object or not a string, convert to None - if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): - litellm_changed_by = None - - ## only allow user to delete keys they own - verbose_proxy_logger.debug( - f"user_api_key_dict.user_role: {user_api_key_dict.user_role}" - ) - - num_keys_to_be_deleted = 0 - deleted_keys = [] - if data.keys: - number_deleted_keys, _keys_being_deleted = await delete_verification_tokens( - tokens=data.keys, - user_api_key_cache=user_api_key_cache, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - num_keys_to_be_deleted = len(data.keys) - deleted_keys = data.keys - elif data.key_aliases: - number_deleted_keys, _keys_being_deleted = await delete_key_aliases( - key_aliases=data.key_aliases, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - num_keys_to_be_deleted = len(data.key_aliases) - deleted_keys = data.key_aliases - else: - raise ValueError("Invalid request type") - - if number_deleted_keys is None: - raise ProxyException( - message="Failed to delete keys got None response from delete_verification_token", - type=ProxyErrorTypes.internal_server_error, - param="keys", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - verbose_proxy_logger.debug(f"/key/delete - deleted_keys={number_deleted_keys}") - - try: - assert num_keys_to_be_deleted == len(deleted_keys) - except Exception: - raise HTTPException( - status_code=400, - detail={ - "error": f"Not all keys passed in were deleted. This probably means you don't have access to delete all the keys passed in. Keys passed in={num_keys_to_be_deleted}, Deleted keys ={number_deleted_keys}" - }, - ) - - verbose_proxy_logger.debug( - f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}" - ) - - asyncio.create_task( - KeyManagementEventHooks.async_key_deleted_hook( - data=data, - keys_being_deleted=_keys_being_deleted, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - response=number_deleted_keys, - ) - ) - - return {"deleted_keys": deleted_keys} - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {}".format( - str(e) - ) - ) - raise handle_exception_on_proxy(e) - - -@router.post( - "/v2/key/info", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], - include_in_schema=False, -) -async def info_key_fn_v2( - data: Optional[KeyRequest] = None, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Retrieve information about a list of keys. - - **New endpoint**. Currently admin only. - Parameters: - keys: Optional[list] = body parameter representing the key(s) in the request - user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key - Returns: - Dict containing the key and its associated information - - Example Curl: - ``` - curl -X GET "http://0.0.0.0:4000/key/info" \ - -H "Authorization: Bearer sk-1234" \ - -d {"keys": ["sk-1", "sk-2", "sk-3"]} - ``` - """ - from litellm.proxy.proxy_server import prisma_client - - try: - if prisma_client is None: - raise Exception( - "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" - ) - if data is None: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={"message": "Malformed request. No keys passed in."}, - ) - - key_info = await prisma_client.get_data( - token=data.keys, table_name="key", query_type="find_all" - ) - if key_info is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"message": "No keys found"}, - ) - filtered_key_info = [] - for k in key_info: - try: - k = k.model_dump() # noqa - except Exception: - # if using pydantic v1 - k = k.dict() - filtered_key_info.append(k) - return {"key": data.keys, "info": filtered_key_info} - - except Exception as e: - raise handle_exception_on_proxy(e) - - -@router.get( - "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) -async def info_key_fn( - key: Optional[str] = fastapi.Query( - default=None, description="Key in the request parameters" - ), - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Retrieve information about a key. - Parameters: - key: Optional[str] = Query parameter representing the key in the request - user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key - Returns: - Dict containing the key and its associated information - - Example Curl: - ``` - curl -X GET "http://0.0.0.0:4000/key/info?key=sk-test-example-key-123" \ --H "Authorization: Bearer sk-1234" - ``` - - Example Curl - if no key is passed, it will use the Key Passed in Authorization Header - ``` - curl -X GET "http://0.0.0.0:4000/key/info" \ --H "Authorization: Bearer sk-test-example-key-123" - ``` - """ - from litellm.proxy.proxy_server import prisma_client - - try: - if prisma_client is None: - raise Exception( - "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" - ) - - # default to using Auth token if no key is passed in - key = key or user_api_key_dict.api_key - hashed_key: Optional[str] = key - if key is not None: - hashed_key = _hash_token_if_needed(token=key) - key_info = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_key}, # type: ignore - include={"litellm_budget_table": True}, - ) - if key_info is None: - raise ProxyException( - message="Key not found in database", - type=ProxyErrorTypes.not_found_error, - param="key", - code=status.HTTP_404_NOT_FOUND, - ) - - if ( - await _can_user_query_key_info( - user_api_key_dict=user_api_key_dict, - key=key, - key_info=key_info, - ) - is not True - ): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="You are not allowed to access this key's info. Your role={}".format( - user_api_key_dict.user_role - ), - ) - ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## - try: - key_info = key_info.model_dump() # noqa - except Exception: - # if using pydantic v1 - key_info = key_info.dict() - key_info.pop("token") - return {"key": key, "info": key_info} - except Exception as e: - raise handle_exception_on_proxy(e) - - -def _check_model_access_group( - models: Optional[List[str]], llm_router: Optional[Router], premium_user: bool -) -> Literal[True]: - """ - if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user - - Return True if user is a premium user, False otherwise - """ - if models is None or llm_router is None: - return True - - for model in models: - if llm_router._is_model_access_group_for_wildcard_route( - model_access_group=model - ): - if not premium_user: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Setting a model access group on a wildcard model is only available for LiteLLM Enterprise users.{}".format( - CommonProxyErrors.not_premium_user.value - ) - }, - ) - - return True - - -async def generate_key_helper_fn( # noqa: PLR0915 - request_type: Literal[ - "user", "key" - ], # identifies if this request is from /user/new or /key/generate - duration: Optional[str] = None, - models: list = [], - aliases: dict = {}, - config: dict = {}, - spend: float = 0.0, - key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key - key_budget_duration: Optional[str] = None, - budget_id: Optional[float] = None, # budget id <-> LiteLLM_BudgetTable - soft_budget: Optional[ - float - ] = None, # soft_budget is used to set soft Budgets Per user - max_budget: Optional[float] = None, # max_budget is used to Budget Per user - blocked: Optional[bool] = None, - budget_duration: Optional[str] = None, # max_budget is used to Budget Per user - token: Optional[str] = None, - key: Optional[ - str - ] = None, # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself. - user_id: Optional[str] = None, - user_alias: Optional[str] = None, - team_id: Optional[str] = None, - user_email: Optional[str] = None, - user_role: Optional[str] = None, - max_parallel_requests: Optional[int] = None, - metadata: Optional[dict] = {}, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, - query_type: Literal["insert_data", "update_data"] = "insert_data", - update_key_values: Optional[dict] = None, - key_alias: Optional[str] = None, - allowed_cache_controls: Optional[list] = [], - permissions: Optional[dict] = {}, - model_max_budget: Optional[dict] = {}, - model_rpm_limit: Optional[dict] = None, - model_tpm_limit: Optional[dict] = None, - guardrails: Optional[list] = None, - policies: Optional[list] = None, - prompts: Optional[list] = None, - teams: Optional[list] = None, - organization_id: Optional[str] = None, - table_name: Optional[Literal["key", "user"]] = None, - send_invite_email: Optional[bool] = None, - created_by: Optional[str] = None, - updated_by: Optional[str] = None, - allowed_routes: Optional[list] = None, - sso_user_id: Optional[str] = None, - object_permission_id: Optional[ - str - ] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, - auto_rotate: Optional[bool] = None, - rotation_interval: Optional[str] = None, - router_settings: Optional[dict] = None, -): - from litellm.proxy.proxy_server import premium_user, prisma_client - - if prisma_client is None: - raise Exception( - "Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys " - ) - - if token is None: - if key is not None: - token = key - else: - token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" - - if duration is None: # allow tokens that never expire - expires = None - else: - # Add duration to current time for exact expiration (not standardized reset time) - duration_seconds = duration_in_seconds(duration) - expires = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds) - - if key_budget_duration is None: # one-time budget - key_reset_at = None - else: - key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration) - - if budget_duration is None: # one-time budget - reset_at = None - else: - reset_at = get_budget_reset_time(budget_duration=budget_duration) - - aliases_json = json.dumps(aliases) - config_json = json.dumps(config) - permissions_json = json.dumps(permissions) - router_settings_json = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) - - # Add model_rpm_limit and model_tpm_limit to metadata - if model_rpm_limit is not None: - metadata = metadata or {} - metadata["model_rpm_limit"] = model_rpm_limit - if model_tpm_limit is not None: - metadata = metadata or {} - metadata["model_tpm_limit"] = model_tpm_limit - if guardrails is not None: - metadata = metadata or {} - metadata["guardrails"] = guardrails - if policies is not None: - metadata = metadata or {} - metadata["policies"] = policies - if prompts is not None: - metadata = metadata or {} - metadata["prompts"] = prompts - - metadata_json = json.dumps(metadata) - validate_model_max_budget(model_max_budget) - model_max_budget_json = json.dumps(model_max_budget) - user_role = user_role - tpm_limit = tpm_limit - rpm_limit = rpm_limit - allowed_cache_controls = allowed_cache_controls - - try: - # Create a new verification token (you may want to enhance this logic based on your needs) - - user_data = { - "max_budget": max_budget, - "user_email": user_email, - "user_id": user_id, - "user_alias": user_alias, - "team_id": team_id, - "organization_id": organization_id, - "user_role": user_role, - "spend": spend, - "models": models, - "metadata": metadata_json, - "max_parallel_requests": max_parallel_requests, - "tpm_limit": tpm_limit, - "rpm_limit": rpm_limit, - "budget_duration": budget_duration, - "budget_reset_at": reset_at, - "allowed_cache_controls": allowed_cache_controls, - "sso_user_id": sso_user_id, - "object_permission_id": object_permission_id, - } - if teams is not None: - user_data["teams"] = teams - key_data = { - "token": token, - "key_alias": key_alias, - "expires": expires, - "models": models, - "aliases": aliases_json, - "config": config_json, - "spend": spend, - "max_budget": key_max_budget, - "user_id": user_id, - "team_id": team_id, - "max_parallel_requests": max_parallel_requests, - "metadata": metadata_json, - "tpm_limit": tpm_limit, - "rpm_limit": rpm_limit, - "budget_duration": key_budget_duration, - "budget_reset_at": key_reset_at, - "allowed_cache_controls": allowed_cache_controls, - "permissions": permissions_json, - "model_max_budget": model_max_budget_json, - "organization_id": organization_id, - "budget_id": budget_id, - "blocked": blocked, - "created_by": created_by, - "updated_by": updated_by, - "allowed_routes": allowed_routes or [], - "object_permission_id": object_permission_id, - "router_settings": router_settings_json, - } - - # Add rotation fields if auto_rotate is enabled - _set_key_rotation_fields( - data=key_data, - auto_rotate=auto_rotate or False, - rotation_interval=rotation_interval, - ) - - if ( - get_secret("DISABLE_KEY_NAME", False) is True - ): # allow user to disable storing abbreviated key name (shown in UI, to help figure out which key spent how much) - pass - else: - key_data["key_name"] = abbreviate_api_key(api_key=token) - saved_token = copy.deepcopy(key_data) - if isinstance(saved_token["aliases"], str): - saved_token["aliases"] = json.loads(saved_token["aliases"]) - if isinstance(saved_token["config"], str): - saved_token["config"] = json.loads(saved_token["config"]) - if isinstance(saved_token["metadata"], str): - saved_token["metadata"] = json.loads(saved_token["metadata"]) - if isinstance(saved_token["permissions"], str): - if ( - "get_spend_routes" in saved_token["permissions"] - and premium_user is not True - ): - raise ValueError( - "get_spend_routes permission is only available for LiteLLM Enterprise users" - ) - - saved_token["permissions"] = json.loads(saved_token["permissions"]) - if isinstance(saved_token["model_max_budget"], str): - saved_token["model_max_budget"] = json.loads( - saved_token["model_max_budget"] - ) - router_settings = cast(Optional[dict], saved_token.get("router_settings")) - if router_settings is not None and isinstance(router_settings, str): - try: - saved_token["router_settings"] = yaml.safe_load(router_settings) - except yaml.YAMLError: - # If it's not valid JSON/YAML, keep as is or set to empty dict - saved_token["router_settings"] = {} - - if saved_token.get("expires", None) is not None and isinstance( - saved_token["expires"], datetime - ): - saved_token["expires"] = saved_token["expires"].isoformat() - if prisma_client is not None: - if ( - table_name is None or table_name == "user" - ): # do not auto-create users for `/key/generate` - ## CREATE USER (If necessary) - if query_type == "insert_data": - user_row = await prisma_client.insert_data( - data=user_data, table_name="user" - ) - - if user_row is None: - raise Exception("Failed to create user") - ## use default user model list if no key-specific model list provided - if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore - key_data["models"] = user_row.models # type: ignore - elif query_type == "update_data": - user_row = await prisma_client.update_data( - data=user_data, - table_name="user", - update_key_values=update_key_values, - ) - if table_name is not None and table_name == "user": - # do not create a key if table name is set to just 'user' - # we only need to ensure this exists in the user table - # the LiteLLM_VerificationToken table will increase in size if we don't do this check - return user_data - - ## CREATE KEY - verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data) - create_key_response = await prisma_client.insert_data( - data=key_data, table_name="key" - ) - - key_data["token_id"] = getattr(create_key_response, "token", None) - key_data["litellm_budget_table"] = getattr( - create_key_response, "litellm_budget_table", None - ) - key_data["created_at"] = getattr(create_key_response, "created_at", None) - key_data["updated_at"] = getattr(create_key_response, "updated_at", None) - - # Deserialize router_settings from JSON string to dict for response - router_settings_value = key_data.get("router_settings") - if router_settings_value is not None and isinstance(router_settings_value, str): - try: - key_data["router_settings"] = yaml.safe_load(router_settings_value) - except yaml.YAMLError: - # If it's not valid JSON/YAML, keep as is or set to empty dict - key_data["router_settings"] = {} - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise e - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": "Internal Server Error."}, - ) - - # Add budget related info in key_data - this ensures it's returned - key_data["budget_id"] = budget_id - - if request_type == "user": - # if this is a /user/new request update the key_date with user_data fields - key_data.update(user_data) - - return key_data - - -async def _team_key_deletion_check( - user_api_key_dict: UserAPIKeyAuth, - key_info: LiteLLM_VerificationToken, - prisma_client: PrismaClient, - user_api_key_cache: DualCache, -): - is_team_key = _is_team_key(data=key_info) - - if is_team_key and key_info.team_id is not None: - team_table = await get_team_object( - team_id=key_info.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - if ( - litellm.key_generation_settings is not None - and "team_key_generation" in litellm.key_generation_settings - ): - _team_key_generation = litellm.key_generation_settings[ - "team_key_generation" - ] - else: - _team_key_generation = TeamUIKeyGenerationConfig( - allowed_team_member_roles=["admin", "user"], - ) - # check if user is team admin - if team_table is not None: - return _team_key_operation_team_member_check( - assigned_user_id=user_api_key_dict.user_id, - team_table=team_table, - user_api_key_dict=user_api_key_dict, - team_key_generation=_team_key_generation, - route=KeyManagementRoutes.KEY_DELETE, - ) - else: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={ - "error": f"Team not found in db, and user not proxy admin. Team id = {key_info.team_id}" - }, - ) - return False - - -async def can_modify_verification_token( - key_info: LiteLLM_VerificationToken, - user_api_key_cache: DualCache, - user_api_key_dict: UserAPIKeyAuth, - prisma_client: PrismaClient, -) -> bool: - """ - Check if user has permission to modify (delete/regenerate) a verification token. - - Rules: - - Proxy admin can modify any key - - For team keys: only team admin or key owner can modify - - For personal keys: only key owner can modify - - Args: - key_info: The verification token to check - user_api_key_cache: Cache for user API keys - user_api_key_dict: The user making the request - prisma_client: Prisma client for database access - - Returns: - True if user can modify the key, False otherwise - """ - is_team_key = _is_team_key(data=key_info) - - # 1. Proxy admin can modify any key - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return True - - # 2. For team keys: only team admin or key owner can modify - if is_team_key and key_info.team_id is not None: - # Get team object to check if user is team admin - team_table = await get_team_object( - team_id=key_info.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - - if team_table is None: - return False - - # Check if user is team admin - if _is_user_team_admin( - user_api_key_dict=user_api_key_dict, - team_obj=team_table, - ): - return True - - # Check if the key belongs to the user (they own it) - if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: - return True - - # Not team admin and doesn't own the key - return False - - # 3. For personal keys: only key owner can modify - if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: - return True - - # Default: deny - return False - - - - -async def delete_verification_tokens( - tokens: List, - user_api_key_cache: DualCache, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - """ - Helper that deletes the list of tokens from the database - - - check if user is proxy admin - - check if user is team admin and key is a team key - - Args: - tokens: List of tokens to delete - user_id: Optional user_id to filter by - - Returns: - Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - Optional[Dict]: - - Number of deleted tokens - List[LiteLLM_VerificationToken]: - - List of keys being deleted, this contains information about the key_alias, token, and user_id being deleted, - this is passed down to the KeyManagementEventHooks to delete the keys from the secret manager and handle audit logs - """ - from litellm.proxy.proxy_server import prisma_client - - try: - if prisma_client: - tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} - ) - ) - - if len(_keys_being_deleted) == 0: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": "No keys found"}, - ) - - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - authorized_keys = _keys_being_deleted - else: - authorized_keys = [] - for key in _keys_being_deleted: - if await can_modify_verification_token( - key_info=key, - user_api_key_cache=user_api_key_cache, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ): - authorized_keys.append(key) - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "You are not authorized to delete this key" - }, - ) - await _persist_deleted_verification_tokens( - keys=authorized_keys, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - deleted_tokens = await prisma_client.delete_data(tokens=tokens) - else: - deletion_tasks = [ - prisma_client.delete_data(tokens=[key.token]) - for key in authorized_keys - ] - await asyncio.gather(*deletion_tasks) - - deleted_tokens = [key.token for key in authorized_keys] - if len(deleted_tokens) != len(tokens): - failed_tokens = [ - token for token in tokens if token not in deleted_tokens - ] - raise Exception( - "Failed to delete all tokens. Failed to delete tokens: " - + str(failed_tokens) - ) - else: - raise Exception("DB not connected. prisma_client is None") - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - raise e - - for key in tokens: - user_api_key_cache.delete_cache(key) - # remove hash token from cache - hashed_token = hash_token(cast(str, key)) - user_api_key_cache.delete_cache(hashed_token) - - return {"deleted_keys": deleted_tokens}, _keys_being_deleted - - -def _transform_verification_tokens_to_deleted_records( - keys: List[LiteLLM_VerificationToken], - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: - """Transform verification tokens into deleted token records ready for persistence.""" - if not keys: - return [] - - deleted_at = datetime.now(timezone.utc) - records = [] - for key in keys: - key_payload = key.model_dump() - deleted_record = LiteLLM_DeletedVerificationToken( - **key_payload, - deleted_at=deleted_at, - deleted_by=user_api_key_dict.user_id, - deleted_by_api_key=user_api_key_dict.api_key, - litellm_changed_by=litellm_changed_by, - ) - record = deleted_record.model_dump() - - # Map org_id to organization_id (model uses org_id, but schema expects organization_id) - org_id_value = record.pop("org_id", None) - if org_id_value is not None: - record["organization_id"] = org_id_value - - for json_field in ["aliases", "config", "permissions", "metadata", "model_spend", "model_max_budget", "router_settings"]: - if json_field in record and record[json_field] is not None: - record[json_field] = json.dumps(record[json_field]) - - for rel_key in ("litellm_budget_table", "litellm_organization_table", "object_permission", "id"): - record.pop(rel_key, None) - - records.append(record) - - return records - - -async def _save_deleted_verification_token_records( - records: List[Dict[str, Any]], - prisma_client: PrismaClient, -) -> None: - """Save deleted verification token records to the database.""" - if not records: - return - await prisma_client.db.litellm_deletedverificationtoken.create_many( - data=records - ) - - -async def _persist_deleted_verification_tokens( - keys: List[LiteLLM_VerificationToken], - prisma_client: PrismaClient, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> None: - """Persist deleted verification token records by transforming and saving them.""" - records = _transform_verification_tokens_to_deleted_records( - keys=keys, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - await _save_deleted_verification_token_records( - records=records, - prisma_client=prisma_client, - ) - - -async def delete_key_aliases( - key_aliases: List[str], - user_api_key_cache: DualCache, - prisma_client: PrismaClient, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - _keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many( - where={"key_alias": {"in": key_aliases}} - ) - - tokens = [key.token for key in _keys_being_deleted] - return await delete_verification_tokens( - tokens=tokens, - user_api_key_cache=user_api_key_cache, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - - -async def _rotate_master_key( - prisma_client: PrismaClient, - user_api_key_dict: UserAPIKeyAuth, - current_master_key: str, - new_master_key: str, -) -> None: - """ - Rotate the master key - - 1. Get the values from the DB - - Get models from DB - - Get config from DB - 2. Decrypt the values - - ModelTable - - [{"model_name": "str", "litellm_params": {}}] - - ConfigTable - 3. Encrypt the values with the new master key - 4. Update the values in the DB - """ - from litellm.proxy.proxy_server import proxy_config - - try: - models: Optional[List] = ( - await prisma_client.db.litellm_proxymodeltable.find_many() - ) - except Exception: - models = None - # 2. process model table - if models: - decrypted_models = proxy_config.decrypt_model_list_from_db(new_models=models) - verbose_proxy_logger.debug( - "ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models) - ) - new_models = [] - for model in decrypted_models: - new_model = await _add_model_to_db( - model_params=Deployment(**model), - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - new_encryption_key=new_master_key, - should_create_model_in_db=False, - ) - if new_model: - new_models.append(jsonify_object(new_model.model_dump())) - verbose_proxy_logger.debug("Resetting proxy model table") - await prisma_client.db.litellm_proxymodeltable.delete_many() - verbose_proxy_logger.debug("Creating %s models", len(new_models)) - await prisma_client.db.litellm_proxymodeltable.create_many( - data=new_models, - ) - # 3. process config table - try: - config = await prisma_client.db.litellm_config.find_many() - except Exception: - config = None - - if config: - """If environment_variables is found, decrypt it and encrypt it with the new master key""" - environment_variables_dict = {} - for c in config: - if c.param_name == "environment_variables": - environment_variables_dict = c.param_value - - if environment_variables_dict: - decrypted_env_vars = proxy_config._decrypt_and_set_db_env_variables( - environment_variables=environment_variables_dict - ) - encrypted_env_vars = proxy_config._encrypt_env_variables( - environment_variables=decrypted_env_vars, - new_encryption_key=new_master_key, - ) - - if encrypted_env_vars: - await prisma_client.db.litellm_config.update( - where={"param_name": "environment_variables"}, - data={"param_value": jsonify_object(encrypted_env_vars)}, - ) - - # 4. process MCP server table - await rotate_mcp_server_credentials_master_key( - prisma_client=prisma_client, - touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, - new_master_key=new_master_key, - ) - - # 5. process credentials table - try: - credentials = await prisma_client.db.litellm_credentialstable.find_many() - except Exception: - credentials = None - if credentials: - from litellm.proxy.credential_endpoints.endpoints import update_db_credential - - for cred in credentials: - try: - decrypted_cred = proxy_config.decrypt_credentials(cred) - encrypted_cred = update_db_credential( - db_credential=cred, - updated_patch=decrypted_cred, - new_encryption_key=new_master_key, - ) - credential_object_jsonified = jsonify_object(encrypted_cred.model_dump()) - await prisma_client.db.litellm_credentialstable.update( - where={"credential_name": cred.credential_name}, - data={ - **credential_object_jsonified, - "updated_by": user_api_key_dict.user_id, - }, - ) - except Exception as e: - verbose_proxy_logger.error( - f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}" - ) - # Continue with next credential instead of failing entire rotation - continue - verbose_proxy_logger.debug( - f"Successfully re-encrypted {len(credentials)} credentials with new master key" - ) - - -def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: - if data and data.new_key is not None: - new_token = data.new_key - if not data.new_key.startswith("sk-"): - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key." - }, - ) - else: - new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" - return new_token - - -@router.post( - "/key/{key:path}/regenerate", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], -) -@router.post( - "/key/regenerate", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], -) -@management_endpoint_wrapper -async def regenerate_key_fn( - key: Optional[str] = None, - data: Optional[RegenerateKeyRequest] = None, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -) -> Optional[GenerateKeyResponse]: - """ - Regenerate an existing API key while optionally updating its parameters. - - Parameters: - - key: str (path parameter) - The key to regenerate - - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update - - key: Optional[str] - The key to regenerate. - - new_master_key: Optional[str] - The new master key to use, if key is the master key. - - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used. - - key_alias: Optional[str] - User-friendly key alias - - user_id: Optional[str] - User ID associated with key - - team_id: Optional[str] - Team ID associated with key - - models: Optional[list] - Model_name's a user is allowed to call - - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) - - spend: Optional[float] - Amount spent by key - - max_budget: Optional[float] - Max budget for key - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. - - max_parallel_requests: Optional[int] - Rate limit for parallel requests - - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} - - tpm_limit: Optional[int] - Tokens per minute limit - - rpm_limit: Optional[int] - Requests per minute limit - - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} - - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} - - allowed_cache_controls: Optional[list] - List of allowed cache control values - - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) - - permissions: Optional[dict] - Key-specific permissions - - guardrails: Optional[List[str]] - List of active guardrails for the key - - blocked: Optional[bool] - Whether the key is blocked - - - Returns: - - GenerateKeyResponse containing the new key and its updated parameters - - Example: - ```bash - curl --location --request POST 'http://localhost:4000/key/sk-1234/regenerate' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data-raw '{ - "max_budget": 100, - "metadata": {"team": "core-infra"}, - "models": ["gpt-4", "gpt-3.5-turbo"] - }' - ``` - - Note: This is an Enterprise feature. It requires a premium license to use. - """ - try: - from litellm.proxy.proxy_server import ( - hash_token, - master_key, - premium_user, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - is_master_key_regeneration = data and data.new_master_key is not None - - if ( - premium_user is not True and not is_master_key_regeneration - ): # allow master key regeneration for non-premium users - raise ValueError( - f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}" - ) - - # Check if key exists, raise exception if key is not in the DB - key = data.key if data and data.key else key - if not key: - raise HTTPException(status_code=400, detail={"error": "No key passed in."}) - ### 1. Create New copy that is duplicate of existing key - ###################################################################### - - # create duplicate of existing key - # set token = new token generated - # insert new token in DB - - # create hash of token - if prisma_client is None: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": "DB not connected. prisma_client is None"}, - ) - - _is_master_key_valid = _is_master_key(api_key=key, _master_key=master_key) - - if master_key is not None and data and _is_master_key_valid: - if data.new_master_key is None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={"error": "New master key is required."}, - ) - await _rotate_master_key( - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - current_master_key=master_key, - new_master_key=data.new_master_key, - ) - return GenerateKeyResponse( - key=data.new_master_key, - token=data.new_master_key, - key_name=data.new_master_key, - expires=None, - ) - - if "sk" not in key: - hashed_api_key = key - else: - hashed_api_key = hash_token(key) - - _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_api_key}, - ) - if _key_in_db is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"Key {key} not found."}, - ) - - # check if user has permission to regenerate key - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_REGENERATE, - prisma_client=prisma_client, - existing_key_row=_key_in_db, - user_api_key_cache=user_api_key_cache, - ) - - # check if user has ownership permission to regenerate key - if not await can_modify_verification_token( - key_info=_key_in_db, - user_api_key_cache=user_api_key_cache, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "You are not authorized to regenerate this key"}, - ) - - verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) - - new_token = get_new_token(data=data) - - new_token_hash = hash_token(new_token) - new_token_key_name = f"sk-...{new_token[-4:]}" - - # Prepare the update data - update_data = { - "token": new_token_hash, - "key_name": new_token_key_name, - } - - non_default_values = {} - if data is not None: - # Update with any provided parameters from GenerateKeyRequest - non_default_values = await prepare_key_update_data( - data=data, existing_key_row=_key_in_db - ) - verbose_proxy_logger.debug("non_default_values: %s", non_default_values) - - update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) - # Update the token in the database - updated_token = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_api_key}, - data=update_data, # type: ignore - ) - - updated_token_dict = {} - if updated_token is not None: - updated_token_dict = dict(updated_token) - - updated_token_dict["key"] = new_token - updated_token_dict["token_id"] = updated_token_dict.pop("token") - - ### 3. remove existing key entry from cache - ###################################################################### - - if hashed_api_key or key: - await _delete_cache_key_object( - hashed_token=hash_token(key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - response = GenerateKeyResponse( - **updated_token_dict, - ) - - asyncio.create_task( - KeyManagementEventHooks.async_key_rotated_hook( - data=data, - existing_key_row=_key_in_db, - response=response, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - ) - - return response - except Exception as e: - verbose_proxy_logger.exception("Error regenerating key: %s", e) - raise handle_exception_on_proxy(e) - - -async def validate_key_list_check( - user_api_key_dict: UserAPIKeyAuth, - user_id: Optional[str], - team_id: Optional[str], - organization_id: Optional[str], - key_alias: Optional[str], - key_hash: Optional[str], - prisma_client: PrismaClient, -) -> Optional[LiteLLM_UserTable]: - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - return None - - if user_api_key_dict.user_id is None: - raise ProxyException( - message="You are not authorized to access this endpoint. No 'user_id' is associated with your API key.", - type=ProxyErrorTypes.bad_request_error, - param="user_id", - code=status.HTTP_403_FORBIDDEN, - ) - complete_user_info_db_obj: Optional[BaseModel] = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, - ) - ) - - if complete_user_info_db_obj is None: - raise ProxyException( - message="You are not authorized to access this endpoint. No 'user_id' is associated with your API key.", - type=ProxyErrorTypes.bad_request_error, - param="user_id", - code=status.HTTP_403_FORBIDDEN, - ) - - complete_user_info = LiteLLM_UserTable(**complete_user_info_db_obj.model_dump()) - - # internal user can only see their own keys - if user_id: - if complete_user_info.user_id != user_id: - raise ProxyException( - message="You are not authorized to check another user's keys", - type=ProxyErrorTypes.bad_request_error, - param="user_id", - code=status.HTTP_403_FORBIDDEN, - ) - - if team_id: - if team_id not in complete_user_info.teams: - raise ProxyException( - message="You are not authorized to check this team's keys", - type=ProxyErrorTypes.bad_request_error, - param="team_id", - code=status.HTTP_403_FORBIDDEN, - ) - - if organization_id: - if ( - complete_user_info.organization_memberships is None - or organization_id - not in [ - membership.organization_id - for membership in complete_user_info.organization_memberships - ] - ): - raise ProxyException( - message="You are not authorized to check this organization's keys", - type=ProxyErrorTypes.bad_request_error, - param="organization_id", - code=status.HTTP_403_FORBIDDEN, - ) - - if key_hash: - try: - key_info = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": key_hash}, - ) - except Exception: - raise ProxyException( - message="Key Hash not found.", - type=ProxyErrorTypes.bad_request_error, - param="key_hash", - code=status.HTTP_403_FORBIDDEN, - ) - can_user_query_key_info = await _can_user_query_key_info( - user_api_key_dict=user_api_key_dict, - key=key_hash, - key_info=key_info, - ) - if not can_user_query_key_info: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="You are not allowed to access this key's info. Your role={}".format( - user_api_key_dict.user_role - ), - ) - return complete_user_info - - -async def get_admin_team_ids( - complete_user_info: Optional[LiteLLM_UserTable], - user_api_key_dict: UserAPIKeyAuth, - prisma_client: PrismaClient, -) -> List[str]: - """ - Get all team IDs where the user is an admin. - """ - if complete_user_info is None: - return [] - # Get all teams that user is an admin of - teams: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} - ) - ) - if teams is None: - return [] - - teams_pydantic_obj = [LiteLLM_TeamTable(**team.model_dump()) for team in teams] - - admin_team_ids = [ - team.team_id - for team in teams_pydantic_obj - if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team) - ] - return admin_team_ids - - -@router.get( - "/key/list", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], -) -@management_endpoint_wrapper -async def list_keys( - request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - page: int = Query(1, description="Page number", ge=1), - size: int = Query(10, description="Page size", ge=1, le=100), - user_id: Optional[str] = Query(None, description="Filter keys by user ID"), - team_id: Optional[str] = Query(None, description="Filter keys by team ID"), - organization_id: Optional[str] = Query( - None, description="Filter keys by organization ID" - ), - key_hash: Optional[str] = Query(None, description="Filter keys by key hash"), - key_alias: Optional[str] = Query(None, description="Filter keys by key alias"), - return_full_object: bool = Query(False, description="Return full key object"), - include_team_keys: bool = Query( - False, description="Include all keys for teams that user is an admin of." - ), - include_created_by_keys: bool = Query( - False, description="Include keys created by the user" - ), - sort_by: Optional[str] = Query( - default=None, - description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", - ), - sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), - expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), - status: Optional[str] = Query(None, description="Filter by status (e.g. 'deleted')"), -) -> KeyListResponseObject: - """ - List all keys for a given user / team / organization. - - Parameters: - expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) - status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. - - Returns: - { - "keys": List[str] or List[UserAPIKeyAuth], - "total_count": int, - "current_page": int, - "total_pages": int, - } - - When expand includes "user", each key object will include a "user" field with the associated user object. - Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter. - """ - try: - from litellm.proxy.proxy_server import prisma_client - - verbose_proxy_logger.debug("Entering list_keys function") - - if prisma_client is None: - verbose_proxy_logger.error("Database not connected") - raise Exception("Database not connected") - - # Validate status parameter - if status is not None and status != "deleted": - raise HTTPException( - status_code=400, - detail={ - "error": "Invalid status value. Currently only 'deleted' is supported." - }, - ) - - complete_user_info = await validate_key_list_check( - user_api_key_dict=user_api_key_dict, - user_id=user_id, - team_id=team_id, - organization_id=organization_id, - key_alias=key_alias, - key_hash=key_hash, - prisma_client=prisma_client, - ) - - if include_team_keys: - admin_team_ids = await get_admin_team_ids( - complete_user_info=complete_user_info, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - else: - admin_team_ids = None - - if user_id is None and user_api_key_dict.user_role not in [ - LitellmUserRoles.PROXY_ADMIN.value, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, - ]: - user_id = user_api_key_dict.user_id - - response = await _list_key_helper( - prisma_client=prisma_client, - page=page, - size=size, - user_id=user_id, - team_id=team_id, - key_alias=key_alias, - key_hash=key_hash, - return_full_object=return_full_object, - organization_id=organization_id, - admin_team_ids=admin_team_ids, - include_created_by_keys=include_created_by_keys, - sort_by=sort_by, - sort_order=sort_order, - expand=expand, - status=status, - ) - - verbose_proxy_logger.debug("Successfully prepared response") - - return response - - except Exception as e: - verbose_proxy_logger.exception(f"Error in list_keys: {e}") - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"error({str(e)})"), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - -@router.get( - "/key/aliases", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], -) -@management_endpoint_wrapper -async def key_aliases() -> Dict[str, List[str]]: - """ - Lists all key aliases - - Returns: - { - "aliases": List[str] - } - """ - try: - from litellm.proxy.proxy_server import prisma_client - - verbose_proxy_logger.debug("Entering key_aliases function") - - if prisma_client is None: - verbose_proxy_logger.error("Database not connected") - raise Exception("Database not connected") - - where: Dict[str, Any] = {} - try: - where.update(_get_condition_to_filter_out_ui_session_tokens()) - except NameError: - # Helper may not exist in some builds; ignore if missing - pass - - rows = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, - order=[{"key_alias": "asc"}], - ) - - seen = set() - aliases: List[str] = [] - for row in rows: - alias = getattr(row, "key_alias", None) - if alias is None and isinstance(row, dict): - alias = row.get("key_alias") - - if not alias: - continue - - alias_str = str(alias).strip() - if alias_str and alias_str not in seen: - seen.add(alias_str) - aliases.append(alias_str) - - verbose_proxy_logger.debug(f"Returning {len(aliases)} key aliases") - - return {"aliases": aliases} - - except Exception as e: - verbose_proxy_logger.exception(f"Error in key_aliases: {e}") - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"error({str(e)})"), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - -def _validate_sort_params( - sort_by: Optional[str], sort_order: str -) -> Optional[Dict[str, str]]: - order_by: Dict[str, str] = {} - - if sort_by is None: - return None - # Validate sort_by is a valid column - valid_columns = [ - "spend", - "max_budget", - "created_at", - "updated_at", - "token", - "key_alias", - ] - if sort_by not in valid_columns: - raise HTTPException( - status_code=400, - detail={ - "error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}" - }, - ) - - # Validate sort_order - if sort_order.lower() not in ["asc", "desc"]: - raise HTTPException( - status_code=400, - detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"}, - ) - - order_by[sort_by] = sort_order.lower() - - return order_by - - -def _build_key_filter_conditions( - user_id: Optional[str], - team_id: Optional[str], - organization_id: Optional[str], - key_alias: Optional[str], - key_hash: Optional[str], - exclude_team_id: Optional[str], - admin_team_ids: Optional[List[str]], - include_created_by_keys: bool, -) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: - """Build filter conditions for key listing.""" - # Prepare filter conditions - where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {} - where.update(_get_condition_to_filter_out_ui_session_tokens()) - - # Build the OR conditions for user's keys and admin team keys - or_conditions: List[Dict[str, Any]] = [] - - # Base conditions for user's own keys - user_condition: Dict[str, Any] = {} - if user_id and isinstance(user_id, str): - user_condition["user_id"] = user_id - if team_id and isinstance(team_id, str): - user_condition["team_id"] = team_id - if key_alias and isinstance(key_alias, str): - user_condition["key_alias"] = key_alias - if exclude_team_id and isinstance(exclude_team_id, str): - user_condition["team_id"] = {"not": exclude_team_id} - if organization_id and isinstance(organization_id, str): - user_condition["organization_id"] = organization_id - if key_hash and isinstance(key_hash, str): - user_condition["token"] = key_hash - - if user_condition: - or_conditions.append(user_condition) - - # Add condition for created by keys if provided - if include_created_by_keys and user_id: - or_conditions.append({"created_by": user_id}) - - # Add condition for admin team keys if provided - if admin_team_ids: - or_conditions.append({"team_id": {"in": admin_team_ids}}) - - # Combine conditions with OR if we have multiple conditions - if len(or_conditions) > 1: - where = {"AND": [where, {"OR": or_conditions}]} - elif len(or_conditions) == 1: - where.update(or_conditions[0]) - - verbose_proxy_logger.debug(f"Filter conditions: {where}") - return where - - -async def _list_key_helper( - prisma_client: PrismaClient, - page: int, - size: int, - user_id: Optional[str], - team_id: Optional[str], - organization_id: Optional[str], - key_alias: Optional[str], - key_hash: Optional[str], - exclude_team_id: Optional[str] = None, - return_full_object: bool = False, - admin_team_ids: Optional[ - List[str] - ] = None, # New parameter for teams where user is admin - include_created_by_keys: bool = False, - sort_by: Optional[str] = None, - sort_order: str = "desc", - expand: Optional[List[str]] = None, - status: Optional[str] = None, -) -> KeyListResponseObject: - """ - Helper function to list keys - Args: - page: int - size: int - user_id: Optional[str] - team_id: Optional[str] - key_alias: Optional[str] - exclude_team_id: Optional[str] # exclude a specific team_id - return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token - admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin - - Returns: - KeyListResponseObject - { - "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types - "total_count": int, - "current_page": int, - "total_pages": int, - } - """ - where = _build_key_filter_conditions( - user_id=user_id, - team_id=team_id, - organization_id=organization_id, - key_alias=key_alias, - key_hash=key_hash, - exclude_team_id=exclude_team_id, - admin_team_ids=admin_team_ids, - include_created_by_keys=include_created_by_keys, - ) - - # Calculate skip for pagination - skip = (page - 1) * size - - verbose_proxy_logger.debug(f"Pagination: skip={skip}, take={size}") - - order_by: Optional[Dict[str, str]] = ( - _validate_sort_params(sort_by, sort_order) - if sort_by is not None and isinstance(sort_by, str) - else None - ) - - # Determine which table to query based on status - use_deleted_table = status == "deleted" - - # Fetch keys with pagination - if use_deleted_table: - keys = await prisma_client.db.litellm_deletedverificationtoken.find_many( - where=where, # type: ignore - skip=skip, # type: ignore - take=size, # type: ignore - order=( - order_by - if order_by - else [ - {"created_at": "desc"}, - {"token": "desc"}, # fallback sort - ] - ), - ) - else: - keys = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, # type: ignore - skip=skip, # type: ignore - take=size, # type: ignore - order=( - order_by - if order_by - else [ - {"created_at": "desc"}, - {"token": "desc"}, # fallback sort - ] - ), - include={"object_permission": True}, - ) - - verbose_proxy_logger.debug(f"Fetched {len(keys)} keys") - - # Get total count of keys - if use_deleted_table: - total_count = await prisma_client.db.litellm_deletedverificationtoken.count( - where=where # type: ignore - ) - else: - total_count = await prisma_client.db.litellm_verificationtoken.count( - where=where # type: ignore - ) - - verbose_proxy_logger.debug(f"Total count of keys: {total_count}") - - # Calculate total pages - total_pages = -(-total_count // size) # Ceiling division - - # Fetch user information if expand includes "user" - user_map = {} - if expand and "user" in expand: - user_ids = [key.user_id for key in keys if key.user_id] - if user_ids: - users = await prisma_client.db.litellm_usertable.find_many( - where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates - ) - user_map = {user.user_id: user for user in users} - - # Prepare response - key_list: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] = [] - for key in keys: - # Convert Prisma model to dict (supports both Pydantic v1 and v2) - try: - key_dict = key.model_dump() - except Exception: - # Fallback for Pydantic v1 compatibility - key_dict = key.dict() - # Attach object_permission if object_permission_id is set (only for non-deleted keys) - if not use_deleted_table: - key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) - - # Include user information if expand includes "user" - if expand and "user" in expand and key.user_id and key.user_id in user_map: - try: - key_dict["user"] = user_map[key.user_id].model_dump() - except Exception: - key_dict["user"] = user_map[key.user_id].dict() - - if return_full_object is True or (expand and "user" in expand): - if use_deleted_table: - # Use deleted key type to preserve deleted_at, deleted_by, etc. - key_list.append(LiteLLM_DeletedVerificationToken(**key_dict)) - else: - key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object - else: - _token = key_dict.get("token") - key_list.append(cast(str, _token)) # Return only the token - - return KeyListResponseObject( - keys=key_list, - total_count=total_count, - current_page=page, - total_pages=total_pages, - ) - - -def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: - """ - Condition to filter out UI session tokens - """ - return { - "OR": [ - {"team_id": None}, # Include records where team_id is null - { - "team_id": {"not": UI_SESSION_TOKEN_TEAM_ID} - }, # Include records where team_id != UI_SESSION_TOKEN_TEAM_ID - ] - } - - -@router.post( - "/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) -@management_endpoint_wrapper -async def block_key( - data: BlockKeyRequest, - http_request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -) -> Optional[LiteLLM_VerificationToken]: - """ - Block an Virtual key from making any requests. - - Parameters: - - key: str - The key to block. Can be either the unhashed key (sk-...) or the hashed key value - - Example: - ```bash - curl --location 'http://0.0.0.0:4000/key/block' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" - }' - ``` - - Note: This is an admin-only endpoint. Only proxy admins can block keys. - """ - from litellm.proxy.proxy_server import ( - create_audit_log_for_update, - hash_token, - litellm_proxy_admin_name, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if prisma_client is None: - raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) - - if not is_valid_api_key(data.key): - raise ProxyException( - message="Invalid key format.", - type=ProxyErrorTypes.bad_request_error, - param="key", - code=status.HTTP_400_BAD_REQUEST, - ) - if data.key.startswith("sk-"): - hashed_token = hash_token(token=data.key) - else: - hashed_token = data.key - - if litellm.store_audit_logs is True: - # make an audit log for key update - record = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) - if record is None: - raise ProxyException( - message=f"Key {data.key} not found", - type=ProxyErrorTypes.bad_request_error, - param="key", - code=status.HTTP_404_NOT_FOUND, - ) - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by - or user_api_key_dict.user_id - or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=hashed_token, - action="blocked", - updated_values="{}", - before_value=record.model_dump_json(), - ) - ) - ) - - record = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_token}, data={"blocked": True} # type: ignore - ) - - ## UPDATE KEY CACHE - - ### get cached object ### - key_object = await get_key_object( - hashed_token=hashed_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=proxy_logging_obj, - ) - - ### update cached object ### - key_object.blocked = True - - ### store cached object ### - await _cache_key_object( - hashed_token=hashed_token, - user_api_key_obj=key_object, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - return record - - -@router.post( - "/key/unblock", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) -@management_endpoint_wrapper -async def unblock_key( - data: BlockKeyRequest, - http_request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( - None, - description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", - ), -): - """ - Unblock a Virtual key to allow it to make requests again. - - Parameters: - - key: str - The key to unblock. Can be either the unhashed key (sk-...) or the hashed key value - - Example: - ```bash - curl --location 'http://0.0.0.0:4000/key/unblock' \ - --header 'Authorization: Bearer sk-1234' \ - --header 'Content-Type: application/json' \ - --data '{ - "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" - }' - ``` - - Note: This is an admin-only endpoint. Only proxy admins can unblock keys. - """ - from litellm.proxy.proxy_server import ( - create_audit_log_for_update, - hash_token, - litellm_proxy_admin_name, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if prisma_client is None: - raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) - - if not is_valid_api_key(data.key): - raise ProxyException( - message="Invalid key format.", - type=ProxyErrorTypes.bad_request_error, - param="key", - code=status.HTTP_400_BAD_REQUEST, - ) - if data.key.startswith("sk-"): - hashed_token = hash_token(token=data.key) - else: - hashed_token = data.key - - if litellm.store_audit_logs is True: - # make an audit log for key update - record = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) - if record is None: - raise ProxyException( - message=f"Key {data.key} not found", - type=ProxyErrorTypes.bad_request_error, - param="key", - code=status.HTTP_404_NOT_FOUND, - ) - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by - or user_api_key_dict.user_id - or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=hashed_token, - action="blocked", - updated_values="{}", - before_value=record.model_dump_json(), - ) - ) - ) - - record = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_token}, data={"blocked": False} # type: ignore - ) - - ## UPDATE KEY CACHE - - ### get cached object ### - key_object = await get_key_object( - hashed_token=hashed_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=proxy_logging_obj, - ) - - ### update cached object ### - key_object.blocked = False - - ### store cached object ### - await _cache_key_object( - hashed_token=hashed_token, - user_api_key_obj=key_object, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - return record - - -@router.post( - "/key/health", - tags=["key management"], - dependencies=[Depends(user_api_key_auth)], - response_model=KeyHealthResponse, -) -@management_endpoint_wrapper -async def key_health( - request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Check the health of the key - - Checks: - - If key based logging is configured correctly - sends a test log - - Usage - - Pass the key in the request header - - ```bash - curl -X POST "http://localhost:4000/key/health" \ - -H "Authorization: Bearer sk-1234" \ - -H "Content-Type: application/json" - ``` - - Response when logging callbacks are setup correctly: - - ```json - { - "key": "healthy", - "logging_callbacks": { - "callbacks": [ - "gcs_bucket" - ], - "status": "healthy", - "details": "No logger exceptions triggered, system is healthy. Manually check if logs were sent to ['gcs_bucket']" - } - } - ``` - - - Response when logging callbacks are not setup correctly: - ```json - { - "key": "unhealthy", - "logging_callbacks": { - "callbacks": [ - "gcs_bucket" - ], - "status": "unhealthy", - "details": "Logger exceptions triggered, system is unhealthy: Failed to load vertex credentials. Check to see if credentials containing partial/invalid information." - } - } - ``` - """ - try: - # Get the key's metadata - key_metadata = user_api_key_dict.metadata - - health_status: KeyHealthResponse = KeyHealthResponse( - key="healthy", - logging_callbacks=None, - ) - - # Check if logging is configured in metadata - if key_metadata and "logging" in key_metadata: - logging_statuses = await test_key_logging( - user_api_key_dict=user_api_key_dict, - request=request, - key_logging=key_metadata["logging"], - ) - health_status["logging_callbacks"] = logging_statuses - - # Check if any logging callback is unhealthy - if logging_statuses.get("status") == "unhealthy": - health_status["key"] = "unhealthy" - - return KeyHealthResponse(**health_status) - - except Exception as e: - raise ProxyException( - message=f"Key health check failed: {str(e)}", - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - -async def _can_user_query_key_info( - user_api_key_dict: UserAPIKeyAuth, - key: Optional[str], - key_info: LiteLLM_VerificationToken, -) -> bool: - """ - Helper to check if the user has access to the key's info - """ - if ( - user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ): - return True - elif user_api_key_dict.api_key == key: - return True - # user can query their own key info - elif key_info.user_id == user_api_key_dict.user_id: - return True - elif await TeamMemberPermissionChecks.user_belongs_to_keys_team( - user_api_key_dict=user_api_key_dict, - existing_key_row=key_info, - ): - return True - return False - - -async def test_key_logging( - user_api_key_dict: UserAPIKeyAuth, - request: Request, - key_logging: List[Dict[str, Any]], -) -> LoggingCallbackStatus: - """ - Test the key-based logging - - - Test that key logging is correctly formatted and all args are passed correctly - - Make a mock completion call -> user can check if it's correctly logged - - Check if any logger.exceptions were triggered -> if they were then returns it to the user client side - """ - import logging - from io import StringIO - - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - from litellm.proxy.proxy_server import general_settings, proxy_config - - logging_callbacks: List[str] = [] - for callback in key_logging: - if callback.get("callback_name") is not None: - logging_callbacks.append(callback["callback_name"]) - else: - raise ValueError("callback_name is required in key_logging") - - log_capture_string = StringIO() - ch = logging.StreamHandler(log_capture_string) - ch.setLevel(logging.ERROR) - logger = logging.getLogger() - logger.addHandler(ch) - - try: - data = { - "model": "openai/litellm-key-health-test", - "messages": [ - { - "role": "user", - "content": "Hello, this is a test from litellm /key/health. No LLM API call was made for this", - } - ], - "mock_response": "test response", - } - data = await add_litellm_data_to_request( - data=data, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - general_settings=general_settings, - request=request, - ) - await litellm.acompletion( - **data - ) # make mock completion call to trigger key based callbacks - except Exception as e: - return LoggingCallbackStatus( - callbacks=logging_callbacks, - status="unhealthy", - details=f"Logging test failed: {str(e)}", - ) - - await asyncio.sleep( - 2 - ) # wait for callbacks to run, callbacks use batching so wait for the flush event - - # Check if any logger exceptions were triggered - log_contents = log_capture_string.getvalue() - logger.removeHandler(ch) - if log_contents: - return LoggingCallbackStatus( - callbacks=logging_callbacks, - status="unhealthy", - details=f"Logger exceptions triggered, system is unhealthy: {log_contents}", - ) - else: - return LoggingCallbackStatus( - callbacks=logging_callbacks, - status="healthy", - details=f"No logger exceptions triggered, system is healthy. Manually check if logs were sent to {logging_callbacks} ", - ) - - -async def _enforce_unique_key_alias( - key_alias: Optional[str], - prisma_client: Any, - existing_key_token: Optional[str] = None, -) -> None: - """ - Helper to enforce unique key aliases across all keys. - - Args: - key_alias (Optional[str]): The key alias to check - prisma_client (Any): Prisma client instance - existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check - (The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating) - - Raises: - ProxyException: If key alias already exists on a different key - """ - if key_alias is not None and prisma_client is not None: - where_clause: dict[str, Any] = {"key_alias": key_alias} - if existing_key_token: - # Exclude the current key from the uniqueness check - where_clause["NOT"] = {"token": existing_key_token} - - existing_key = await prisma_client.db.litellm_verificationtoken.find_first( - where=where_clause - ) - if existing_key is not None: - raise ProxyException( - message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", - type=ProxyErrorTypes.bad_request_error, - param="key_alias", - code=status.HTTP_400_BAD_REQUEST, - ) - - -def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: - """ - Validate the model_max_budget is GenericBudgetConfigType + enforce user has an enterprise license - - Raises: - Exception: If model_max_budget is not a valid GenericBudgetConfigType - """ - try: - if model_max_budget is None: - return - if len(model_max_budget) == 0: - return - if model_max_budget is not None: - from litellm.proxy.proxy_server import CommonProxyErrors, premium_user - - if premium_user is not True: - raise ValueError( - f"You must have an enterprise license to set model_max_budget. {CommonProxyErrors.not_premium_user.value}" - ) - for _model, _budget_info in model_max_budget.items(): - assert isinstance(_model, str) - - # /CRUD endpoints can pass budget_limit as a string, so we need to convert it to a float - if "budget_limit" in _budget_info: - _budget_info["budget_limit"] = float(_budget_info["budget_limit"]) - BudgetConfig(**_budget_info) - except Exception as e: - raise ValueError( - f"Invalid model_max_budget: {str(e)}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users" - ) +""" +KEY MANAGEMENT + +All /key management endpoints + +/key/generate +/key/info +/key/update +/key/delete +""" + +import asyncio +import copy +import json +import secrets +import traceback +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List, Literal, Optional, Tuple, cast + +import fastapi +import yaml +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.caching import DualCache +from litellm.constants import ( + LENGTH_OF_LITELLM_GENERATED_KEY, + LITELLM_PROXY_ADMIN_NAME, + UI_SESSION_TOKEN_TEAM_ID, +) +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._experimental.mcp_server.db import ( + rotate_mcp_server_credentials_master_key, +) +from litellm.proxy._types import * +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + BulkUpdateKeyResponse, + FailedKeyUpdate, + SuccessfulKeyUpdate, +) +from litellm.proxy.auth.auth_checks import ( + _cache_key_object, + _delete_cache_key_object, + can_team_access_model, + get_key_object, + get_org_object, + get_team_object, +) +from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks +from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, + _set_object_metadata_field, +) +from litellm.proxy.management_endpoints.model_management_endpoints import ( + _add_model_to_db, +) +from litellm.proxy.management_helpers.object_permission_utils import ( + _set_object_permission, + attach_object_permission_to_dict, + handle_update_object_permission_common, +) +from litellm.proxy.management_helpers.team_member_permission_checks import ( + TeamMemberPermissionChecks, +) +from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key +from litellm.proxy.utils import ( + PrismaClient, + _hash_token_if_needed, + handle_exception_on_proxy, + is_valid_api_key, + jsonify_object, +) +from litellm.router import Router +from litellm.secret_managers.main import get_secret +from litellm.types.router import Deployment +from litellm.types.utils import ( + BudgetConfig, + PersonalUIKeyGenerationConfig, + TeamUIKeyGenerationConfig, +) + + +def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]): + return data.team_id is not None + + +def _get_user_in_team( + team_table: LiteLLM_TeamTableCachedObj, user_id: Optional[str] +) -> Optional[Member]: + if user_id is None: + return None + for member in team_table.members_with_roles: + if member.user_id is not None and member.user_id == user_id: + return member + + return None + + +def _calculate_key_rotation_time(rotation_interval: str) -> datetime: + """ + Helper function to calculate the next rotation time for a key based on the rotation interval. + + Args: + rotation_interval: String representing the rotation interval (e.g., '30d', '90d', '1h') + + Returns: + datetime: The calculated next rotation time in UTC + """ + now = datetime.now(timezone.utc) + interval_seconds = duration_in_seconds(rotation_interval) + return now + timedelta(seconds=interval_seconds) + + +def _set_key_rotation_fields( + data: dict, auto_rotate: bool, rotation_interval: Optional[str] +) -> None: + """ + Helper function to set rotation fields in key data if auto_rotate is enabled. + + Args: + data: Dictionary to update with rotation fields + auto_rotate: Whether auto rotation is enabled + rotation_interval: The rotation interval string (required if auto_rotate is True) + """ + if auto_rotate and rotation_interval: + data.update( + { + "auto_rotate": auto_rotate, + "rotation_interval": rotation_interval, + "key_rotation_at": _calculate_key_rotation_time(rotation_interval), + } + ) + + +def _is_allowed_to_make_key_request( + user_api_key_dict: UserAPIKeyAuth, + user_id: Optional[str], + team_id: Optional[str], +) -> bool: + """ + Assert user only creates/updates keys for themselves + + Relevant issue: https://github.com/BerriAI/litellm/issues/7336 + """ + ## BASE CASE - PROXY ADMIN + if ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ): + return True + + if user_id is not None: + assert ( + user_id == user_api_key_dict.user_id + ), "User can only create keys for themselves. Got user_id={}, Your ID={}".format( + user_id, user_api_key_dict.user_id + ) + + if team_id is not None: + if ( + user_api_key_dict.team_id is not None + and user_api_key_dict.team_id == UI_TEAM_ID + ): + return True # handle https://github.com/BerriAI/litellm/issues/7482 + + return True + + +def _team_key_operation_team_member_check( + assigned_user_id: Optional[str], + team_table: LiteLLM_TeamTableCachedObj, + user_api_key_dict: UserAPIKeyAuth, + team_key_generation: TeamUIKeyGenerationConfig, + route: KeyManagementRoutes, +): + if assigned_user_id is not None: + key_assigned_user_in_team = _get_user_in_team( + team_table=team_table, user_id=assigned_user_id + ) + + if key_assigned_user_in_team is None: + raise HTTPException( + status_code=400, + detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}", + ) + + team_member_object = _get_user_in_team( + team_table=team_table, user_id=user_api_key_dict.user_id + ) + + is_admin = ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + + if is_admin: + return True + elif team_member_object is None: + raise HTTPException( + status_code=400, + detail=f"User={user_api_key_dict.user_id} not assigned to team={team_table.team_id}", + ) + elif ( + "allowed_team_member_roles" in team_key_generation + and team_member_object.role + not in team_key_generation["allowed_team_member_roles"] + ): + raise HTTPException( + status_code=400, + detail=f"Team member role {team_member_object.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}", + ) + + TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( + team_member_object=team_member_object, + team_table=team_table, + route=route, + ) + return True + + +def _key_generation_required_param_check( + data: GenerateKeyRequest, required_params: Optional[List[str]] +): + if required_params is None: + return True + + data_dict = data.model_dump(exclude_unset=True) + for param in required_params: + if param not in data_dict: + raise HTTPException( + status_code=400, + detail=f"Required param {param} not in data", + ) + return True + + +def _team_key_generation_check( + team_table: LiteLLM_TeamTableCachedObj, + user_api_key_dict: UserAPIKeyAuth, + data: GenerateKeyRequest, + route: KeyManagementRoutes, +): + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return True + if ( + litellm.key_generation_settings is not None + and "team_key_generation" in litellm.key_generation_settings + ): + _team_key_generation = litellm.key_generation_settings["team_key_generation"] + else: + _team_key_generation = TeamUIKeyGenerationConfig( + allowed_team_member_roles=["admin", "user"], + ) + + _team_key_operation_team_member_check( + assigned_user_id=data.user_id, + team_table=team_table, + user_api_key_dict=user_api_key_dict, + team_key_generation=_team_key_generation, + route=route, + ) + _key_generation_required_param_check( + data, + _team_key_generation.get("required_params"), + ) + + return True + + +def _personal_key_membership_check( + user_api_key_dict: UserAPIKeyAuth, + personal_key_generation: Optional[PersonalUIKeyGenerationConfig], +): + if ( + personal_key_generation is None + or "allowed_user_roles" not in personal_key_generation + ): + return True + + if user_api_key_dict.user_role not in personal_key_generation["allowed_user_roles"]: + raise HTTPException( + status_code=400, + detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}", # type: ignore + ) + + return True + + +def _personal_key_generation_check( + user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest +): + if ( + litellm.key_generation_settings is None + or litellm.key_generation_settings.get("personal_key_generation") is None + ): + return True + + _personal_key_generation = litellm.key_generation_settings["personal_key_generation"] # type: ignore + + _personal_key_membership_check( + user_api_key_dict, + personal_key_generation=_personal_key_generation, + ) + + _key_generation_required_param_check( + data, + _personal_key_generation.get("required_params"), + ) + + return True + + +def key_generation_check( + team_table: Optional[LiteLLM_TeamTableCachedObj], + user_api_key_dict: UserAPIKeyAuth, + data: GenerateKeyRequest, + route: KeyManagementRoutes, +) -> bool: + """ + Check if admin has restricted key creation to certain roles for teams or individuals + """ + + ## check if key is for team or individual + is_team_key = _is_team_key(data=data) + if is_team_key: + if team_table is None and litellm.key_generation_settings is not None: + raise HTTPException( + status_code=400, + detail=f"Unable to find team object in database. Team ID: {data.team_id}", + ) + elif team_table is None: + return True # assume user is assigning team_id without using the team table + return _team_key_generation_check( + team_table=team_table, + user_api_key_dict=user_api_key_dict, + data=data, + route=route, + ) + else: + return _personal_key_generation_check( + user_api_key_dict=user_api_key_dict, data=data + ) + + +def common_key_access_checks( + user_api_key_dict: UserAPIKeyAuth, + data: Union[GenerateKeyRequest, UpdateKeyRequest], + llm_router: Optional[Router], + premium_user: bool, + user_id: Optional[str] = None, +) -> Literal[True]: + """ + Check if user is allowed to make a key request, for this key + """ + try: + _is_allowed_to_make_key_request( + user_api_key_dict=user_api_key_dict, + user_id=user_id or data.user_id, + team_id=data.team_id, + ) + except AssertionError as e: + raise HTTPException( + status_code=403, + detail=str(e), + ) + except Exception as e: + raise HTTPException( + status_code=500, + detail=str(e), + ) + + _check_model_access_group( + models=data.models, + llm_router=llm_router, + premium_user=premium_user, + ) + return True + + +router = APIRouter() + + +def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: + """ + Handle the key type. + """ + key_type = data.key_type + data_json.pop("key_type", None) + if key_type == LiteLLMKeyType.LLM_API: + data_json["allowed_routes"] = ["llm_api_routes"] + elif key_type == LiteLLMKeyType.MANAGEMENT: + data_json["allowed_routes"] = ["management_routes"] + elif key_type == LiteLLMKeyType.READ_ONLY: + data_json["allowed_routes"] = ["info_routes"] + return data_json + + +async def validate_team_id_used_in_service_account_request( + team_id: Optional[str], + prisma_client: Optional[PrismaClient], +): + """ + Validate team_id is used in the request body for generating a service account key + """ + if team_id is None: + raise HTTPException( + status_code=400, + detail="team_id is required for service account keys. Please specify `team_id` in the request body.", + ) + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail="prisma_client is required for service account keys. Please specify `prisma_client` in the request body.", + ) + + # check if team_id exists in the database + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + if team is None: + raise HTTPException( + status_code=400, + detail="team_id does not exist in the database. Please specify a valid `team_id` in the request body.", + ) + return True + + +async def _common_key_generation_helper( # noqa: PLR0915 + data: GenerateKeyRequest, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str], + team_table: Optional[LiteLLM_TeamTableCachedObj], +) -> GenerateKeyResponse: + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + llm_router, + premium_user, + prisma_client, + ) + + common_key_access_checks( + user_api_key_dict=user_api_key_dict, + data=data, + llm_router=llm_router, + premium_user=premium_user, + ) + + if ( + data.metadata is not None + and data.metadata.get("service_account_id") is not None + and data.team_id is None + ): + await validate_team_id_used_in_service_account_request( + team_id=data.team_id, + prisma_client=prisma_client, + ) + + # check if user set default key/generate params on config.yaml + if litellm.default_key_generate_params is not None: + for elem in data: + key, value = elem + if value is None and key in [ + "max_budget", + "user_id", + "team_id", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "budget_duration", + ]: + setattr(data, key, litellm.default_key_generate_params.get(key, None)) + elif key == "models" and value == []: + setattr(data, key, litellm.default_key_generate_params.get(key, [])) + elif key == "metadata" and value == {}: + setattr(data, key, litellm.default_key_generate_params.get(key, {})) + + # check if user set default key/generate params on config.yaml + if litellm.upperbound_key_generate_params is not None: + for elem in data: + key, value = elem + upperbound_value = getattr( + litellm.upperbound_key_generate_params, key, None + ) + if upperbound_value is not None: + if value is None: + # Use the upperbound value if user didn't provide a value + setattr(data, key, upperbound_value) + else: + # Compare with upperbound for numeric fields + if key in [ + "max_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + ]: + if value > upperbound_value: + raise HTTPException( + status_code=400, + detail={ + "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" + }, + ) + # Compare durations + elif key in ["budget_duration", "duration"]: + upperbound_duration = duration_in_seconds( + duration=upperbound_value + ) + # Handle special case where duration is "-1" (never expires) + if value == "-1": + user_duration = float("inf") # Infinite duration + else: + user_duration = duration_in_seconds(duration=value) + if user_duration > upperbound_duration: + raise HTTPException( + status_code=400, + detail={ + "error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}" + }, + ) + + # APPLY ENTERPRISE KEY MANAGEMENT PARAMS + try: + from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import ( + apply_enterprise_key_management_params, + ) + + data = apply_enterprise_key_management_params(data, team_table) + except Exception as e: + verbose_proxy_logger.debug( + "litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - {}".format( + str(e) + ) + ) + + # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable + _budget_id = data.budget_id + if prisma_client is not None and data.soft_budget is not None: + # create the Budget Row for the LiteLLM Verification Token + budget_row = LiteLLM_BudgetTable( + soft_budget=data.soft_budget, + model_max_budget=data.model_max_budget or {}, + ) + new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + + _budget = await prisma_client.db.litellm_budgettable.create( + data={ + **new_budget, # type: ignore + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + ) + _budget_id = getattr(_budget, "budget_id", None) + + # ADD METADATA FIELDS + # Set Management Endpoint Metadata Fields + for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=data, + field_name=field, + value=getattr(data, field), + ) + delattr(data, field) + + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=data, + field_name=field, + value=getattr(data, field), + ) + delattr(data, field) + + data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore + + data_json = handle_key_type(data, data_json) + + # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users + if "max_budget" in data_json: + data_json["key_max_budget"] = data_json.pop("max_budget", None) + if _budget_id is not None: + data_json["budget_id"] = _budget_id + + if "budget_duration" in data_json: + data_json["key_budget_duration"] = data_json.pop("budget_duration", None) + + if user_api_key_dict.user_id is not None: + data_json["created_by"] = user_api_key_dict.user_id + data_json["updated_by"] = user_api_key_dict.user_id + + # Set tags on the new key + if "tags" in data_json: + from litellm.proxy.proxy_server import premium_user + + if premium_user is not True and data_json["tags"] is not None: + raise ValueError( + f"Only premium users can add tags to keys. {CommonProxyErrors.not_premium_user.value}" + ) + + _metadata = data_json.get("metadata") + if not _metadata: + data_json["metadata"] = {"tags": data_json["tags"]} + else: + data_json["metadata"]["tags"] = data_json["tags"] + + data_json.pop("tags") + + data_json = await _set_object_permission( + data_json=data_json, + prisma_client=prisma_client, + ) + + await _enforce_unique_key_alias( + key_alias=data_json.get("key_alias", None), + prisma_client=prisma_client, + ) + + # Validate user-provided key format + if data.key is not None and not data.key.startswith("sk-"): + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {data.key}" + }, + ) + + # check org key limits - done here to handle inheriting org id from team + if data.organization_id is not None: + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client: + org_table = await get_org_object( + org_id=data.organization_id, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + if org_table is None: + raise HTTPException( + status_code=400, + detail=f"Organization not found for organization_id={data.organization_id}", + ) + await _check_org_key_limits( + org_table=org_table, + data=data, + prisma_client=prisma_client, + ) + + response = await generate_key_helper_fn( + request_type="key", **data_json, table_name="key" + ) + + response[ + "soft_budget" + ] = data.soft_budget # include the user-input soft budget in the response + + response = GenerateKeyResponse(**response) + + response.token = ( + response.token_id + ) # remap token to use the hash, and leave the key in the `key` field [TODO]: clean up generate_key_helper_fn to do this + + asyncio.create_task( + KeyManagementEventHooks.async_key_generated_hook( + data=data, + response=response, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + ) + + return response + + +def _check_key_model_specific_limits( + keys: List[LiteLLM_VerificationToken], + data: Union[GenerateKeyRequest, UpdateKeyRequest], + entity_rpm_limit: Optional[int], + entity_tpm_limit: Optional[int], + entity_model_rpm_limit_dict: Dict[str, int], + entity_model_tpm_limit_dict: Dict[str, int], + entity_type: str, # "team" or "organization" +) -> None: + """ + Generic function to check if a key is allocating model specific limits. + Raises an error if we're overallocating. + """ + model_rpm_limit = getattr(data, "model_rpm_limit", None) or ( + data.metadata.get("model_rpm_limit", None) if data.metadata else None + ) + model_tpm_limit = getattr(data, "model_tpm_limit", None) or ( + data.metadata.get("model_tpm_limit", None) if data.metadata else None + ) + if model_rpm_limit is None and model_tpm_limit is None: + return + + # get total model specific tpm/rpm limit + model_specific_rpm_limit: Dict[str, int] = {} + model_specific_tpm_limit: Dict[str, int] = {} + + for key in keys: + if key.metadata.get("model_rpm_limit", None) is not None: + for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items(): + model_specific_rpm_limit[model] = ( + model_specific_rpm_limit.get(model, 0) + rpm_limit + ) + if key.metadata.get("model_tpm_limit", None) is not None: + for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items(): + model_specific_tpm_limit[model] = ( + model_specific_tpm_limit.get(model, 0) + tpm_limit + ) + + if model_rpm_limit is not None: + for model, rpm_limit in model_rpm_limit.items(): + if ( + entity_rpm_limit is not None + and model_specific_rpm_limit.get(model, 0) + rpm_limit + > entity_rpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", + ) + elif entity_model_rpm_limit_dict: + entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model) + if ( + entity_model_specific_rpm_limit + and model_specific_rpm_limit.get(model, 0) + rpm_limit + > entity_model_specific_rpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit}", + ) + + if model_tpm_limit is not None: + for model, tpm_limit in model_tpm_limit.items(): + if ( + entity_tpm_limit is not None + and model_specific_tpm_limit.get(model, 0) + tpm_limit + > entity_tpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", + ) + elif entity_model_tpm_limit_dict: + entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model) + if ( + entity_model_specific_tpm_limit + and model_specific_tpm_limit.get(model, 0) + tpm_limit + > entity_model_specific_tpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit}", + ) + + +def _check_key_rpm_tpm_limits( + keys: List[LiteLLM_VerificationToken], + data: Union[GenerateKeyRequest, UpdateKeyRequest], + entity_rpm_limit: Optional[int], + entity_tpm_limit: Optional[int], + entity_type: str, # "team" or "organization" +) -> None: + """ + Generic function to check if a key is allocating rpm/tpm limits. + Raises an error if we're overallocating. + """ + if keys is not None and len(keys) > 0: + allocated_tpm = sum(key.tpm_limit for key in keys if key.tpm_limit is not None) + allocated_rpm = sum(key.rpm_limit for key in keys if key.rpm_limit is not None) + else: + allocated_tpm = 0 + allocated_rpm = 0 + + if ( + data.tpm_limit is not None + and entity_tpm_limit is not None + and data.tpm_limit + allocated_tpm > entity_tpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated TPM limit={allocated_tpm} + Key TPM limit={data.tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", + ) + if ( + data.rpm_limit is not None + and entity_rpm_limit is not None + and data.rpm_limit + allocated_rpm > entity_rpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated RPM limit={allocated_rpm} + Key RPM limit={data.rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", + ) + + +def check_team_key_model_specific_limits( + keys: List[LiteLLM_VerificationToken], + team_table: LiteLLM_TeamTableCachedObj, + data: Union[GenerateKeyRequest, UpdateKeyRequest], +) -> None: + """ + Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating. + """ + entity_model_rpm_limit_dict = {} + entity_model_tpm_limit_dict = {} + if team_table.metadata: + entity_model_rpm_limit_dict = team_table.metadata.get("model_rpm_limit", {}) + entity_model_tpm_limit_dict = team_table.metadata.get("model_tpm_limit", {}) + + _check_key_model_specific_limits( + keys=keys, + data=data, + entity_rpm_limit=team_table.rpm_limit, + entity_tpm_limit=team_table.tpm_limit, + entity_model_rpm_limit_dict=entity_model_rpm_limit_dict, + entity_model_tpm_limit_dict=entity_model_tpm_limit_dict, + entity_type="team", + ) + + +def check_team_key_rpm_tpm_limits( + keys: List[LiteLLM_VerificationToken], + team_table: LiteLLM_TeamTableCachedObj, + data: Union[GenerateKeyRequest, UpdateKeyRequest], +) -> None: + """ + Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. + """ + _check_key_rpm_tpm_limits( + keys=keys, + data=data, + entity_rpm_limit=team_table.rpm_limit, + entity_tpm_limit=team_table.tpm_limit, + entity_type="team", + ) + + +async def _check_team_key_limits( + team_table: LiteLLM_TeamTableCachedObj, + data: Union[GenerateKeyRequest, UpdateKeyRequest], + prisma_client: PrismaClient, +) -> None: + """ + Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + + Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" + """ + if ( + data.tpm_limit_type != "guaranteed_throughput" + and data.rpm_limit_type != "guaranteed_throughput" + ): + return + # get all team keys + # calculate allocated tpm/rpm limit + # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit + + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": team_table.team_id}, + ) + check_team_key_model_specific_limits( + keys=keys, + team_table=team_table, + data=data, + ) + check_team_key_rpm_tpm_limits( + keys=keys, + team_table=team_table, + data=data, + ) + + +def check_org_key_model_specific_limits( + keys: List[LiteLLM_VerificationToken], + org_table: LiteLLM_OrganizationTable, + data: Union[GenerateKeyRequest, UpdateKeyRequest], +) -> None: + """ + Check if the organization key is allocating model specific limits. If so, raise an error if we're overallocating. + """ + # Get org limits from budget table if available + entity_rpm_limit = None + entity_tpm_limit = None + entity_model_rpm_limit_dict = {} + entity_model_tpm_limit_dict = {} + + if org_table.litellm_budget_table is not None: + entity_rpm_limit = org_table.litellm_budget_table.rpm_limit + entity_tpm_limit = org_table.litellm_budget_table.tpm_limit + + if org_table.metadata: + entity_model_rpm_limit_dict = org_table.metadata.get("model_rpm_limit", {}) + entity_model_tpm_limit_dict = org_table.metadata.get("model_tpm_limit", {}) + + _check_key_model_specific_limits( + keys=keys, + data=data, + entity_rpm_limit=entity_rpm_limit, + entity_tpm_limit=entity_tpm_limit, + entity_model_rpm_limit_dict=entity_model_rpm_limit_dict, + entity_model_tpm_limit_dict=entity_model_tpm_limit_dict, + entity_type="organization", + ) + + +def check_org_key_rpm_tpm_limits( + keys: List[LiteLLM_VerificationToken], + org_table: LiteLLM_OrganizationTable, + data: Union[GenerateKeyRequest, UpdateKeyRequest], +) -> None: + """ + Check if the organization key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. + """ + # Get org limits from budget table if available + entity_rpm_limit = None + entity_tpm_limit = None + + if org_table.litellm_budget_table is not None: + entity_rpm_limit = org_table.litellm_budget_table.rpm_limit + entity_tpm_limit = org_table.litellm_budget_table.tpm_limit + + _check_key_rpm_tpm_limits( + keys=keys, + data=data, + entity_rpm_limit=entity_rpm_limit, + entity_tpm_limit=entity_tpm_limit, + entity_type="organization", + ) + + +async def _check_org_key_limits( + org_table: LiteLLM_OrganizationTable, + data: Union[GenerateKeyRequest, UpdateKeyRequest], + prisma_client: PrismaClient, +) -> None: + """ + Check if the organization key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + + Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" + """ + + rpm_limit_type = getattr(data, "rpm_limit_type", None) or ( + data.metadata.get("rpm_limit_type", None) if data.metadata else None + ) + tpm_limit_type = getattr(data, "tpm_limit_type", None) or ( + data.metadata.get("tpm_limit_type", None) if data.metadata else None + ) + + if ( + tpm_limit_type != "guaranteed_throughput" + and rpm_limit_type != "guaranteed_throughput" + ): + return + # get all organization keys + # calculate allocated tpm/rpm limit + # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"organization_id": org_table.organization_id}, + ) + check_org_key_model_specific_limits( + keys=keys, + org_table=org_table, + data=data, + ) + check_org_key_rpm_tpm_limits( + keys=keys, + org_table=org_table, + data=data, + ) + + +@router.post( + "/key/generate", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], + response_model=GenerateKeyResponse, +) +@management_endpoint_wrapper +async def generate_key_fn( + data: GenerateKeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Generate an API key based on the provided data. + + Docs: https://docs.litellm.ai/docs/proxy/virtual_keys + + Parameters: + - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). + - key_alias: Optional[str] - User defined key alias + - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + - team_id: Optional[str] - The team id of the key + - user_id: Optional[str] - The user id of the key + - organization_id: Optional[str] - The organization id of the key. If not set, and team_id is set, the organization id will be the same as the team id. If conflict, an error will be raised. + - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. + - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) + - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models + - config: Optional[dict] - any key-specific configs, overrides config in config.yaml + - spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend + - send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key + - max_budget: Optional[float] - Specify max budget for a given key. + - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). + - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. + - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } + - guardrails: Optional[List[str]] - List of active guardrails for the key + - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. + - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. + - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. + - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. + - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". + - rpm_limit_type: Optional[str] - Type of rpm limit. Options: "best_effort_throughput" (no error if we're overallocating rpm), "guaranteed_throughput" (raise an error if we're overallocating rpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". + - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request + - blocked: Optional[bool] - Whether the key is blocked. + - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) + - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. + - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). + - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. + - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) + - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. + - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] + - allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored. + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. + - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default". + - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. + - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated) + - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings. + + Examples: + + 1. Allow users to turn on/off pii masking + + ```bash + curl --location 'http://0.0.0.0:4000/key/generate' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "permissions": {"allow_pii_controls": true} + }' + ``` + + Returns: + - key: (str) The generated api key + - expires: (datetime) Datetime object for when key expires. + - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. + """ + try: + from litellm.proxy._types import CommonProxyErrors + from litellm.proxy.proxy_server import ( + prisma_client, + user_api_key_cache, + user_custom_key_generate, + ) + + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + verbose_proxy_logger.debug("entered /key/generate") + + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, + ) + if data.soft_budget is not None and data.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, + ) + + if user_custom_key_generate is not None: + if asyncio.iscoroutinefunction(user_custom_key_generate): + result = await user_custom_key_generate(data) # type: ignore + else: + raise ValueError("user_custom_key_generate must be a coroutine") + decision = result.get("decision", True) + message = result.get("message", "Authentication Failed - Custom Auth Rule") + if not decision: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=message + ) + team_table: Optional[LiteLLM_TeamTableCachedObj] = None + if data.team_id is not None: + try: + team_table = await get_team_object( + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + check_db_only=True, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Error getting team object in `/key/generate`: {e}" + ) + + key_generation_check( + team_table=team_table, + user_api_key_dict=user_api_key_dict, + data=data, + route=KeyManagementRoutes.KEY_GENERATE, + ) + + if team_table is not None: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=prisma_client, + ) + + return await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + team_table=team_table, + ) + + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + +@router.post( + "/key/service-account/generate", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def generate_service_account_key_fn( + data: GenerateKeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Generate a Service Account API key based on the provided data. This key does not belong to any user. It belongs to the team. + + Why use a service account key? + - Prevent key from being deleted when user is deleted. + - Apply team limits, not team member limits to key. + + Docs: https://docs.litellm.ai/docs/proxy/virtual_keys + + Parameters: + - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). + - key_alias: Optional[str] - User defined key alias + - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + - team_id: Optional[str] - The team id of the key + - user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key + - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. + - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) + - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models + - config: Optional[dict] - any key-specific configs, overrides config in config.yaml + - spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend + - send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key + - max_budget: Optional[float] - Specify max budget for a given key. + - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). + - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. + - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } + - guardrails: Optional[List[str]] - List of active guardrails for the key + - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. + - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. + - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" + - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" + - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request + - blocked: Optional[bool] - Whether the key is blocked. + - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) + - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. + - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). + - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) + - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. + Examples: + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + + + 1. Allow users to turn on/off pii masking + + ```bash + curl --location 'http://0.0.0.0:4000/key/generate' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "permissions": {"allow_pii_controls": true} + }' + ``` + + Returns: + - key: (str) The generated api key + - expires: (datetime) Datetime object for when key expires. + - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. + + """ + from litellm.proxy._types import CommonProxyErrors + from litellm.proxy.proxy_server import ( + prisma_client, + user_api_key_cache, + user_custom_key_generate, + ) + + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + await validate_team_id_used_in_service_account_request( + team_id=data.team_id, + prisma_client=prisma_client, + ) + + verbose_proxy_logger.debug("entered /key/generate") + + if user_custom_key_generate is not None: + if asyncio.iscoroutinefunction(user_custom_key_generate): + result = await user_custom_key_generate(data) # type: ignore + else: + raise ValueError("user_custom_key_generate must be a coroutine") + decision = result.get("decision", True) + message = result.get("message", "Authentication Failed - Custom Auth Rule") + if not decision: + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) + team_table: Optional[LiteLLM_TeamTableCachedObj] = None + if data.team_id is not None: + try: + team_table = await get_team_object( + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + check_db_only=True, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Error getting team object in `/key/generate`: {e}" + ) + team_table = None + + if team_table is not None: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=prisma_client, + ) + + key_generation_check( + team_table=team_table, + user_api_key_dict=user_api_key_dict, + data=data, + route=KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT, + ) + + data.user_id = None # do not allow user_id to be set for service account keys + + return await _common_key_generation_helper( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + team_table=team_table, + ) + + +def prepare_metadata_fields( + data: BaseModel, non_default_values: dict, existing_metadata: dict +) -> dict: + """ + Check LiteLLM_ManagementEndpoint_MetadataFields (proxy/_types.py) for fields that are allowed to be updated + """ + if "metadata" not in non_default_values: # allow user to set metadata to none + non_default_values["metadata"] = existing_metadata.copy() + + casted_metadata = cast(dict, non_default_values["metadata"]) + + data_json = data.model_dump(exclude_unset=True, exclude_none=True) + + try: + for k, v in data_json.items(): + if k in LiteLLM_ManagementEndpoint_MetadataFields: + if isinstance(v, datetime): + casted_metadata[k] = v.isoformat() + else: + casted_metadata[k] = v + if k in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + from litellm.proxy.utils import _premium_user_check + + _premium_user_check(k) + casted_metadata[k] = v + + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - {}".format( + str(e) + ) + ) + + non_default_values["metadata"] = casted_metadata + return non_default_values + + +async def prepare_key_update_data( + data: Union[UpdateKeyRequest, RegenerateKeyRequest], + existing_key_row: LiteLLM_VerificationToken, +): + data_json: dict = data.model_dump(exclude_unset=True) + data_json.pop("key", None) + data_json.pop("new_key", None) + if ( + data.metadata is not None + and data.metadata.get("service_account_id") is not None + and (data.team_id or existing_key_row.team_id) is None + ): + raise HTTPException( + status_code=400, + detail="team_id is required for service account keys. Please specify `team_id` in the request body.", + ) + non_default_values = {} + # ADD METADATA FIELDS + # Set Management Endpoint Metadata Fields + for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=data, + field_name=field, + value=getattr(data, field), + ) + for k, v in data_json.items(): + if ( + k in LiteLLM_ManagementEndpoint_MetadataFields + or k in LiteLLM_ManagementEndpoint_MetadataFields_Premium + ): + continue + non_default_values[k] = v + + if "duration" in non_default_values: + duration = non_default_values.pop("duration") + if duration == "-1": + # Set expires to None to indicate the key never expires + non_default_values["expires"] = None + elif duration and (isinstance(duration, str)) and len(duration) > 0: + duration_s = duration_in_seconds(duration=duration) + expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + non_default_values["expires"] = expires + + if "budget_duration" in non_default_values: + budget_duration = non_default_values.pop("budget_duration") + if ( + budget_duration + and (isinstance(budget_duration, str)) + and len(budget_duration) > 0 + ): + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + key_reset_at = get_budget_reset_time(budget_duration=budget_duration) + non_default_values["budget_reset_at"] = key_reset_at + non_default_values["budget_duration"] = budget_duration + + if "object_permission" in non_default_values: + non_default_values = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + ) + + _metadata = existing_key_row.metadata or {} + + # validate model_max_budget + if "model_max_budget" in non_default_values: + validate_model_max_budget(non_default_values["model_max_budget"]) + + # Serialize router_settings to JSON if present + if ( + "router_settings" in non_default_values + and non_default_values["router_settings"] is not None + ): + non_default_values["router_settings"] = safe_dumps( + non_default_values["router_settings"] + ) + + non_default_values = prepare_metadata_fields( + data=data, non_default_values=non_default_values, existing_metadata=_metadata + ) + + return non_default_values + + +async def _handle_update_object_permission( + data_json: dict, + existing_key_row: LiteLLM_VerificationToken, +) -> dict: + """ + Handle the update of object permission. + """ + from litellm.proxy.proxy_server import prisma_client + + # Use the common helper to handle the object permission update + object_permission_id = await handle_update_object_permission_common( + data_json=data_json, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + ) + + # Add the object_permission_id to data_json if one was created/updated + if object_permission_id is not None: + data_json["object_permission_id"] = object_permission_id + verbose_proxy_logger.debug( + f"updated object_permission_id: {object_permission_id}" + ) + + return data_json + + +def is_different_team( + data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken +) -> bool: + if data.team_id is None: + return False + if existing_key_row.team_id is None: + return True + return data.team_id != existing_key_row.team_id + + +def _validate_max_budget(max_budget: Optional[float]) -> None: + """ + Validate that max_budget is not negative. + + Args: + max_budget: The max_budget value to validate + + Raises: + HTTPException: If max_budget is negative + """ + if max_budget is not None and max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {max_budget}"}, + ) + + +async def _get_and_validate_existing_key( + token: str, prisma_client: Optional[PrismaClient] +) -> LiteLLM_VerificationToken: + """ + Get existing key from database and validate it exists. + + Args: + token: The key token to look up + prisma_client: Prisma client instance + + Returns: + LiteLLM_VerificationToken: The existing key row + + Raises: + HTTPException: If key is not found + """ + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + + existing_key_row = await prisma_client.get_data( + token=token, + table_name="key", + query_type="find_unique", + ) + + if existing_key_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Key not found: {token}"}, + ) + + return existing_key_row + + +async def _process_single_key_update( + key_update_item: BulkUpdateKeyRequestItem, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: Any, + llm_router: Optional[Router], +) -> Dict[str, Any]: + """ + Process a single key update with all validations and checks. + + This function encapsulates all the logic for updating a single key, + including validation, permission checks, team checks, and database updates. + + Args: + key_update_item: The key update request item + user_api_key_dict: The authenticated user's API key info + litellm_changed_by: Optional header for tracking who made the change + prisma_client: Prisma client instance + user_api_key_cache: User API key cache + proxy_logging_obj: Proxy logging object + llm_router: LLM router instance + + Returns: + Dict containing the updated key information + + Raises: + HTTPException: For various validation and permission errors + """ + # Validate max_budget + _validate_max_budget(key_update_item.max_budget) + + # Get and validate existing key + existing_key_row = await _get_and_validate_existing_key( + token=key_update_item.key, + prisma_client=prisma_client, + ) + + # Check team member permissions + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=existing_key_row, + user_api_key_cache=user_api_key_cache, + ) + + # Create UpdateKeyRequest from BulkUpdateKeyRequestItem + update_key_request = UpdateKeyRequest( + key=key_update_item.key, + budget_id=key_update_item.budget_id, + max_budget=key_update_item.max_budget, + team_id=key_update_item.team_id, + tags=key_update_item.tags, + ) + + # Get team object and check team limits if team_id is provided + team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + if update_key_request.team_id is not None: + team_obj = await get_team_object( + team_id=update_key_request.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_obj is not None and prisma_client is not None: + await _check_team_key_limits( + team_table=team_obj, + data=update_key_request, + prisma_client=prisma_client, + ) + + # Validate team change if team is being changed + if is_different_team(data=update_key_request, existing_key_row=existing_key_row): + if llm_router is None: + raise HTTPException( + status_code=400, + detail={ + "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." + }, + ) + if team_obj is None: + raise HTTPException( + status_code=500, + detail={"error": "Team object not found for team change validation"}, + ) + validate_key_team_change( + key=existing_key_row, + team=team_obj, + change_initiated_by=user_api_key_dict, + llm_router=llm_router, + ) + + # Prepare update data + non_default_values = await prepare_key_update_data( + data=update_key_request, existing_key_row=existing_key_row + ) + + # Update key in database + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + + _data = {**non_default_values, "token": key_update_item.key} + response = await prisma_client.update_data(token=key_update_item.key, data=_data) + + # Delete cache + await _delete_cache_key_object( + hashed_token=hash_token(key_update_item.key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # Trigger async hook + asyncio.create_task( + KeyManagementEventHooks.async_key_updated_hook( + data=update_key_request, + existing_key_row=existing_key_row, + response=response, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + ) + + if response is None: + raise ValueError("Failed to update key got response = None") + + # Extract and format updated key info + updated_key_info = response.get("data", {}) + if hasattr(updated_key_info, "model_dump"): + updated_key_info = updated_key_info.model_dump() + elif hasattr(updated_key_info, "dict"): + updated_key_info = updated_key_info.dict() + + updated_key_info.pop("token", None) + + return updated_key_info + + +@router.post( + "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] +) +@management_endpoint_wrapper +async def update_key_fn( + request: Request, + data: UpdateKeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Update an existing API key's parameters. + + Parameters: + - key: str - The key to update + - key_alias: Optional[str] - User-friendly key alias + - user_id: Optional[str] - User ID associated with key + - team_id: Optional[str] - Team ID associated with key + - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. + - models: Optional[list] - Model_name's a user is allowed to call + - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) + - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. + - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) + - spend: Optional[float] - Amount spent by key + - max_budget: Optional[float] - Max budget for key + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. + - max_parallel_requests: Optional[int] - Rate limit for parallel requests + - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} + - tpm_limit: Optional[int] - Tokens per minute limit + - rpm_limit: Optional[int] - Requests per minute limit + - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} + - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} + - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" + - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" + - allowed_cache_controls: Optional[list] - List of allowed cache control values + - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) or "-1" to never expire + - permissions: Optional[dict] - Key-specific permissions + - send_invite_email: Optional[bool] - Send invite email to user_id + - guardrails: Optional[List[str]] - List of active guardrails for the key + - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. + - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. + - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. + - blocked: Optional[bool] - Whether the key is blocked + - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) + - config: Optional[dict] - [DEPRECATED PARAM] Key-specific config. + - temp_budget_increase: Optional[float] - Temporary budget increase for the key (Enterprise only). + - temp_budget_expiry: Optional[str] - Expiry time for the temporary budget increase (Enterprise only). + - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] + - allowed_passthrough_routes: Optional[list] - List of allowed pass through routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through routes the key can access, without specifying the routes. If allowed_routes is specified, allowed_passthrough_routes is ignored. + - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. + - auto_rotate: Optional[bool] - Whether this key should be automatically rotated + - rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings. + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/key/update' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "key": "sk-1234", + "key_alias": "my-key", + "user_id": "user-1234", + "team_id": "team-1234", + "max_budget": 100, + "metadata": {"any_key": "any-val"}, + }' + ``` + """ + from litellm.proxy.proxy_server import ( + llm_router, + premium_user, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + try: + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, + ) + + data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) + key = data_json.pop("key") + + # get the row from db + if prisma_client is None: + raise Exception("Not connected to DB!") + + existing_key_row = await prisma_client.get_data( + token=data.key, table_name="key", query_type="find_unique" + ) + + if existing_key_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) + + ## sanity check - prevent non-proxy admin user from updating key to belong to a different user + if ( + data.user_id is not None + and data.user_id != existing_key_row.user_id + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + ): + raise HTTPException( + status_code=403, + detail=f"User={data.user_id} is not allowed to update key={key} to belong to user={existing_key_row.user_id}", + ) + + common_key_access_checks( + user_api_key_dict=user_api_key_dict, + data=data, + user_id=existing_key_row.user_id, + llm_router=llm_router, + premium_user=premium_user, + ) + + # check if user has permission to update key + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=existing_key_row, + user_api_key_cache=user_api_key_cache, + ) + + # Only check team limits if key has a team_id + team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + if data.team_id is not None: + team_obj = await get_team_object( + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_obj is not None: + await _check_team_key_limits( + team_table=team_obj, + data=data, + prisma_client=prisma_client, + ) + + # if team change - check if this is possible + if is_different_team(data=data, existing_key_row=existing_key_row): + if llm_router is None: + raise HTTPException( + status_code=400, + detail={ + "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." + }, + ) + # team_obj should be set since is_different_team() returns True only when data.team_id is not None + if team_obj is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Team object not found for team change validation" + }, + ) + validate_key_team_change( + key=existing_key_row, + team=team_obj, + change_initiated_by=user_api_key_dict, + llm_router=llm_router, + ) + + # Set Management Endpoint Metadata Fields + + non_default_values = await prepare_key_update_data( + data=data, existing_key_row=existing_key_row + ) + + await _enforce_unique_key_alias( + key_alias=non_default_values.get("key_alias", None), + prisma_client=prisma_client, + existing_key_token=existing_key_row.token, + ) + + # Handle rotation fields if auto_rotate is being enabled + _set_key_rotation_fields( + non_default_values, + non_default_values.get("auto_rotate", False), + non_default_values.get("rotation_interval"), + ) + + _data = {**non_default_values, "token": key} + response = await prisma_client.update_data(token=key, data=_data) + + # Delete - key from cache, since it's been updated! + # key updated - a new model could have been added to this key. it should not block requests after this is done + await _delete_cache_key_object( + hashed_token=hash_token(key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + asyncio.create_task( + KeyManagementEventHooks.async_key_updated_hook( + data=data, + existing_key_row=existing_key_row, + response=response, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + ) + + if response is None: + raise ValueError("Failed to update key got response = None") + + return {"key": key, **response["data"]} + # update based on remaining passed in values + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.update_key_fn(): Exception occured - {}".format( + str(e) + ) + ) + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Authentication Error({str(e)})"), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=status.HTTP_400_BAD_REQUEST, + ) + + +@router.post( + "/key/bulk_update", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], + response_model=BulkUpdateKeyResponse, +) +@management_endpoint_wrapper +async def bulk_update_keys( + data: BulkUpdateKeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Bulk update multiple keys at once. + + This endpoint allows updating multiple keys in a single request. Each key update + is processed independently - if some updates fail, others will still succeed. + + Parameters: + - keys: List[BulkUpdateKeyRequestItem] - List of key update requests, each containing: + - key: str - The key identifier (token) to update + - budget_id: Optional[str] - Budget ID associated with the key + - max_budget: Optional[float] - Max budget for key + - team_id: Optional[str] - Team ID associated with key + - tags: Optional[List[str]] - Tags for organizing keys + + Returns: + - total_requested: int - Total number of keys requested for update + - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info + - failed_updates: List[FailedKeyUpdate] - List of failed updates with key_info and failed_reason + + Example request: + ```bash + curl --location 'http://0.0.0.0:4000/key/bulk_update' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "keys": [ + { + "key": "sk-1234", + "max_budget": 100.0, + "team_id": "team-123", + "tags": ["production", "api"] + }, + { + "key": "sk-5678", + "budget_id": "budget-456", + "tags": ["staging"] + } + ] + }' + ``` + """ + from litellm.proxy.proxy_server import ( + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins can perform bulk key updates"}, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + + if not data.keys: + raise HTTPException( + status_code=400, + detail={"error": "No keys provided for update"}, + ) + + MAX_BATCH_SIZE = 500 + if len(data.keys) > MAX_BATCH_SIZE: + raise HTTPException( + status_code=400, + detail={ + "error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys." + }, + ) + + successful_updates: List[SuccessfulKeyUpdate] = [] + failed_updates: List[FailedKeyUpdate] = [] + + for key_update_item in data.keys: + try: + # Process single key update using reusable function + updated_key_info = await _process_single_key_update( + key_update_item=key_update_item, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + ) + + successful_updates.append( + SuccessfulKeyUpdate( + key=key_update_item.key, + key_info=updated_key_info, + ) + ) + + except Exception as e: + verbose_proxy_logger.exception( + f"Failed to update key {key_update_item.key}: {e}" + ) + + if isinstance(e, HTTPException): + error_detail = e.detail + if isinstance(error_detail, dict): + error_message = error_detail.get("error", str(e)) + else: + error_message = str(error_detail) + else: + error_message = str(e) + + key_info = None + try: + existing_key_row = await prisma_client.get_data( + token=key_update_item.key, + table_name="key", + query_type="find_unique", + ) + if existing_key_row is not None: + if hasattr(existing_key_row, "model_dump"): + key_info = existing_key_row.model_dump() + elif hasattr(existing_key_row, "dict"): + key_info = existing_key_row.dict() + if key_info: + key_info.pop("token", None) + except Exception: + pass + + failed_updates.append( + FailedKeyUpdate( + key=key_update_item.key, + key_info=key_info, + failed_reason=error_message, + ) + ) + + return BulkUpdateKeyResponse( + total_requested=len(data.keys), + successful_updates=successful_updates, + failed_updates=failed_updates, + ) + + +def validate_key_team_change( + key: LiteLLM_VerificationToken, + team: LiteLLM_TeamTable, + change_initiated_by: UserAPIKeyAuth, + llm_router: Router, +): + """ + Validate that a key can be moved to a new team. + + - The team must have access to the key's models + - The key's user_id must be a member of the team + - The key's tpm/rpm limit must be less than the team's tpm/rpm limit + - The person initiating the change must be either Proxy Admin or Team Admin + """ + # Check if the team has access to the key's models + if len(key.models) > 0: + for model in key.models: + can_team_access_model( + model=model, + team_object=team, + llm_router=llm_router, + ) + + # Check if the key's user_id is a member of the team + member_object = _get_user_in_team( + team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id + ) + if key.user_id is not None: + if not member_object: + raise HTTPException( + status_code=403, + detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", + ) + + # Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit + if key.tpm_limit is not None: + if team.tpm_limit and key.tpm_limit > team.tpm_limit: + raise HTTPException( + status_code=403, + detail=f"Key={key.token} has a tpm_limit={key.tpm_limit} which is greater than the team's tpm_limit={team.tpm_limit}.", + ) + if team.rpm_limit and key.rpm_limit and key.rpm_limit > team.rpm_limit: + raise HTTPException( + status_code=403, + detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.", + ) + + # Check if the person initiating the change is a Proxy Admin or Team Admin + if change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + elif _is_user_team_admin( + user_api_key_dict=change_initiated_by, + team_obj=team, + ): + return + # this teams member permissions allow updating a + elif TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( + team_member_object=member_object, + team_table=cast(LiteLLM_TeamTableCachedObj, team), + route=KeyManagementRoutes.KEY_UPDATE.value, + ): + return + else: + raise HTTPException( + status_code=403, + detail=f"User={change_initiated_by.user_id} is not a Proxy Admin or Team Admin for team={team.team_id}. Please ask your Proxy Admin to allow this action under 'Member Permissions' for this team.", + ) + + +@router.post( + "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)] +) +@management_endpoint_wrapper +async def delete_key_fn( + data: KeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Delete a key from the key management system. + + Parameters:: + - keys (List[str]): A list of keys or hashed keys to delete. Example {"keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]} + - key_aliases (List[str]): A list of key aliases to delete. Can be passed instead of `keys`.Example {"key_aliases": ["alias1", "alias2"]} + + Returns: + - deleted_keys (List[str]): A list of deleted keys. Example {"deleted_keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]} + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/key/delete' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "keys": ["sk-QWrxEynunsNpV1zT48HIrw"] + }' + ``` + + Raises: + HTTPException: If an error occurs during key deletion. + """ + try: + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise Exception("Not connected to DB!") + + # Normalize litellm_changed_by: if it's a Header object or not a string, convert to None + if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): + litellm_changed_by = None + + ## only allow user to delete keys they own + verbose_proxy_logger.debug( + f"user_api_key_dict.user_role: {user_api_key_dict.user_role}" + ) + + num_keys_to_be_deleted = 0 + deleted_keys = [] + if data.keys: + number_deleted_keys, _keys_being_deleted = await delete_verification_tokens( + tokens=data.keys, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + num_keys_to_be_deleted = len(data.keys) + deleted_keys = data.keys + elif data.key_aliases: + number_deleted_keys, _keys_being_deleted = await delete_key_aliases( + key_aliases=data.key_aliases, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + num_keys_to_be_deleted = len(data.key_aliases) + deleted_keys = data.key_aliases + else: + raise ValueError("Invalid request type") + + if number_deleted_keys is None: + raise ProxyException( + message="Failed to delete keys got None response from delete_verification_token", + type=ProxyErrorTypes.internal_server_error, + param="keys", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + verbose_proxy_logger.debug(f"/key/delete - deleted_keys={number_deleted_keys}") + + try: + assert num_keys_to_be_deleted == len(deleted_keys) + except Exception: + raise HTTPException( + status_code=400, + detail={ + "error": f"Not all keys passed in were deleted. This probably means you don't have access to delete all the keys passed in. Keys passed in={num_keys_to_be_deleted}, Deleted keys ={number_deleted_keys}" + }, + ) + + verbose_proxy_logger.debug( + f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}" + ) + + asyncio.create_task( + KeyManagementEventHooks.async_key_deleted_hook( + data=data, + keys_being_deleted=_keys_being_deleted, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + response=number_deleted_keys, + ) + ) + + return {"deleted_keys": deleted_keys} + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + +@router.post( + "/v2/key/info", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], + include_in_schema=False, +) +async def info_key_fn_v2( + data: Optional[KeyRequest] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Retrieve information about a list of keys. + + **New endpoint**. Currently admin only. + Parameters: + keys: Optional[list] = body parameter representing the key(s) in the request + user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + Returns: + Dict containing the key and its associated information + + Example Curl: + ``` + curl -X GET "http://0.0.0.0:4000/key/info" \ + -H "Authorization: Bearer sk-1234" \ + -d {"keys": ["sk-1", "sk-2", "sk-3"]} + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if prisma_client is None: + raise Exception( + "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + ) + if data is None: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": "Malformed request. No keys passed in."}, + ) + + key_info = await prisma_client.get_data( + token=data.keys, table_name="key", query_type="find_all" + ) + if key_info is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"message": "No keys found"}, + ) + filtered_key_info = [] + for k in key_info: + try: + k = k.model_dump() # noqa + except Exception: + # if using pydantic v1 + k = k.dict() + filtered_key_info.append(k) + return {"key": data.keys, "info": filtered_key_info} + + except Exception as e: + raise handle_exception_on_proxy(e) + + +@router.get( + "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)] +) +async def info_key_fn( + key: Optional[str] = fastapi.Query( + default=None, description="Key in the request parameters" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Retrieve information about a key. + Parameters: + key: Optional[str] = Query parameter representing the key in the request + user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + Returns: + Dict containing the key and its associated information + + Example Curl: + ``` + curl -X GET "http://0.0.0.0:4000/key/info?key=sk-test-example-key-123" \ +-H "Authorization: Bearer sk-1234" + ``` + + Example Curl - if no key is passed, it will use the Key Passed in Authorization Header + ``` + curl -X GET "http://0.0.0.0:4000/key/info" \ +-H "Authorization: Bearer sk-test-example-key-123" + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if prisma_client is None: + raise Exception( + "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + ) + + # default to using Auth token if no key is passed in + key = key or user_api_key_dict.api_key + hashed_key: Optional[str] = key + if key is not None: + hashed_key = _hash_token_if_needed(token=key) + key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_key}, # type: ignore + include={"litellm_budget_table": True}, + ) + if key_info is None: + raise ProxyException( + message="Key not found in database", + type=ProxyErrorTypes.not_found_error, + param="key", + code=status.HTTP_404_NOT_FOUND, + ) + + if ( + await _can_user_query_key_info( + user_api_key_dict=user_api_key_dict, + key=key, + key_info=key_info, + ) + is not True + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You are not allowed to access this key's info. Your role={}".format( + user_api_key_dict.user_role + ), + ) + ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## + try: + key_info = key_info.model_dump() # noqa + except Exception: + # if using pydantic v1 + key_info = key_info.dict() + key_info.pop("token") + return {"key": key, "info": key_info} + except Exception as e: + raise handle_exception_on_proxy(e) + + +def _check_model_access_group( + models: Optional[List[str]], llm_router: Optional[Router], premium_user: bool +) -> Literal[True]: + """ + if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user + + Return True if user is a premium user, False otherwise + """ + if models is None or llm_router is None: + return True + + for model in models: + if llm_router._is_model_access_group_for_wildcard_route( + model_access_group=model + ): + if not premium_user: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Setting a model access group on a wildcard model is only available for LiteLLM Enterprise users.{}".format( + CommonProxyErrors.not_premium_user.value + ) + }, + ) + + return True + + +async def generate_key_helper_fn( # noqa: PLR0915 + request_type: Literal[ + "user", "key" + ], # identifies if this request is from /user/new or /key/generate + duration: Optional[str] = None, + models: list = [], + aliases: dict = {}, + config: dict = {}, + spend: float = 0.0, + key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key + key_budget_duration: Optional[str] = None, + budget_id: Optional[float] = None, # budget id <-> LiteLLM_BudgetTable + soft_budget: Optional[ + float + ] = None, # soft_budget is used to set soft Budgets Per user + max_budget: Optional[float] = None, # max_budget is used to Budget Per user + blocked: Optional[bool] = None, + budget_duration: Optional[str] = None, # max_budget is used to Budget Per user + token: Optional[str] = None, + key: Optional[ + str + ] = None, # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself. + user_id: Optional[str] = None, + user_alias: Optional[str] = None, + team_id: Optional[str] = None, + user_email: Optional[str] = None, + user_role: Optional[str] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[dict] = {}, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + query_type: Literal["insert_data", "update_data"] = "insert_data", + update_key_values: Optional[dict] = None, + key_alias: Optional[str] = None, + allowed_cache_controls: Optional[list] = [], + permissions: Optional[dict] = {}, + model_max_budget: Optional[dict] = {}, + model_rpm_limit: Optional[dict] = None, + model_tpm_limit: Optional[dict] = None, + guardrails: Optional[list] = None, + policies: Optional[list] = None, + prompts: Optional[list] = None, + teams: Optional[list] = None, + organization_id: Optional[str] = None, + table_name: Optional[Literal["key", "user"]] = None, + send_invite_email: Optional[bool] = None, + created_by: Optional[str] = None, + updated_by: Optional[str] = None, + allowed_routes: Optional[list] = None, + sso_user_id: Optional[str] = None, + object_permission_id: Optional[ + str + ] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, + auto_rotate: Optional[bool] = None, + rotation_interval: Optional[str] = None, + router_settings: Optional[dict] = None, +): + from litellm.proxy.proxy_server import premium_user, prisma_client + + if prisma_client is None: + raise Exception( + "Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys " + ) + + if token is None: + if key is not None: + token = key + else: + token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" + + if duration is None: # allow tokens that never expire + expires = None + else: + # Add duration to current time for exact expiration (not standardized reset time) + duration_seconds = duration_in_seconds(duration) + expires = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds) + + if key_budget_duration is None: # one-time budget + key_reset_at = None + else: + key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration) + + if budget_duration is None: # one-time budget + reset_at = None + else: + reset_at = get_budget_reset_time(budget_duration=budget_duration) + + aliases_json = json.dumps(aliases) + config_json = json.dumps(config) + permissions_json = json.dumps(permissions) + router_settings_json = ( + safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) + ) + + # Add model_rpm_limit and model_tpm_limit to metadata + if model_rpm_limit is not None: + metadata = metadata or {} + metadata["model_rpm_limit"] = model_rpm_limit + if model_tpm_limit is not None: + metadata = metadata or {} + metadata["model_tpm_limit"] = model_tpm_limit + if guardrails is not None: + metadata = metadata or {} + metadata["guardrails"] = guardrails + if policies is not None: + metadata = metadata or {} + metadata["policies"] = policies + if prompts is not None: + metadata = metadata or {} + metadata["prompts"] = prompts + + metadata_json = json.dumps(metadata) + validate_model_max_budget(model_max_budget) + model_max_budget_json = json.dumps(model_max_budget) + user_role = user_role + tpm_limit = tpm_limit + rpm_limit = rpm_limit + allowed_cache_controls = allowed_cache_controls + + try: + # Create a new verification token (you may want to enhance this logic based on your needs) + + user_data = { + "max_budget": max_budget, + "user_email": user_email, + "user_id": user_id, + "user_alias": user_alias, + "team_id": team_id, + "organization_id": organization_id, + "user_role": user_role, + "spend": spend, + "models": models, + "metadata": metadata_json, + "max_parallel_requests": max_parallel_requests, + "tpm_limit": tpm_limit, + "rpm_limit": rpm_limit, + "budget_duration": budget_duration, + "budget_reset_at": reset_at, + "allowed_cache_controls": allowed_cache_controls, + "sso_user_id": sso_user_id, + "object_permission_id": object_permission_id, + } + if teams is not None: + user_data["teams"] = teams + key_data = { + "token": token, + "key_alias": key_alias, + "expires": expires, + "models": models, + "aliases": aliases_json, + "config": config_json, + "spend": spend, + "max_budget": key_max_budget, + "user_id": user_id, + "team_id": team_id, + "max_parallel_requests": max_parallel_requests, + "metadata": metadata_json, + "tpm_limit": tpm_limit, + "rpm_limit": rpm_limit, + "budget_duration": key_budget_duration, + "budget_reset_at": key_reset_at, + "allowed_cache_controls": allowed_cache_controls, + "permissions": permissions_json, + "model_max_budget": model_max_budget_json, + "organization_id": organization_id, + "budget_id": budget_id, + "blocked": blocked, + "created_by": created_by, + "updated_by": updated_by, + "allowed_routes": allowed_routes or [], + "object_permission_id": object_permission_id, + "router_settings": router_settings_json, + } + + # Add rotation fields if auto_rotate is enabled + _set_key_rotation_fields( + data=key_data, + auto_rotate=auto_rotate or False, + rotation_interval=rotation_interval, + ) + + if ( + get_secret("DISABLE_KEY_NAME", False) is True + ): # allow user to disable storing abbreviated key name (shown in UI, to help figure out which key spent how much) + pass + else: + key_data["key_name"] = abbreviate_api_key(api_key=token) + saved_token = copy.deepcopy(key_data) + if isinstance(saved_token["aliases"], str): + saved_token["aliases"] = json.loads(saved_token["aliases"]) + if isinstance(saved_token["config"], str): + saved_token["config"] = json.loads(saved_token["config"]) + if isinstance(saved_token["metadata"], str): + saved_token["metadata"] = json.loads(saved_token["metadata"]) + if isinstance(saved_token["permissions"], str): + if ( + "get_spend_routes" in saved_token["permissions"] + and premium_user is not True + ): + raise ValueError( + "get_spend_routes permission is only available for LiteLLM Enterprise users" + ) + + saved_token["permissions"] = json.loads(saved_token["permissions"]) + if isinstance(saved_token["model_max_budget"], str): + saved_token["model_max_budget"] = json.loads( + saved_token["model_max_budget"] + ) + router_settings = cast(Optional[dict], saved_token.get("router_settings")) + if router_settings is not None and isinstance(router_settings, str): + try: + saved_token["router_settings"] = yaml.safe_load(router_settings) + except yaml.YAMLError: + # If it's not valid JSON/YAML, keep as is or set to empty dict + saved_token["router_settings"] = {} + + if saved_token.get("expires", None) is not None and isinstance( + saved_token["expires"], datetime + ): + saved_token["expires"] = saved_token["expires"].isoformat() + if prisma_client is not None: + if ( + table_name is None or table_name == "user" + ): # do not auto-create users for `/key/generate` + ## CREATE USER (If necessary) + if query_type == "insert_data": + user_row = await prisma_client.insert_data( + data=user_data, table_name="user" + ) + + if user_row is None: + raise Exception("Failed to create user") + ## use default user model list if no key-specific model list provided + if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore + key_data["models"] = user_row.models # type: ignore + elif query_type == "update_data": + user_row = await prisma_client.update_data( + data=user_data, + table_name="user", + update_key_values=update_key_values, + ) + if table_name is not None and table_name == "user": + # do not create a key if table name is set to just 'user' + # we only need to ensure this exists in the user table + # the LiteLLM_VerificationToken table will increase in size if we don't do this check + return user_data + + ## CREATE KEY + verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data) + create_key_response = await prisma_client.insert_data( + data=key_data, table_name="key" + ) + + key_data["token_id"] = getattr(create_key_response, "token", None) + key_data["litellm_budget_table"] = getattr( + create_key_response, "litellm_budget_table", None + ) + key_data["created_at"] = getattr(create_key_response, "created_at", None) + key_data["updated_at"] = getattr(create_key_response, "updated_at", None) + + # Deserialize router_settings from JSON string to dict for response + router_settings_value = key_data.get("router_settings") + if router_settings_value is not None and isinstance( + router_settings_value, str + ): + try: + key_data["router_settings"] = yaml.safe_load(router_settings_value) + except yaml.YAMLError: + # If it's not valid JSON/YAML, keep as is or set to empty dict + key_data["router_settings"] = {} + except Exception as e: + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {}".format( + str(e) + ) + ) + verbose_proxy_logger.debug(traceback.format_exc()) + if isinstance(e, HTTPException): + raise e + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": "Internal Server Error."}, + ) + + # Add budget related info in key_data - this ensures it's returned + key_data["budget_id"] = budget_id + + if request_type == "user": + # if this is a /user/new request update the key_date with user_data fields + key_data.update(user_data) + + return key_data + + +async def _team_key_deletion_check( + user_api_key_dict: UserAPIKeyAuth, + key_info: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: DualCache, +): + is_team_key = _is_team_key(data=key_info) + + if is_team_key and key_info.team_id is not None: + team_table = await get_team_object( + team_id=key_info.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + if ( + litellm.key_generation_settings is not None + and "team_key_generation" in litellm.key_generation_settings + ): + _team_key_generation = litellm.key_generation_settings[ + "team_key_generation" + ] + else: + _team_key_generation = TeamUIKeyGenerationConfig( + allowed_team_member_roles=["admin", "user"], + ) + # check if user is team admin + if team_table is not None: + return _team_key_operation_team_member_check( + assigned_user_id=user_api_key_dict.user_id, + team_table=team_table, + user_api_key_dict=user_api_key_dict, + team_key_generation=_team_key_generation, + route=KeyManagementRoutes.KEY_DELETE, + ) + else: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={ + "error": f"Team not found in db, and user not proxy admin. Team id = {key_info.team_id}" + }, + ) + return False + + +async def can_modify_verification_token( + key_info: LiteLLM_VerificationToken, + user_api_key_cache: DualCache, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, +) -> bool: + """ + Check if user has permission to modify (delete/regenerate) a verification token. + + Rules: + - Proxy admin can modify any key + - For team keys: only team admin or key owner can modify + - For personal keys: only key owner can modify + + Args: + key_info: The verification token to check + user_api_key_cache: Cache for user API keys + user_api_key_dict: The user making the request + prisma_client: Prisma client for database access + + Returns: + True if user can modify the key, False otherwise + """ + is_team_key = _is_team_key(data=key_info) + + # 1. Proxy admin can modify any key + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return True + + # 2. For team keys: only team admin or key owner can modify + if is_team_key and key_info.team_id is not None: + # Get team object to check if user is team admin + team_table = await get_team_object( + team_id=key_info.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_table is None: + return False + + # Check if user is team admin + if _is_user_team_admin( + user_api_key_dict=user_api_key_dict, + team_obj=team_table, + ): + return True + + # Check if the key belongs to the user (they own it) + if ( + key_info.user_id is not None + and key_info.user_id == user_api_key_dict.user_id + ): + return True + + # Not team admin and doesn't own the key + return False + + # 3. For personal keys: only key owner can modify + if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: + return True + + # Default: deny + return False + + +async def delete_verification_tokens( + tokens: List, + user_api_key_cache: DualCache, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str] = None, +) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: + """ + Helper that deletes the list of tokens from the database + + - check if user is proxy admin + - check if user is team admin and key is a team key + + Args: + tokens: List of tokens to delete + user_id: Optional user_id to filter by + + Returns: + Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: + Optional[Dict]: + - Number of deleted tokens + List[LiteLLM_VerificationToken]: + - List of keys being deleted, this contains information about the key_alias, token, and user_id being deleted, + this is passed down to the KeyManagementEventHooks to delete the keys from the secret manager and handle audit logs + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if prisma_client: + tokens = [_hash_token_if_needed(token=key) for key in tokens] + _keys_being_deleted: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} + ) + + if len(_keys_being_deleted) == 0: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": "No keys found"}, + ) + + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + authorized_keys = _keys_being_deleted + else: + authorized_keys = [] + for key in _keys_being_deleted: + if await can_modify_verification_token( + key_info=key, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ): + authorized_keys.append(key) + else: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "You are not authorized to delete this key" + }, + ) + await _persist_deleted_verification_tokens( + keys=authorized_keys, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + deleted_tokens = await prisma_client.delete_data(tokens=tokens) + else: + deletion_tasks = [ + prisma_client.delete_data(tokens=[key.token]) + for key in authorized_keys + ] + await asyncio.gather(*deletion_tasks) + + deleted_tokens = [key.token for key in authorized_keys] + if len(deleted_tokens) != len(tokens): + failed_tokens = [ + token for token in tokens if token not in deleted_tokens + ] + raise Exception( + "Failed to delete all tokens. Failed to delete tokens: " + + str(failed_tokens) + ) + else: + raise Exception("DB not connected. prisma_client is None") + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - {}".format( + str(e) + ) + ) + verbose_proxy_logger.debug(traceback.format_exc()) + raise e + + for key in tokens: + user_api_key_cache.delete_cache(key) + # remove hash token from cache + hashed_token = hash_token(cast(str, key)) + user_api_key_cache.delete_cache(hashed_token) + + return {"deleted_keys": deleted_tokens}, _keys_being_deleted + + +def _transform_verification_tokens_to_deleted_records( + keys: List[LiteLLM_VerificationToken], + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str] = None, +) -> List[Dict[str, Any]]: + """Transform verification tokens into deleted token records ready for persistence.""" + if not keys: + return [] + + deleted_at = datetime.now(timezone.utc) + records = [] + for key in keys: + key_payload = key.model_dump() + deleted_record = LiteLLM_DeletedVerificationToken( + **key_payload, + deleted_at=deleted_at, + deleted_by=user_api_key_dict.user_id, + deleted_by_api_key=user_api_key_dict.api_key, + litellm_changed_by=litellm_changed_by, + ) + record = deleted_record.model_dump() + + # Map org_id to organization_id (model uses org_id, but schema expects organization_id) + org_id_value = record.pop("org_id", None) + if org_id_value is not None: + record["organization_id"] = org_id_value + + for json_field in [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ]: + if json_field in record and record[json_field] is not None: + record[json_field] = json.dumps(record[json_field]) + + for rel_key in ( + "litellm_budget_table", + "litellm_organization_table", + "object_permission", + "id", + ): + record.pop(rel_key, None) + + records.append(record) + + return records + + +async def _save_deleted_verification_token_records( + records: List[Dict[str, Any]], + prisma_client: PrismaClient, +) -> None: + """Save deleted verification token records to the database.""" + if not records: + return + await prisma_client.db.litellm_deletedverificationtoken.create_many(data=records) + + +async def _persist_deleted_verification_tokens( + keys: List[LiteLLM_VerificationToken], + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str] = None, +) -> None: + """Persist deleted verification token records by transforming and saving them.""" + records = _transform_verification_tokens_to_deleted_records( + keys=keys, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + await _save_deleted_verification_token_records( + records=records, + prisma_client=prisma_client, + ) + + +async def delete_key_aliases( + key_aliases: List[str], + user_api_key_cache: DualCache, + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str] = None, +) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: + _keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many( + where={"key_alias": {"in": key_aliases}} + ) + + tokens = [key.token for key in _keys_being_deleted] + return await delete_verification_tokens( + tokens=tokens, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + +async def _rotate_master_key( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + current_master_key: str, + new_master_key: str, +) -> None: + """ + Rotate the master key + + 1. Get the values from the DB + - Get models from DB + - Get config from DB + 2. Decrypt the values + - ModelTable + - [{"model_name": "str", "litellm_params": {}}] + - ConfigTable + 3. Encrypt the values with the new master key + 4. Update the values in the DB + """ + from litellm.proxy.proxy_server import proxy_config + + try: + models: Optional[ + List + ] = await prisma_client.db.litellm_proxymodeltable.find_many() + except Exception: + models = None + # 2. process model table + if models: + decrypted_models = proxy_config.decrypt_model_list_from_db(new_models=models) + verbose_proxy_logger.debug( + "ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models) + ) + new_models = [] + for model in decrypted_models: + new_model = await _add_model_to_db( + model_params=Deployment(**model), + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + new_encryption_key=new_master_key, + should_create_model_in_db=False, + ) + if new_model: + new_models.append(jsonify_object(new_model.model_dump())) + verbose_proxy_logger.debug("Resetting proxy model table") + await prisma_client.db.litellm_proxymodeltable.delete_many() + verbose_proxy_logger.debug("Creating %s models", len(new_models)) + await prisma_client.db.litellm_proxymodeltable.create_many( + data=new_models, + ) + # 3. process config table + try: + config = await prisma_client.db.litellm_config.find_many() + except Exception: + config = None + + if config: + """If environment_variables is found, decrypt it and encrypt it with the new master key""" + environment_variables_dict = {} + for c in config: + if c.param_name == "environment_variables": + environment_variables_dict = c.param_value + + if environment_variables_dict: + decrypted_env_vars = proxy_config._decrypt_and_set_db_env_variables( + environment_variables=environment_variables_dict + ) + encrypted_env_vars = proxy_config._encrypt_env_variables( + environment_variables=decrypted_env_vars, + new_encryption_key=new_master_key, + ) + + if encrypted_env_vars: + await prisma_client.db.litellm_config.update( + where={"param_name": "environment_variables"}, + data={"param_value": jsonify_object(encrypted_env_vars)}, + ) + + # 4. process MCP server table + await rotate_mcp_server_credentials_master_key( + prisma_client=prisma_client, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + new_master_key=new_master_key, + ) + + # 5. process credentials table + try: + credentials = await prisma_client.db.litellm_credentialstable.find_many() + except Exception: + credentials = None + if credentials: + from litellm.proxy.credential_endpoints.endpoints import update_db_credential + + for cred in credentials: + try: + decrypted_cred = proxy_config.decrypt_credentials(cred) + encrypted_cred = update_db_credential( + db_credential=cred, + updated_patch=decrypted_cred, + new_encryption_key=new_master_key, + ) + credential_object_jsonified = jsonify_object( + encrypted_cred.model_dump() + ) + await prisma_client.db.litellm_credentialstable.update( + where={"credential_name": cred.credential_name}, + data={ + **credential_object_jsonified, + "updated_by": user_api_key_dict.user_id, + }, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}" + ) + # Continue with next credential instead of failing entire rotation + continue + verbose_proxy_logger.debug( + f"Successfully re-encrypted {len(credentials)} credentials with new master key" + ) + + +def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: + if data and data.new_key is not None: + new_token = data.new_key + if not data.new_key.startswith("sk-"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key." + }, + ) + else: + new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" + return new_token + + +@router.post( + "/key/{key:path}/regenerate", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/key/regenerate", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def regenerate_key_fn( + key: Optional[str] = None, + data: Optional[RegenerateKeyRequest] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +) -> Optional[GenerateKeyResponse]: + """ + Regenerate an existing API key while optionally updating its parameters. + + Parameters: + - key: str (path parameter) - The key to regenerate + - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update + - key: Optional[str] - The key to regenerate. + - new_master_key: Optional[str] - The new master key to use, if key is the master key. + - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used. + - key_alias: Optional[str] - User-friendly key alias + - user_id: Optional[str] - User ID associated with key + - team_id: Optional[str] - Team ID associated with key + - models: Optional[list] - Model_name's a user is allowed to call + - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) + - spend: Optional[float] - Amount spent by key + - max_budget: Optional[float] - Max budget for key + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. + - max_parallel_requests: Optional[int] - Rate limit for parallel requests + - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} + - tpm_limit: Optional[int] - Tokens per minute limit + - rpm_limit: Optional[int] - Requests per minute limit + - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} + - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} + - allowed_cache_controls: Optional[list] - List of allowed cache control values + - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) + - permissions: Optional[dict] - Key-specific permissions + - guardrails: Optional[List[str]] - List of active guardrails for the key + - blocked: Optional[bool] - Whether the key is blocked + + + Returns: + - GenerateKeyResponse containing the new key and its updated parameters + + Example: + ```bash + curl --location --request POST 'http://localhost:4000/key/sk-1234/regenerate' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data-raw '{ + "max_budget": 100, + "metadata": {"team": "core-infra"}, + "models": ["gpt-4", "gpt-3.5-turbo"] + }' + ``` + + Note: This is an Enterprise feature. It requires a premium license to use. + """ + try: + from litellm.proxy.proxy_server import ( + hash_token, + master_key, + premium_user, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + is_master_key_regeneration = data and data.new_master_key is not None + + if ( + premium_user is not True and not is_master_key_regeneration + ): # allow master key regeneration for non-premium users + raise ValueError( + f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}" + ) + + # Check if key exists, raise exception if key is not in the DB + key = data.key if data and data.key else key + if not key: + raise HTTPException(status_code=400, detail={"error": "No key passed in."}) + ### 1. Create New copy that is duplicate of existing key + ###################################################################### + + # create duplicate of existing key + # set token = new token generated + # insert new token in DB + + # create hash of token + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": "DB not connected. prisma_client is None"}, + ) + + _is_master_key_valid = _is_master_key(api_key=key, _master_key=master_key) + + if master_key is not None and data and _is_master_key_valid: + if data.new_master_key is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "New master key is required."}, + ) + await _rotate_master_key( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + current_master_key=master_key, + new_master_key=data.new_master_key, + ) + return GenerateKeyResponse( + key=data.new_master_key, + token=data.new_master_key, + key_name=data.new_master_key, + expires=None, + ) + + if "sk" not in key: + hashed_api_key = key + else: + hashed_api_key = hash_token(key) + + _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_api_key}, + ) + if _key_in_db is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Key {key} not found."}, + ) + + # check if user has permission to regenerate key + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_REGENERATE, + prisma_client=prisma_client, + existing_key_row=_key_in_db, + user_api_key_cache=user_api_key_cache, + ) + + # check if user has ownership permission to regenerate key + if not await can_modify_verification_token( + key_info=_key_in_db, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "You are not authorized to regenerate this key"}, + ) + + verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) + + new_token = get_new_token(data=data) + + new_token_hash = hash_token(new_token) + new_token_key_name = f"sk-...{new_token[-4:]}" + + # Prepare the update data + update_data = { + "token": new_token_hash, + "key_name": new_token_key_name, + } + + non_default_values = {} + if data is not None: + # Update with any provided parameters from GenerateKeyRequest + non_default_values = await prepare_key_update_data( + data=data, existing_key_row=_key_in_db + ) + verbose_proxy_logger.debug("non_default_values: %s", non_default_values) + + update_data.update(non_default_values) + update_data = prisma_client.jsonify_object(data=update_data) + # Update the token in the database + updated_token = await prisma_client.db.litellm_verificationtoken.update( + where={"token": hashed_api_key}, + data=update_data, # type: ignore + ) + + updated_token_dict = {} + if updated_token is not None: + updated_token_dict = dict(updated_token) + + updated_token_dict["key"] = new_token + updated_token_dict["token_id"] = updated_token_dict.pop("token") + + ### 3. remove existing key entry from cache + ###################################################################### + + if hashed_api_key or key: + await _delete_cache_key_object( + hashed_token=hash_token(key), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + response = GenerateKeyResponse( + **updated_token_dict, + ) + + asyncio.create_task( + KeyManagementEventHooks.async_key_rotated_hook( + data=data, + existing_key_row=_key_in_db, + response=response, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + ) + + return response + except Exception as e: + verbose_proxy_logger.exception("Error regenerating key: %s", e) + raise handle_exception_on_proxy(e) + + +async def validate_key_list_check( + user_api_key_dict: UserAPIKeyAuth, + user_id: Optional[str], + team_id: Optional[str], + organization_id: Optional[str], + key_alias: Optional[str], + key_hash: Optional[str], + prisma_client: PrismaClient, +) -> Optional[LiteLLM_UserTable]: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return None + + if user_api_key_dict.user_id is None: + raise ProxyException( + message="You are not authorized to access this endpoint. No 'user_id' is associated with your API key.", + type=ProxyErrorTypes.bad_request_error, + param="user_id", + code=status.HTTP_403_FORBIDDEN, + ) + complete_user_info_db_obj: Optional[ + BaseModel + ] = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, + ) + + if complete_user_info_db_obj is None: + raise ProxyException( + message="You are not authorized to access this endpoint. No 'user_id' is associated with your API key.", + type=ProxyErrorTypes.bad_request_error, + param="user_id", + code=status.HTTP_403_FORBIDDEN, + ) + + complete_user_info = LiteLLM_UserTable(**complete_user_info_db_obj.model_dump()) + + # internal user can only see their own keys + if user_id: + if complete_user_info.user_id != user_id: + raise ProxyException( + message="You are not authorized to check another user's keys", + type=ProxyErrorTypes.bad_request_error, + param="user_id", + code=status.HTTP_403_FORBIDDEN, + ) + + if team_id: + if team_id not in complete_user_info.teams: + raise ProxyException( + message="You are not authorized to check this team's keys", + type=ProxyErrorTypes.bad_request_error, + param="team_id", + code=status.HTTP_403_FORBIDDEN, + ) + + if organization_id: + if ( + complete_user_info.organization_memberships is None + or organization_id + not in [ + membership.organization_id + for membership in complete_user_info.organization_memberships + ] + ): + raise ProxyException( + message="You are not authorized to check this organization's keys", + type=ProxyErrorTypes.bad_request_error, + param="organization_id", + code=status.HTTP_403_FORBIDDEN, + ) + + if key_hash: + try: + key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": key_hash}, + ) + except Exception: + raise ProxyException( + message="Key Hash not found.", + type=ProxyErrorTypes.bad_request_error, + param="key_hash", + code=status.HTTP_403_FORBIDDEN, + ) + can_user_query_key_info = await _can_user_query_key_info( + user_api_key_dict=user_api_key_dict, + key=key_hash, + key_info=key_info, + ) + if not can_user_query_key_info: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You are not allowed to access this key's info. Your role={}".format( + user_api_key_dict.user_role + ), + ) + return complete_user_info + + +async def get_admin_team_ids( + complete_user_info: Optional[LiteLLM_UserTable], + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, +) -> List[str]: + """ + Get all team IDs where the user is an admin. + """ + if complete_user_info is None: + return [] + # Get all teams that user is an admin of + teams: Optional[ + List[BaseModel] + ] = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} + ) + if teams is None: + return [] + + teams_pydantic_obj = [LiteLLM_TeamTable(**team.model_dump()) for team in teams] + + admin_team_ids = [ + team.team_id + for team in teams_pydantic_obj + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team) + ] + return admin_team_ids + + +@router.get( + "/key/list", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def list_keys( + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + page: int = Query(1, description="Page number", ge=1), + size: int = Query(10, description="Page size", ge=1, le=100), + user_id: Optional[str] = Query(None, description="Filter keys by user ID"), + team_id: Optional[str] = Query(None, description="Filter keys by team ID"), + organization_id: Optional[str] = Query( + None, description="Filter keys by organization ID" + ), + key_hash: Optional[str] = Query(None, description="Filter keys by key hash"), + key_alias: Optional[str] = Query(None, description="Filter keys by key alias"), + return_full_object: bool = Query(False, description="Return full key object"), + include_team_keys: bool = Query( + False, description="Include all keys for teams that user is an admin of." + ), + include_created_by_keys: bool = Query( + False, description="Include keys created by the user" + ), + sort_by: Optional[str] = Query( + default=None, + description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", + ), + sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), + expand: Optional[List[str]] = Query( + None, description="Expand related objects (e.g. 'user')" + ), + status: Optional[str] = Query( + None, description="Filter by status (e.g. 'deleted')" + ), +) -> KeyListResponseObject: + """ + List all keys for a given user / team / organization. + + Parameters: + expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) + status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. + + Returns: + { + "keys": List[str] or List[UserAPIKeyAuth], + "total_count": int, + "current_page": int, + "total_pages": int, + } + + When expand includes "user", each key object will include a "user" field with the associated user object. + Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter. + """ + try: + from litellm.proxy.proxy_server import prisma_client + + verbose_proxy_logger.debug("Entering list_keys function") + + if prisma_client is None: + verbose_proxy_logger.error("Database not connected") + raise Exception("Database not connected") + + # Validate status parameter + if status is not None and status != "deleted": + raise HTTPException( + status_code=400, + detail={ + "error": "Invalid status value. Currently only 'deleted' is supported." + }, + ) + + complete_user_info = await validate_key_list_check( + user_api_key_dict=user_api_key_dict, + user_id=user_id, + team_id=team_id, + organization_id=organization_id, + key_alias=key_alias, + key_hash=key_hash, + prisma_client=prisma_client, + ) + + if include_team_keys: + admin_team_ids = await get_admin_team_ids( + complete_user_info=complete_user_info, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + else: + admin_team_ids = None + + if user_id is None and user_api_key_dict.user_role not in [ + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ]: + user_id = user_api_key_dict.user_id + + response = await _list_key_helper( + prisma_client=prisma_client, + page=page, + size=size, + user_id=user_id, + team_id=team_id, + key_alias=key_alias, + key_hash=key_hash, + return_full_object=return_full_object, + organization_id=organization_id, + admin_team_ids=admin_team_ids, + include_created_by_keys=include_created_by_keys, + sort_by=sort_by, + sort_order=sort_order, + expand=expand, + status=status, + ) + + verbose_proxy_logger.debug("Successfully prepared response") + + return response + + except Exception as e: + verbose_proxy_logger.exception(f"Error in list_keys: {e}") + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"error({str(e)})"), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=getattr( + e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR + ), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@router.get( + "/key/aliases", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def key_aliases() -> Dict[str, List[str]]: + """ + Lists all key aliases + + Returns: + { + "aliases": List[str] + } + """ + try: + from litellm.proxy.proxy_server import prisma_client + + verbose_proxy_logger.debug("Entering key_aliases function") + + if prisma_client is None: + verbose_proxy_logger.error("Database not connected") + raise Exception("Database not connected") + + where: Dict[str, Any] = {} + try: + where.update(_get_condition_to_filter_out_ui_session_tokens()) + except NameError: + # Helper may not exist in some builds; ignore if missing + pass + + rows = await prisma_client.db.litellm_verificationtoken.find_many( + where=where, + order=[{"key_alias": "asc"}], + ) + + seen = set() + aliases: List[str] = [] + for row in rows: + alias = getattr(row, "key_alias", None) + if alias is None and isinstance(row, dict): + alias = row.get("key_alias") + + if not alias: + continue + + alias_str = str(alias).strip() + if alias_str and alias_str not in seen: + seen.add(alias_str) + aliases.append(alias_str) + + verbose_proxy_logger.debug(f"Returning {len(aliases)} key aliases") + + return {"aliases": aliases} + + except Exception as e: + verbose_proxy_logger.exception(f"Error in key_aliases: {e}") + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"error({str(e)})"), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +def _validate_sort_params( + sort_by: Optional[str], sort_order: str +) -> Optional[Dict[str, str]]: + order_by: Dict[str, str] = {} + + if sort_by is None: + return None + # Validate sort_by is a valid column + valid_columns = [ + "spend", + "max_budget", + "created_at", + "updated_at", + "token", + "key_alias", + ] + if sort_by not in valid_columns: + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}" + }, + ) + + # Validate sort_order + if sort_order.lower() not in ["asc", "desc"]: + raise HTTPException( + status_code=400, + detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"}, + ) + + order_by[sort_by] = sort_order.lower() + + return order_by + + +def _build_key_filter_conditions( + user_id: Optional[str], + team_id: Optional[str], + organization_id: Optional[str], + key_alias: Optional[str], + key_hash: Optional[str], + exclude_team_id: Optional[str], + admin_team_ids: Optional[List[str]], + include_created_by_keys: bool, +) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: + """Build filter conditions for key listing.""" + # Prepare filter conditions + where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {} + where.update(_get_condition_to_filter_out_ui_session_tokens()) + + # Build the OR conditions for user's keys and admin team keys + or_conditions: List[Dict[str, Any]] = [] + + # Base conditions for user's own keys + user_condition: Dict[str, Any] = {} + if user_id and isinstance(user_id, str): + user_condition["user_id"] = user_id + if team_id and isinstance(team_id, str): + user_condition["team_id"] = team_id + if key_alias and isinstance(key_alias, str): + user_condition["key_alias"] = key_alias + if exclude_team_id and isinstance(exclude_team_id, str): + user_condition["team_id"] = {"not": exclude_team_id} + if organization_id and isinstance(organization_id, str): + user_condition["organization_id"] = organization_id + if key_hash and isinstance(key_hash, str): + user_condition["token"] = key_hash + + if user_condition: + or_conditions.append(user_condition) + + # Add condition for created by keys if provided + if include_created_by_keys and user_id: + or_conditions.append({"created_by": user_id}) + + # Add condition for admin team keys if provided + if admin_team_ids: + or_conditions.append({"team_id": {"in": admin_team_ids}}) + + # Combine conditions with OR if we have multiple conditions + if len(or_conditions) > 1: + where = {"AND": [where, {"OR": or_conditions}]} + elif len(or_conditions) == 1: + where.update(or_conditions[0]) + + verbose_proxy_logger.debug(f"Filter conditions: {where}") + return where + + +async def _list_key_helper( + prisma_client: PrismaClient, + page: int, + size: int, + user_id: Optional[str], + team_id: Optional[str], + organization_id: Optional[str], + key_alias: Optional[str], + key_hash: Optional[str], + exclude_team_id: Optional[str] = None, + return_full_object: bool = False, + admin_team_ids: Optional[ + List[str] + ] = None, # New parameter for teams where user is admin + include_created_by_keys: bool = False, + sort_by: Optional[str] = None, + sort_order: str = "desc", + expand: Optional[List[str]] = None, + status: Optional[str] = None, +) -> KeyListResponseObject: + """ + Helper function to list keys + Args: + page: int + size: int + user_id: Optional[str] + team_id: Optional[str] + key_alias: Optional[str] + exclude_team_id: Optional[str] # exclude a specific team_id + return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token + admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin + + Returns: + KeyListResponseObject + { + "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types + "total_count": int, + "current_page": int, + "total_pages": int, + } + """ + where = _build_key_filter_conditions( + user_id=user_id, + team_id=team_id, + organization_id=organization_id, + key_alias=key_alias, + key_hash=key_hash, + exclude_team_id=exclude_team_id, + admin_team_ids=admin_team_ids, + include_created_by_keys=include_created_by_keys, + ) + + # Calculate skip for pagination + skip = (page - 1) * size + + verbose_proxy_logger.debug(f"Pagination: skip={skip}, take={size}") + + order_by: Optional[Dict[str, str]] = ( + _validate_sort_params(sort_by, sort_order) + if sort_by is not None and isinstance(sort_by, str) + else None + ) + + # Determine which table to query based on status + use_deleted_table = status == "deleted" + + # Fetch keys with pagination + if use_deleted_table: + keys = await prisma_client.db.litellm_deletedverificationtoken.find_many( + where=where, # type: ignore + skip=skip, # type: ignore + take=size, # type: ignore + order=( + order_by + if order_by + else [ + {"created_at": "desc"}, + {"token": "desc"}, # fallback sort + ] + ), + ) + else: + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where=where, # type: ignore + skip=skip, # type: ignore + take=size, # type: ignore + order=( + order_by + if order_by + else [ + {"created_at": "desc"}, + {"token": "desc"}, # fallback sort + ] + ), + include={"object_permission": True}, + ) + + verbose_proxy_logger.debug(f"Fetched {len(keys)} keys") + + # Get total count of keys + if use_deleted_table: + total_count = await prisma_client.db.litellm_deletedverificationtoken.count( + where=where # type: ignore + ) + else: + total_count = await prisma_client.db.litellm_verificationtoken.count( + where=where # type: ignore + ) + + verbose_proxy_logger.debug(f"Total count of keys: {total_count}") + + # Calculate total pages + total_pages = -(-total_count // size) # Ceiling division + + # Fetch user information if expand includes "user" + user_map = {} + if expand and "user" in expand: + user_ids = [key.user_id for key in keys if key.user_id] + if user_ids: + users = await prisma_client.db.litellm_usertable.find_many( + where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates + ) + user_map = {user.user_id: user for user in users} + + # Prepare response + key_list: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] = [] + for key in keys: + # Convert Prisma model to dict (supports both Pydantic v1 and v2) + try: + key_dict = key.model_dump() + except Exception: + # Fallback for Pydantic v1 compatibility + key_dict = key.dict() + # Attach object_permission if object_permission_id is set (only for non-deleted keys) + if not use_deleted_table: + key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) + + # Include user information if expand includes "user" + if expand and "user" in expand and key.user_id and key.user_id in user_map: + try: + key_dict["user"] = user_map[key.user_id].model_dump() + except Exception: + key_dict["user"] = user_map[key.user_id].dict() + + if return_full_object is True or (expand and "user" in expand): + if use_deleted_table: + # Use deleted key type to preserve deleted_at, deleted_by, etc. + key_list.append(LiteLLM_DeletedVerificationToken(**key_dict)) + else: + key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object + else: + _token = key_dict.get("token") + key_list.append(cast(str, _token)) # Return only the token + + return KeyListResponseObject( + keys=key_list, + total_count=total_count, + current_page=page, + total_pages=total_pages, + ) + + +def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: + """ + Condition to filter out UI session tokens + """ + return { + "OR": [ + {"team_id": None}, # Include records where team_id is null + { + "team_id": {"not": UI_SESSION_TOKEN_TEAM_ID} + }, # Include records where team_id != UI_SESSION_TOKEN_TEAM_ID + ] + } + + +@router.post( + "/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)] +) +@management_endpoint_wrapper +async def block_key( + data: BlockKeyRequest, + http_request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +) -> Optional[LiteLLM_VerificationToken]: + """ + Block an Virtual key from making any requests. + + Parameters: + - key: str - The key to block. Can be either the unhashed key (sk-...) or the hashed key value + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/key/block' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + }' + ``` + + Note: This is an admin-only endpoint. Only proxy admins can block keys. + """ + from litellm.proxy.proxy_server import ( + create_audit_log_for_update, + hash_token, + litellm_proxy_admin_name, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) + + if not is_valid_api_key(data.key): + raise ProxyException( + message="Invalid key format.", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=status.HTTP_400_BAD_REQUEST, + ) + if data.key.startswith("sk-"): + hashed_token = hash_token(token=data.key) + else: + hashed_token = data.key + + if litellm.store_audit_logs is True: + # make an audit log for key update + record = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_token} + ) + if record is None: + raise ProxyException( + message=f"Key {data.key} not found", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=status.HTTP_404_NOT_FOUND, + ) + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=hashed_token, + action="blocked", + updated_values="{}", + before_value=record.model_dump_json(), + ) + ) + ) + + record = await prisma_client.db.litellm_verificationtoken.update( + where={"token": hashed_token}, data={"blocked": True} # type: ignore + ) + + ## UPDATE KEY CACHE + + ### get cached object ### + key_object = await get_key_object( + hashed_token=hashed_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + ### update cached object ### + key_object.blocked = True + + ### store cached object ### + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=key_object, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return record + + +@router.post( + "/key/unblock", tags=["key management"], dependencies=[Depends(user_api_key_auth)] +) +@management_endpoint_wrapper +async def unblock_key( + data: BlockKeyRequest, + http_request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +): + """ + Unblock a Virtual key to allow it to make requests again. + + Parameters: + - key: str - The key to unblock. Can be either the unhashed key (sk-...) or the hashed key value + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/key/unblock' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "key": "sk-Fn8Ej39NxjAXrvpUGKghGw" + }' + ``` + + Note: This is an admin-only endpoint. Only proxy admins can unblock keys. + """ + from litellm.proxy.proxy_server import ( + create_audit_log_for_update, + hash_token, + litellm_proxy_admin_name, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) + + if not is_valid_api_key(data.key): + raise ProxyException( + message="Invalid key format.", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=status.HTTP_400_BAD_REQUEST, + ) + if data.key.startswith("sk-"): + hashed_token = hash_token(token=data.key) + else: + hashed_token = data.key + + if litellm.store_audit_logs is True: + # make an audit log for key update + record = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_token} + ) + if record is None: + raise ProxyException( + message=f"Key {data.key} not found", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=status.HTTP_404_NOT_FOUND, + ) + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=hashed_token, + action="blocked", + updated_values="{}", + before_value=record.model_dump_json(), + ) + ) + ) + + record = await prisma_client.db.litellm_verificationtoken.update( + where={"token": hashed_token}, data={"blocked": False} # type: ignore + ) + + ## UPDATE KEY CACHE + + ### get cached object ### + key_object = await get_key_object( + hashed_token=hashed_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + ### update cached object ### + key_object.blocked = False + + ### store cached object ### + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=key_object, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return record + + +@router.post( + "/key/health", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], + response_model=KeyHealthResponse, +) +@management_endpoint_wrapper +async def key_health( + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Check the health of the key + + Checks: + - If key based logging is configured correctly - sends a test log + + Usage + + Pass the key in the request header + + ```bash + curl -X POST "http://localhost:4000/key/health" \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" + ``` + + Response when logging callbacks are setup correctly: + + ```json + { + "key": "healthy", + "logging_callbacks": { + "callbacks": [ + "gcs_bucket" + ], + "status": "healthy", + "details": "No logger exceptions triggered, system is healthy. Manually check if logs were sent to ['gcs_bucket']" + } + } + ``` + + + Response when logging callbacks are not setup correctly: + ```json + { + "key": "unhealthy", + "logging_callbacks": { + "callbacks": [ + "gcs_bucket" + ], + "status": "unhealthy", + "details": "Logger exceptions triggered, system is unhealthy: Failed to load vertex credentials. Check to see if credentials containing partial/invalid information." + } + } + ``` + """ + try: + # Get the key's metadata + key_metadata = user_api_key_dict.metadata + + health_status: KeyHealthResponse = KeyHealthResponse( + key="healthy", + logging_callbacks=None, + ) + + # Check if logging is configured in metadata + if key_metadata and "logging" in key_metadata: + logging_statuses = await test_key_logging( + user_api_key_dict=user_api_key_dict, + request=request, + key_logging=key_metadata["logging"], + ) + health_status["logging_callbacks"] = logging_statuses + + # Check if any logging callback is unhealthy + if logging_statuses.get("status") == "unhealthy": + health_status["key"] = "unhealthy" + + return KeyHealthResponse(**health_status) + + except Exception as e: + raise ProxyException( + message=f"Key health check failed: {str(e)}", + type=ProxyErrorTypes.internal_server_error, + param=getattr(e, "param", "None"), + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +async def _can_user_query_key_info( + user_api_key_dict: UserAPIKeyAuth, + key: Optional[str], + key_info: LiteLLM_VerificationToken, +) -> bool: + """ + Helper to check if the user has access to the key's info + """ + if ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value + ): + return True + elif user_api_key_dict.api_key == key: + return True + # user can query their own key info + elif key_info.user_id == user_api_key_dict.user_id: + return True + elif await TeamMemberPermissionChecks.user_belongs_to_keys_team( + user_api_key_dict=user_api_key_dict, + existing_key_row=key_info, + ): + return True + return False + + +async def test_key_logging( + user_api_key_dict: UserAPIKeyAuth, + request: Request, + key_logging: List[Dict[str, Any]], +) -> LoggingCallbackStatus: + """ + Test the key-based logging + + - Test that key logging is correctly formatted and all args are passed correctly + - Make a mock completion call -> user can check if it's correctly logged + - Check if any logger.exceptions were triggered -> if they were then returns it to the user client side + """ + import logging + from io import StringIO + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import general_settings, proxy_config + + logging_callbacks: List[str] = [] + for callback in key_logging: + if callback.get("callback_name") is not None: + logging_callbacks.append(callback["callback_name"]) + else: + raise ValueError("callback_name is required in key_logging") + + log_capture_string = StringIO() + ch = logging.StreamHandler(log_capture_string) + ch.setLevel(logging.ERROR) + logger = logging.getLogger() + logger.addHandler(ch) + + try: + data = { + "model": "openai/litellm-key-health-test", + "messages": [ + { + "role": "user", + "content": "Hello, this is a test from litellm /key/health. No LLM API call was made for this", + } + ], + "mock_response": "test response", + } + data = await add_litellm_data_to_request( + data=data, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + general_settings=general_settings, + request=request, + ) + await litellm.acompletion( + **data + ) # make mock completion call to trigger key based callbacks + except Exception as e: + return LoggingCallbackStatus( + callbacks=logging_callbacks, + status="unhealthy", + details=f"Logging test failed: {str(e)}", + ) + + await asyncio.sleep( + 2 + ) # wait for callbacks to run, callbacks use batching so wait for the flush event + + # Check if any logger exceptions were triggered + log_contents = log_capture_string.getvalue() + logger.removeHandler(ch) + if log_contents: + return LoggingCallbackStatus( + callbacks=logging_callbacks, + status="unhealthy", + details=f"Logger exceptions triggered, system is unhealthy: {log_contents}", + ) + else: + return LoggingCallbackStatus( + callbacks=logging_callbacks, + status="healthy", + details=f"No logger exceptions triggered, system is healthy. Manually check if logs were sent to {logging_callbacks} ", + ) + + +async def _enforce_unique_key_alias( + key_alias: Optional[str], + prisma_client: Any, + existing_key_token: Optional[str] = None, +) -> None: + """ + Helper to enforce unique key aliases across all keys. + + Args: + key_alias (Optional[str]): The key alias to check + prisma_client (Any): Prisma client instance + existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check + (The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating) + + Raises: + ProxyException: If key alias already exists on a different key + """ + if key_alias is not None and prisma_client is not None: + where_clause: dict[str, Any] = {"key_alias": key_alias} + if existing_key_token: + # Exclude the current key from the uniqueness check + where_clause["NOT"] = {"token": existing_key_token} + + existing_key = await prisma_client.db.litellm_verificationtoken.find_first( + where=where_clause + ) + if existing_key is not None: + raise ProxyException( + message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", + type=ProxyErrorTypes.bad_request_error, + param="key_alias", + code=status.HTTP_400_BAD_REQUEST, + ) + + +def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: + """ + Validate the model_max_budget is GenericBudgetConfigType + enforce user has an enterprise license + + Raises: + Exception: If model_max_budget is not a valid GenericBudgetConfigType + """ + try: + if model_max_budget is None: + return + if len(model_max_budget) == 0: + return + if model_max_budget is not None: + from litellm.proxy.proxy_server import CommonProxyErrors, premium_user + + if premium_user is not True: + raise ValueError( + f"You must have an enterprise license to set model_max_budget. {CommonProxyErrors.not_premium_user.value}" + ) + for _model, _budget_info in model_max_budget.items(): + assert isinstance(_model, str) + + # Normalize to dict (Pydantic may already parse nested values as BudgetConfig) + _info = ( + _budget_info.model_dump() + if hasattr(_budget_info, "model_dump") + else dict(_budget_info) + ) + # /CRUD endpoints can pass budget_limit as a string, so we need to convert it to a float + if "budget_limit" in _info: + _info["budget_limit"] = float(_info["budget_limit"]) + BudgetConfig(**_info) + except Exception as e: + raise ValueError( + f"Invalid model_max_budget: {str(e)}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users" + ) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 64f1ec24234..54b9e31a6da 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -22,7 +22,6 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, add_litellm_data_to_request, ) -from litellm.types.utils import SupportedCacheControls @pytest.fixture @@ -496,9 +495,7 @@ def test_add_litellm_data_for_backend_llm_call( from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - user_api_key_dict = UserAPIKeyAuth( - api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" - ) + UserAPIKeyAuth(api_key="test_api_key", user_id="test_user_id", org_id="test_org_id") data = LiteLLMProxyRequestSetup.get_user_from_headers( headers=headers, @@ -1059,7 +1056,7 @@ def test_update_config_fields_default_internal_user_params(monkeypatch): }, }, } - updated_config = proxy_config._update_config_fields(**args) + proxy_config._update_config_fields(**args) assert litellm.default_internal_user_params == { "user_role": "proxy_admin", @@ -1320,6 +1317,61 @@ def test_litellm_verification_token_view_response_with_budget_table( ) +def test_litellm_verification_token_view_budget_does_not_override_key_model_max_budget(): + """ + When key has non-empty model_max_budget, budget's model_max_budget is NOT applied. + Regression test for per-model budget: only apply budget's model_max_budget when key's is empty. + """ + from litellm.proxy._types import LiteLLM_VerificationTokenView + + key_model_max_budget = {"gpt-4": {"max_budget": 50.0, "budget_duration": "1d"}} + args = { + "token": "sk-test-mock-token-303", + "key_name": "sk-...if_g", + "key_alias": None, + "soft_budget_cooldown": False, + "spend": 0.0, + "expires": None, + "models": [], + "aliases": {}, + "config": {}, + "user_id": None, + "team_id": "test", + "permissions": {}, + "max_parallel_requests": None, + "metadata": {}, + "blocked": None, + "tpm_limit": None, + "rpm_limit": None, + "max_budget": None, + "budget_duration": None, + "budget_reset_at": None, + "allowed_cache_controls": [], + "model_spend": {}, + "model_max_budget": key_model_max_budget, + "budget_id": "my-test-tier", + "created_at": "2024-12-26T02:28:52.615+00:00", + "updated_at": "2024-12-26T03:01:51.159+00:00", + "team_spend": None, + "team_max_budget": None, + "team_tpm_limit": None, + "team_rpm_limit": None, + "team_models": [], + "team_metadata": {}, + "team_blocked": False, + "team_alias": None, + "team_members_with_roles": [], + "team_member_spend": None, + "team_model_aliases": None, + "team_member": None, + "litellm_budget_table_model_max_budget": { + "gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"} + }, + } + resp = LiteLLM_VerificationTokenView(**args) + assert resp.model_max_budget == key_model_max_budget + + def test_is_allowed_to_make_key_request(): from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -1381,13 +1433,6 @@ def test_get_model_group_info(): assert len(model_list) == 1 -import asyncio -import json -from unittest.mock import AsyncMock, patch - -import pytest - - @pytest.fixture def mock_team_data(): return [ @@ -1444,7 +1489,6 @@ async def test_get_user_info_for_proxy_admin(mock_team_data, mock_key_data): "litellm.proxy.proxy_server.prisma_client", MockPrismaClientDB(mock_team_data, mock_key_data), ): - from litellm.proxy.management_endpoints.internal_user_endpoints import ( _get_user_info_for_proxy_admin, ) @@ -1558,9 +1602,6 @@ def test_update_key_budget_with_temp_budget_increase(): assert _update_key_budget_with_temp_budget_increase(valid_token).max_budget == 200 -from unittest.mock import AsyncMock, MagicMock - - @pytest.mark.asyncio async def test_health_check_not_called_when_disabled(monkeypatch): from litellm.proxy.proxy_server import ProxyStartupEvent @@ -1603,18 +1644,12 @@ async def test_health_check_not_called_when_disabled(monkeypatch): }, ) def test_custom_openapi(mock_get_openapi_schema): - from litellm.proxy.proxy_server import app, custom_openapi + from litellm.proxy.proxy_server import custom_openapi openapi_schema = custom_openapi() assert openapi_schema is not None -import asyncio -from datetime import timedelta -from unittest.mock import AsyncMock, MagicMock - -import pytest - from litellm.proxy.utils import ProxyUpdateSpend @@ -1639,6 +1674,7 @@ async def test_end_user_transactions_reset(): async def test_spend_logs_cleanup_after_error(): # Setup test data import asyncio + mock_client = MagicMock() mock_client.spend_log_transactions = [ {"id": 1, "amount": 10.0}, @@ -1826,7 +1862,7 @@ def test_provider_specific_header_in_request(custom_llm_provider, expected_resul client = HTTPHandler() with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: - resp = litellm.completion( + litellm.completion( model="anthropic/claude-3-5-sonnet-v2@20241022", messages=[{"role": "user", "content": "Hello world"}], provider_specific_header=ProviderSpecificHeader( @@ -2063,7 +2099,7 @@ async def test_post_call_failure_hook_auth_error_key_info_route(): Test that post_call_failure_hook does NOT call _handle_logging_proxy_only_error when we get an auth error from /key/info route (since it's not an LLM API route). """ - from unittest.mock import AsyncMock, Mock, patch + from unittest.mock import AsyncMock, patch from fastapi import HTTPException @@ -2117,7 +2153,7 @@ async def test_post_call_failure_hook_auth_error_llm_api_route(): Test that post_call_failure_hook DOES call _handle_logging_proxy_only_error when we get an auth error from /v1/chat/completions route (since it is an LLM API route). """ - from unittest.mock import AsyncMock, Mock, patch + from unittest.mock import AsyncMock, patch from fastapi import HTTPException @@ -2182,27 +2218,27 @@ async def test_during_call_hook_parallel_execution(): cache = DualCache() proxy_logging = ProxyLogging(user_api_key_cache=cache) execution_order = [] - + class TestGuardrail(CustomGuardrail): def __init__(self, name): super().__init__( guardrail_name=name, event_hook=GuardrailEventHooks.during_call, - default_on=True + default_on=True, ) self.name = name - + async def async_moderation_hook(self, data, user_api_key_dict, call_type): execution_order.append(f"{self.name}_start") await asyncio.sleep(0.1) execution_order.append(f"{self.name}_end") return data - + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] - + try: litellm.callbacks = [TestGuardrail(f"g{i}") for i in range(3)] - + start_time = asyncio.get_event_loop().time() result = await proxy_logging.during_call_hook( data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]}, @@ -2210,14 +2246,22 @@ async def test_during_call_hook_parallel_execution(): call_type="completion", ) execution_time = asyncio.get_event_loop().time() - start_time - + # Verify parallel execution: all start before any end - first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) - starts_before_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) - assert starts_before_end == 3, f"Expected 3 starts before first end, got {starts_before_end}" - + first_end_idx = next( + i for i, item in enumerate(execution_order) if "end" in item + ) + starts_before_end = sum( + 1 for item in execution_order[:first_end_idx] if "start" in item + ) + assert ( + starts_before_end == 3 + ), f"Expected 3 starts before first end, got {starts_before_end}" + # Verify timing: parallel ~0.1s vs sequential ~0.3s - assert execution_time < 0.2, f"Parallel execution took {execution_time}s, expected < 0.2s" + assert ( + execution_time < 0.2 + ), f"Parallel execution took {execution_time}s, expected < 0.2s" assert result["model"] == "gpt-4" finally: litellm.callbacks = original_callbacks @@ -2235,30 +2279,35 @@ async def test_during_call_hook_parallel_execution_with_error(): cache = DualCache() proxy_logging = ProxyLogging(user_api_key_cache=cache) - + class FailingGuardrail(CustomGuardrail): def __init__(self): super().__init__( guardrail_name="failing_guardrail", event_hook=GuardrailEventHooks.during_call, - default_on=True + default_on=True, ) - + async def async_moderation_hook(self, data, user_api_key_dict, call_type): raise ValueError("Guardrail violation detected!") - + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] - + try: litellm.callbacks = [FailingGuardrail()] - + with pytest.raises(ValueError) as exc_info: await proxy_logging.during_call_hook( - data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]}, - user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + data={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "test"}], + }, + user_api_key_dict=UserAPIKeyAuth( + api_key="test_key", user_id="test_user" + ), call_type="completion", ) - + assert "Guardrail violation detected!" in str(exc_info.value) finally: - litellm.callbacks = original_callbacks \ No newline at end of file + litellm.callbacks = original_callbacks diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index fc8373a1746..352db384c88 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -1,30 +1,20 @@ -import json import os import sys -from datetime import datetime -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system-path -from datetime import datetime as dt_object -import time -import pytest -import litellm -import json -from litellm.types.utils import BudgetConfig as GenericBudgetInfo -import os -import sys -from datetime import datetime -from unittest.mock import AsyncMock, patch import pytest + +import litellm from litellm.caching.caching import DualCache from litellm.proxy.hooks.model_max_budget_limiter import ( _PROXY_VirtualKeyModelMaxBudgetLimiter, ) from litellm.proxy._types import UserAPIKeyAuth -import litellm +from litellm.types.utils import BudgetConfig as GenericBudgetInfo # Test class setup @@ -123,3 +113,48 @@ async def test_get_virtual_key_spend_for_model(budget_limiter): key_budget_config=budget_config, ) assert spend == 50.0 + + +@pytest.mark.asyncio +async def test_async_log_success_event_uses_per_model_budget_duration(budget_limiter): + """ + async_log_success_event must use the per-model budget_duration for the cache key + so spend is tracked per model correctly. Regression test for per-model budget implementation. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + ) + + virtual_key = "test-key-hash" + model = "gpt-4" + budget_duration = "1d" + user_api_key_model_max_budget = { + model: {"budget_limit": 100.0, "time_period": budget_duration}, + } + kwargs = { + "standard_logging_object": { + "response_cost": 0.05, + "model": model, + "metadata": {"user_api_key_hash": virtual_key}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": user_api_key_model_max_budget + }, + }, + } + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + call_kwargs = mock_increment.call_args.kwargs + spend_key = call_kwargs["spend_key"] + assert spend_key == ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" + ) + assert call_kwargs["response_cost"] == 0.05 diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index d5c3ecae7d6..d8c505223d9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -11,7 +11,6 @@ import litellm.proxy.proxy_server as ps from litellm.proxy.proxy_server import app from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, CommonProxyErrors -import litellm.proxy.management_endpoints.budget_management_endpoints as bm sys.path.insert( 0, os.path.abspath("../../../") @@ -22,13 +21,12 @@ sys.path.insert( def client_and_mocks(monkeypatch): # Setup MagicMock Prisma mock_prisma = MagicMock() - mock_table = MagicMock() mock_table.create = AsyncMock(side_effect=lambda *, data: data) mock_table.update = AsyncMock(side_effect=lambda *, where, data: {**where, **data}) mock_prisma.db = types.SimpleNamespace( - litellm_budgettable = mock_table, - litellm_dailyspend = mock_table, + litellm_budgettable=mock_table, + litellm_dailyspend=mock_table, ) # Monkeypatch Mocked Prisma client into the server module @@ -79,6 +77,7 @@ async def test_new_budget_db_not_connected(client_and_mocks, monkeypatch): # override the prisma_client that the handler imports at runtime import litellm.proxy.proxy_server as ps + monkeypatch.setattr(ps, "prisma_client", None) # Call /budget/new endpoint @@ -123,6 +122,7 @@ async def test_update_budget_db_not_connected(client_and_mocks, monkeypatch): # override the prisma_client that the handler imports at runtime import litellm.proxy.proxy_server as ps + monkeypatch.setattr(ps, "prisma_client", None) payload = {"budget_id": "any", "max_budget": 1.0} @@ -136,7 +136,7 @@ async def test_update_budget_db_not_connected(client_and_mocks, monkeypatch): async def test_update_budget_allows_null_max_budget(client_and_mocks): """ Test that /budget/update allows setting max_budget to null. - + Previously, using exclude_none=True would drop null values, making it impossible to remove a budget limit. With exclude_unset=True, explicitly setting max_budget to null should include it in the update. @@ -144,11 +144,11 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks): client, _, mock_table = client_and_mocks captured_data = {} - + async def capture_update(*, where, data): captured_data.update(data) return {**where, **data} - + mock_table.update = AsyncMock(side_effect=capture_update) payload = { @@ -159,9 +159,11 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks): assert resp.status_code == 200, resp.text # Verify that max_budget=None was included in the update data - assert "max_budget" in captured_data, "max_budget should be included when explicitly set to null" + assert ( + "max_budget" in captured_data + ), "max_budget should be included when explicitly set to null" assert captured_data["max_budget"] is None, "max_budget should be None" - + mock_table.update.assert_awaited_once() @@ -169,7 +171,7 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks): async def test_new_budget_negative_max_budget(client_and_mocks): """ Test that /budget/new rejects negative max_budget values. - + This prevents the issue where negative budgets would always trigger budget exceeded errors. """ @@ -181,7 +183,7 @@ async def test_new_budget_negative_max_budget(client_and_mocks): } resp = client.post("/budget/new", json=payload) assert resp.status_code == 400, resp.text - + detail = resp.json()["detail"] assert "max_budget cannot be negative" in str(detail) @@ -199,7 +201,7 @@ async def test_new_budget_negative_soft_budget(client_and_mocks): } resp = client.post("/budget/new", json=payload) assert resp.status_code == 400, resp.text - + detail = resp.json()["detail"] assert "soft_budget cannot be negative" in str(detail) @@ -217,7 +219,7 @@ async def test_update_budget_negative_max_budget(client_and_mocks): } resp = client.post("/budget/update", json=payload) assert resp.status_code == 400, resp.text - + detail = resp.json()["detail"] assert "max_budget cannot be negative" in str(detail) @@ -235,6 +237,30 @@ async def test_update_budget_negative_soft_budget(client_and_mocks): } resp = client.post("/budget/update", json=payload) assert resp.status_code == 400, resp.text - + detail = resp.json()["detail"] assert "soft_budget cannot be negative" in str(detail) + + +@pytest.mark.asyncio +async def test_new_budget_invalid_model_max_budget(client_and_mocks, monkeypatch): + """ + Test that /budget/new validates model_max_budget and returns 400 for invalid structure. + Per-model budget implementation: validate_model_max_budget is called in new_budget. + """ + import litellm.proxy.proxy_server as ps + + monkeypatch.setattr(ps, "premium_user", True) + + client, _, _ = client_and_mocks + + payload = { + "budget_id": "budget_invalid_mmb", + "max_budget": 10.0, + "model_max_budget": {"gpt-4": "not-a-dict"}, + } + resp = client.post("/budget/new", json=payload) + # Pydantic may reject invalid structure with 422 before our validator runs + assert resp.status_code in (400, 422), resp.text + detail = resp.json()["detail"] + assert "model_max_budget" in str(detail) or "dictionary" in str(detail).lower()