phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 02:34:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(proxy_server.py): expose new permissions field for keys

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-15 20:03:32 -08:00
parent 81e782c67e
commit cccd577e75
4 changed files with 11 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

4
.pre-commit-config.yaml
View file

@ -1,10 +1,10 @@
repos: repos:
- repo: https://github.com/psf/black - repo: https://github.com/psf/black
rev: stable rev: 24.2.0
hooks: hooks:
- id: black - id: black
- repo: https://github.com/pycqa/flake8 - repo: https://github.com/pycqa/flake8
rev: 3.8.4 # The version of flake8 to use rev: 7.0.0 # The version of flake8 to use
hooks: hooks:
- id: flake8 - id: flake8
exclude: ^litellm/tests/|^litellm/proxy/proxy_cli.py|^litellm/integrations/|^litellm/proxy/tests/ exclude: ^litellm/tests/|^litellm/proxy/proxy_cli.py|^litellm/integrations/|^litellm/proxy/tests/

4
litellm/proxy/_types.py
View file

@ -154,7 +154,7 @@ class GenerateKeyRequest(GenerateRequestBase):
duration: Optional[str] = None duration: Optional[str] = None
aliases: Optional[dict] = {} aliases: Optional[dict] = {}
config: Optional[dict] = {} config: Optional[dict] = {}
permissions: Optional[dict] = None permissions: Optional[dict] = {}
class GenerateKeyResponse(GenerateKeyRequest): class GenerateKeyResponse(GenerateKeyRequest):
@ -167,7 +167,7 @@ class GenerateKeyResponse(GenerateKeyRequest):
def set_model_info(cls, values): def set_model_info(cls, values):
if values.get("token") is not None: if values.get("token") is not None:
values.update({"key": values.get("token")}) values.update({"key": values.get("token")})
dict_fields = ["metadata", "aliases", "config"] dict_fields = ["metadata", "aliases", "config", "permissions"]
for field in dict_fields: for field in dict_fields:
value = values.get(field) value = values.get(field)
if value is not None and isinstance(value, str): if value is not None and isinstance(value, str):

14
litellm/proxy/proxy_server.py
View file

@ -1567,7 +1567,7 @@ async def generate_key_helper_fn(
update_key_values: Optional[dict] = None, update_key_values: Optional[dict] = None,
key_alias: Optional[str] = None, key_alias: Optional[str] = None,
allowed_cache_controls: Optional[list] = [], allowed_cache_controls: Optional[list] = [],
permissions: Optional[dict] = None, permissions: Optional[dict] = {},
): ):
global prisma_client, custom_db_client, user_api_key_cache global prisma_client, custom_db_client, user_api_key_cache
@ -1597,11 +1597,9 @@ async def generate_key_helper_fn(
duration_s = _duration_in_seconds(duration=budget_duration) duration_s = _duration_in_seconds(duration=budget_duration)
reset_at = datetime.utcnow() + timedelta(seconds=duration_s) reset_at = datetime.utcnow() + timedelta(seconds=duration_s)
if permissions is not None and isinstance(permissions, dict):
permissions = json.dumps(permissions) # type: ignore
aliases_json = json.dumps(aliases) aliases_json = json.dumps(aliases)
config_json = json.dumps(config) config_json = json.dumps(config)
permissions_json = json.dumps(permissions)
metadata_json = json.dumps(metadata) metadata_json = json.dumps(metadata)
user_id = user_id or str(uuid.uuid4()) user_id = user_id or str(uuid.uuid4())
user_role = user_role or "app_user" user_role = user_role or "app_user"
@ -1644,7 +1642,7 @@ async def generate_key_helper_fn(
"budget_duration": key_budget_duration, "budget_duration": key_budget_duration,
"budget_reset_at": key_reset_at, "budget_reset_at": key_reset_at,
"allowed_cache_controls": allowed_cache_controls, "allowed_cache_controls": allowed_cache_controls,
"permissions": permissions, "permissions": permissions_json,
} }
if ( if (
general_settings.get("allow_user_auth", False) == True general_settings.get("allow_user_auth", False) == True
@ -1824,9 +1822,9 @@ async def initialize(
user_api_base = api_base user_api_base = api_base
dynamic_config[user_model]["api_base"] = api_base dynamic_config[user_model]["api_base"] = api_base
if api_version: if api_version:
os.environ[ os.environ["AZURE_API_VERSION"] = (
"AZURE_API_VERSION" api_version # set this for azure - litellm can read this from the env
] = api_version # set this for azure - litellm can read this from the env )
if max_tokens: # model-specific param if max_tokens: # model-specific param
user_max_tokens = max_tokens user_max_tokens = max_tokens
dynamic_config[user_model]["max_tokens"] = max_tokens dynamic_config[user_model]["max_tokens"] = max_tokens

1
litellm/proxy/schema.prisma
View file

@ -55,6 +55,7 @@ model LiteLLM_VerificationToken {
config Json @default("{}") config Json @default("{}")
user_id String? user_id String?
team_id String? team_id String?
permissions Json @default("{}")
max_parallel_requests Int? max_parallel_requests Int?
metadata Json @default("{}") metadata Json @default("{}")
tpm_limit BigInt? tpm_limit BigInt?