diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7a6e400c8..411b61a40 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -154,6 +154,7 @@ class GenerateKeyRequest(GenerateRequestBase): duration: Optional[str] = None aliases: Optional[dict] = {} config: Optional[dict] = {} + permissions: Optional[dict] = None class GenerateKeyResponse(GenerateKeyRequest): @@ -381,6 +382,7 @@ class LiteLLM_VerificationToken(LiteLLMBase): budget_duration: Optional[str] = None budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] + permissions: Optional[dict] = None class UserAPIKeyAuth( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 293185191..e3efb5e15 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1567,6 +1567,7 @@ async def generate_key_helper_fn( update_key_values: Optional[dict] = None, key_alias: Optional[str] = None, allowed_cache_controls: Optional[list] = [], + permissions: Optional[dict] = None, ): global prisma_client, custom_db_client, user_api_key_cache @@ -1596,6 +1597,9 @@ async def generate_key_helper_fn( duration_s = _duration_in_seconds(duration=budget_duration) 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) config_json = json.dumps(config) metadata_json = json.dumps(metadata) @@ -1604,6 +1608,7 @@ async def generate_key_helper_fn( tpm_limit = tpm_limit rpm_limit = rpm_limit allowed_cache_controls = allowed_cache_controls + try: # Create a new verification token (you may want to enhance this logic based on your needs) user_data = { @@ -1639,6 +1644,7 @@ async def generate_key_helper_fn( "budget_duration": key_budget_duration, "budget_reset_at": key_reset_at, "allowed_cache_controls": allowed_cache_controls, + "permissions": permissions, } if ( general_settings.get("allow_user_auth", False) == True @@ -1652,6 +1658,10 @@ async def generate_key_helper_fn( saved_token["config"] = json.loads(saved_token["config"]) if isinstance(saved_token["metadata"], str): saved_token["metadata"] = json.loads(saved_token["metadata"]) + if saved_token["permissions"] is not None and isinstance( + saved_token["permissions"], str + ): + saved_token["permissions"] = json.loads(saved_token["permissions"]) if saved_token.get("expires", None) is not None and isinstance( saved_token["expires"], datetime ): @@ -2965,6 +2975,7 @@ async def generate_key_fn( - max_budget: Optional[float] - Specify max budget for a given key. - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } + - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} Returns: - key: (str) The generated api key diff --git a/schema.prisma b/schema.prisma index e04ed580a..dd473bb69 100644 --- a/schema.prisma +++ b/schema.prisma @@ -63,6 +63,7 @@ model LiteLLM_VerificationToken { budget_duration String? budget_reset_at DateTime? allowed_cache_controls String[] @default([]) + permissions Json? } // store proxy config.yaml