litellm-mirror/litellm/proxy/example_config_yaml/custom_auth.py
Krish Dholakia d57be47b0f
Litellm ruff linting enforcement (#5992)
* ci(config.yml): add a 'check_code_quality' step

Addresses https://github.com/BerriAI/litellm/issues/5991

* ci(config.yml): check why circle ci doesn't pick up this test

* ci(config.yml): fix to run 'check_code_quality' tests

* fix(__init__.py): fix unprotected import

* fix(__init__.py): don't remove unused imports

* build(ruff.toml): update ruff.toml to ignore unused imports

* fix: fix: ruff + pyright - fix linting + type-checking errors

* fix: fix linting errors

* fix(lago.py): fix module init error

* fix: fix linting errors

* ci(config.yml): cd into correct dir for checks

* fix(proxy_server.py): fix linting error

* fix(utils.py): fix bare except

causes ruff linting errors

* fix: ruff - fix remaining linting errors

* fix(clickhouse.py): use standard logging object

* fix(__init__.py): fix unprotected import

* fix: ruff - fix linting errors

* fix: fix linting errors

* ci(config.yml): cleanup code qa step (formatting handled in local_testing)

* fix(_health_endpoints.py): fix ruff linting errors

* ci(config.yml): just use ruff in check_code_quality pipeline for now

* build(custom_guardrail.py): include missing file

* style(embedding_handler.py): fix ruff check
2024-10-01 19:44:20 -04:00

52 lines
1.4 KiB
Python

import os
from fastapi import Request
from litellm.proxy._types import GenerateKeyRequest, UserAPIKeyAuth
async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
try:
modified_master_key = f"{os.getenv('PROXY_MASTER_KEY')}-1234"
if api_key == modified_master_key:
return UserAPIKeyAuth(api_key=api_key)
raise Exception
except Exception:
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")
data_json.get("duration")
data_json.get("models")
data_json.get("aliases")
data_json.get("config")
data_json.get("spend")
data_json.get("user_id")
data_json.get("max_parallel_requests")
data_json.get("metadata")
data_json.get("tpm_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.",
}