diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 1cb28a2e3d4..c20bbef1127 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -392,6 +392,97 @@ general_settings: $ litellm --config /path/to/config.yaml ``` +## Custom /key/generate + +If you need to add custom logic before generating a Proxy API Key (Example Validating `team_id`) + +### 1. Write a custom `custom_generate_key_fn` + + +The input to the custom_generate_key_fn function is a single parameter: `data` [(Type: GenerateKeyRequest)](https://github.com/BerriAI/litellm/blob/main/litellm/proxy/_types.py#L125) + +The output of your `custom_generate_key_fn` should be a dictionary with the following structure +```python +{ + "decision": False, + "message": "This violates LiteLLM Proxy Rules. No team id provided.", +} + +``` + +- decision (Type: bool): A boolean value indicating whether the key generation is allowed (True) or not (False). + +- message (Type: str, Optional): An optional message providing additional information about the decision. This field is included when the decision is False. + + +```python +async def custom_generate_key_fn(data: GenerateKeyRequest)-> dict: + """ + Asynchronous function for generating a key based on the input data. + + Args: + data (GenerateKeyRequest): The input data for key generation. + + Returns: + dict: A dictionary containing the decision and an optional message. + { + "decision": False, + "message": "This violates LiteLLM Proxy Rules. No team id provided.", + } + """ + + # decide if a key should be generated or not + print("using custom auth function!") + data_json = data.json() # type: ignore + + # Unpacking variables + team_id = data_json.get("team_id") + duration = data_json.get("duration") + models = data_json.get("models") + aliases = data_json.get("aliases") + config = data_json.get("config") + spend = data_json.get("spend") + user_id = data_json.get("user_id") + max_parallel_requests = data_json.get("max_parallel_requests") + metadata = data_json.get("metadata") + tpm_limit = data_json.get("tpm_limit") + rpm_limit = data_json.get("rpm_limit") + + if team_id is not None and team_id == "litellm-core-infra@gmail.com": + # only team_id="litellm-core-infra@gmail.com" can make keys + return { + "decision": True, + } + else: + print("Failed custom auth") + return { + "decision": False, + "message": "This violates LiteLLM Proxy Rules. No team id provided.", + } +``` + + +### 2. Pass the filepath (relative to the config.yaml) + +Pass the filepath to the config.yaml + +e.g. if they're both in the same dir - `./config.yaml` and `./custom_auth.py`, this is what it looks like: +```yaml +model_list: + - model_name: "openai-model" + litellm_params: + model: "gpt-3.5-turbo" + +litellm_settings: + drop_params: True + set_verbose: True + +general_settings: + custom_key_generate: custom_auth.custom_generate_key_fn +``` + + + ## [BETA] Dynamo DB diff --git a/litellm/proxy/example_config_yaml/custom_auth.py b/litellm/proxy/example_config_yaml/custom_auth.py index 416b66682a9..a764a647a96 100644 --- a/litellm/proxy/example_config_yaml/custom_auth.py +++ b/litellm/proxy/example_config_yaml/custom_auth.py @@ -1,4 +1,4 @@ -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import UserAPIKeyAuth, GenerateKeyRequest from fastapi import Request from dotenv import load_dotenv import os @@ -14,3 +14,40 @@ async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth: raise Exception except: raise Exception + + +async def generate_key_fn(data: GenerateKeyRequest): + """ + Asynchronously decides if a key should be generated or not based on the provided data. + + Args: + data (GenerateKeyRequest): The data to be used for decision making. + + Returns: + bool: True if a key should be generated, False otherwise. + """ + # decide if a key should be generated or not + data_json = data.json() # type: ignore + + # Unpacking variables + team_id = data_json.get("team_id") + duration = data_json.get("duration") + models = data_json.get("models") + aliases = data_json.get("aliases") + config = data_json.get("config") + spend = data_json.get("spend") + user_id = data_json.get("user_id") + max_parallel_requests = data_json.get("max_parallel_requests") + metadata = data_json.get("metadata") + tpm_limit = data_json.get("tpm_limit") + rpm_limit = data_json.get("rpm_limit") + + if team_id is not None and len(team_id) > 0: + return { + "decision": True, + } + else: + return { + "decision": True, + "message": "This violates LiteLLM Proxy Rules. No team id provided.", + } diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 417b4c6f10d..29aa3cf4ff7 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -62,8 +62,9 @@ litellm_settings: # setting callback class # callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance] -# general_settings: -# master_key: sk-1234 +general_settings: + master_key: sk-1234 + custom_key_generate: custom_auth.generate_key_fn # database_type: "dynamo_db" # database_args: { # 👈 all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190 # "billing_mode": "PAY_PER_REQUEST", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cf75a250c5f..71f6a444703 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -187,6 +187,7 @@ prisma_client: Optional[PrismaClient] = None custom_db_client: Optional[DBClient] = None user_api_key_cache = DualCache() user_custom_auth = None +user_custom_key_generate = None use_background_health_checks = None use_queue = False health_check_interval = None @@ -874,7 +875,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, use_background_health_checks, health_check_interval, use_queue, custom_db_client + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client # Load existing config config = await self.get_config(config_file_path=config_file_path) @@ -1052,6 +1053,12 @@ class ProxyConfig: user_custom_auth = get_instance_fn( value=custom_auth, config_file_path=config_file_path ) + + custom_key_generate = general_settings.get("custom_key_generate", None) + if custom_key_generate is not None: + user_custom_key_generate = get_instance_fn( + value=custom_key_generate, config_file_path=config_file_path + ) ## dynamodb database_type = general_settings.get("database_type", None) if database_type is not None and ( @@ -2164,7 +2171,16 @@ async def generate_key_fn( - 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. """ + global user_custom_key_generate verbose_proxy_logger.debug("entered /key/generate") + + if user_custom_key_generate is not None: + result = await user_custom_key_generate(data) + 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) + data_json = data.json() # type: ignore response = await generate_key_helper_fn(**data_json) return GenerateKeyResponse( @@ -2948,7 +2964,7 @@ async def get_routes(): @router.on_event("shutdown") async def shutdown_event(): - global prisma_client, master_key, user_custom_auth + global prisma_client, master_key, user_custom_auth, user_custom_key_generate if prisma_client: verbose_proxy_logger.debug("Disconnecting from Prisma") await prisma_client.disconnect() @@ -2958,7 +2974,7 @@ async def shutdown_event(): def cleanup_router_config_variables(): - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, use_background_health_checks, health_check_interval, prisma_client, custom_db_client + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, prisma_client, custom_db_client # Set all variables to None master_key = None @@ -2966,6 +2982,7 @@ def cleanup_router_config_variables(): otel_logging = None user_custom_auth = None user_custom_auth_path = None + user_custom_key_generate = None use_background_health_checks = None health_check_interval = None prisma_client = None diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 004aef63a1c..5ecfc89d75c 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -33,6 +33,7 @@ import pytest, logging, asyncio import litellm, asyncio from litellm.proxy.proxy_server import ( new_user, + generate_key_fn, user_api_key_auth, user_update, delete_key_fn, @@ -49,6 +50,7 @@ from litellm.proxy._types import ( DynamoDBArgs, DeleteKeyRequest, UpdateKeyRequest, + GenerateKeyRequest, ) from litellm.proxy.utils import DBClient from starlette.datastructures import URL @@ -593,3 +595,87 @@ def test_generate_and_update_key(prisma_client): print("Got Exception", e) print(e.detail) pytest.fail(f"An exception occurred - {str(e)}") + + +def test_key_generate_with_custom_auth(prisma_client): + # custom - generate key function + async def custom_generate_key_fn(data: GenerateKeyRequest) -> dict: + """ + Asynchronous function for generating a key based on the input data. + + Args: + data (GenerateKeyRequest): The input data for key generation. + + Returns: + dict: A dictionary containing the decision and an optional message. + { + "decision": False, + "message": "This violates LiteLLM Proxy Rules. No team id provided.", + } + """ + + # decide if a key should be generated or not + print("using custom auth function!") + data_json = data.json() # type: ignore + + # Unpacking variables + team_id = data_json.get("team_id") + duration = data_json.get("duration") + models = data_json.get("models") + aliases = data_json.get("aliases") + config = data_json.get("config") + spend = data_json.get("spend") + user_id = data_json.get("user_id") + max_parallel_requests = data_json.get("max_parallel_requests") + metadata = data_json.get("metadata") + tpm_limit = data_json.get("tpm_limit") + rpm_limit = data_json.get("rpm_limit") + + if team_id is not None and team_id == "litellm-core-infra@gmail.com": + # only team_id="litellm-core-infra@gmail.com" can make keys + return { + "decision": True, + } + else: + print("Failed custom auth") + return { + "decision": False, + "message": "This violates LiteLLM Proxy Rules. No team id provided.", + } + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr( + litellm.proxy.proxy_server, "user_custom_key_generate", custom_generate_key_fn + ) + try: + + async def test(): + await litellm.proxy.proxy_server.prisma_client.connect() + try: + request = GenerateKeyRequest() + key = await generate_key_fn(request) + pytest.fail(f"Expected an exception. Got {key}") + except Exception as e: + # this should fail + print("Got Exception", e) + print(e.detail) + print("First request failed!. This is expected") + assert ( + "This violates LiteLLM Proxy Rules. No team id provided." + in e.detail + ) + + request_2 = GenerateKeyRequest( + team_id="litellm-core-infra@gmail.com", + ) + + key = await generate_key_fn(request_2) + print(key) + generated_key = key.key + + asyncio.run(test()) + except Exception as e: + print("Got Exception", e) + print(e.detail) + pytest.fail(f"An exception occurred - {str(e)}")