[Feature]: Override user_api_key_auth with custom auth checking

[krrishdholakia]

The Feature

allow a user to specify their own custom auth and override user api key auth

Motivation, pitch

user request

Twitter / LinkedIn details

No response

[krrishdholakia]

cc: @jeromeroussin

[krrishdholakia]

Did some investigation, here's how we could enable this:

1. Create a custom auth handler file

# custom_auth.py
def handle_request(request):
    # User-defined rules logic
    # Return True if the request is accepted, False otherwise
    return True

2. Point file to proxy in config.yaml

custom_rules_file: path/to/custom_rules.py

3. Load the custom rules

import yaml
import importlib.util

def load_config():
    with open('config.yaml', 'r') as config_file:
        config = yaml.safe_load(config_file)
    return config

def load_custom_rules_handler(file_path):
    spec = importlib.util.spec_from_file_location("custom_rules", file_path)
    custom_rules = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(custom_rules)
    return custom_rules.handle_request

def process_request(request):
    config = load_config()
    custom_rules_handler = load_custom_rules_handler(config['custom_rules_file'])

    # Check request against custom rules
    if custom_rules_handler(request):
        # Accept the request
        print("Request accepted")
    else:
        # Reject the request
        print("Request rejected")

# Example usage
request_example = {'some_data': 'example'}
process_request(request_example)

[krrishdholakia]

cc: @jeromeroussin

is the idea to accept/reject a request like how auth works today, or just run a series of tasks in parallel

(i ask to understand if this blocks execution).

[krrishdholakia]

given the introduction of spend tracking per key, i think the proxy will need some key/db interface to interact with, to update spend per call.

But i'm open to your suggestions on a good implementation for you.

[krrishdholakia]

initial commit pushed - 030bd22

@jeromeroussin you can now override the default api key auth.

Here's how:

1. Create a custom auth file.

Make sure the response type follows the UserAPIKeyAuth pydantic object.

from litellm.proxy.types import UserAPIKeyAuth

async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth: 
    try: 
        modified_master_key = "sk-my-master-key"
        if api_key == modified_master_key:
            return UserAPIKeyAuth(api_key=api_key)
        raise Exception
    except: 
        raise Exception

2. Pass the filepath (relative to the config.yaml) in the config.yaml

e.g. if they're both in the same dir, this is what it looks like:

model_list: 
  - model_name: "openai-model"
    litellm_params: 
      model: "gpt-3.5-turbo"

litellm_settings:
  drop_params: True
  set_verbose: True

general_settings:
  custom_auth: custom_auth.user_api_key_auth

3. Start the proxy

$ litellm --config /path/to/config.yaml 

[jeromeroussin]

This works but it rewrites the exceptions coming from the custom auth function. Can you perhaps leave exceptions untouched if they are HTTPException already?

litellm/litellm/proxy/proxy_server.py

Line 280 in 1463cc6

except Exception as e:

[krrishdholakia]

yes - just added 4d7ff1b

[krrishdholakia]

Will update ticket once it's in prod.

[krrishdholakia]

@jeromeroussin should be live in v1.10.11 Let me know if this works as expected.

[krrishdholakia]

Added docs - 37251d3