Merge pull request #1538 from BerriAI/litellm_use_custom_key_gen

[Feat] Proxy Auth - Use custom_key_generate

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 
 

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.",
+        }

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", 

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

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)}")