diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 3a88c3b3d572..dc48b21aea90 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -353,7 +353,8 @@ Request Params: } ``` -## Upperbound /key/generate params +## Advanced +### Upperbound /key/generate params Use this, if you need to control the upperbound that users can use for `max_budget`, `budget_duration` or any `key/generate` param per key. Set `litellm_settings:upperbound_key_generate_params`: @@ -369,7 +370,7 @@ litellm_settings: - Send a `/key/generate` request with `max_budget=200` - Key will be created with `max_budget=100` since 100 is the upper bound -## Default /key/generate params +### Default /key/generate params Use this, if you need to control the default `max_budget` or any `key/generate` param per key. When a `/key/generate` request does not specify `max_budget`, it will use the `max_budget` specified in `default_key_generate_params` @@ -384,7 +385,46 @@ litellm_settings: metadata: {"setting":"default"} team_id: "core-infra" ``` -## Set Budgets - Per Key + +### Restrict models by `team_id` +`litellm-dev` can only access `azure-gpt-3.5` + +```yaml +litellm_settings: + default_team_settings: + - team_id: litellm-dev + models: ["azure-gpt-3.5"] +``` + +#### Create key with team_id="litellm-dev" +```shell +curl --location 'http://localhost:4000/key/generate' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data-raw '{"team_id": "litellm-dev"}' +``` + +#### Use Key to call invalid model - Fails +```shell +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-qo992IjKOC2CHKZGRoJIGA' \ + --data '{ + "model": "BEDROCK_GROUP", + "messages": [ + { + "role": "user", + "content": "hi" + } + ] + }' +``` + +```shell +{"error":{"message":"Invalid model for team litellm-dev: BEDROCK_GROUP. Valid models for team are: ['azure-gpt-3.5']\n\n\nTraceback (most recent call last):\n File \"/Users/ishaanjaffer/Github/litellm/litellm/proxy/proxy_server.py\", line 2298, in chat_completion\n _is_valid_team_configs(\n File \"/Users/ishaanjaffer/Github/litellm/litellm/proxy/utils.py\", line 1296, in _is_valid_team_configs\n raise Exception(\nException: Invalid model for team litellm-dev: BEDROCK_GROUP. Valid models for team are: ['azure-gpt-3.5']\n\n","type":"None","param":"None","code":500}}% +``` + +### Set Budgets - Per Key Set `max_budget` in (USD $) param in the `key/generate` request. By default the `max_budget` is set to `null` and is not checked for keys @@ -430,7 +470,7 @@ Expected Response from `/chat/completions` when key has crossed budget ``` -## Set Budgets - Per User +### Set Budgets - Per User LiteLLM exposes a `/user/new` endpoint to create budgets for users, that persist across multiple keys. @@ -455,7 +495,7 @@ The request is a normal `/key/generate` request body + a `max_budget` field. } ``` -## Tracking Spend +### Tracking Spend You can get spend for a key by using the `/key/info` endpoint. @@ -490,13 +530,13 @@ This is automatically updated (in USD) when calls are made to /completions, /cha ``` -## Custom Auth +### Custom Auth You can now override the default api key auth. Here's how: -### 1. Create a custom auth file. +#### 1. Create a custom auth file. Make sure the response type follows the `UserAPIKeyAuth` pydantic object. This is used by for logging usage specific to that user key. @@ -513,7 +553,7 @@ async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth: raise Exception ``` -### 2. Pass the filepath (relative to the config.yaml) +#### 2. Pass the filepath (relative to the config.yaml) Pass the filepath to the config.yaml @@ -534,16 +574,16 @@ general_settings: [**Implementation Code**](https://github.com/BerriAI/litellm/blob/caf2a6b279ddbe89ebd1d8f4499f65715d684851/litellm/proxy/utils.py#L122) -### 3. Start the proxy +#### 3. Start the proxy ```shell $ litellm --config /path/to/config.yaml ``` -## Custom /key/generate +### 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` +#### 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) @@ -609,7 +649,7 @@ async def custom_generate_key_fn(data: GenerateKeyRequest)-> dict: ``` -### 2. Pass the filepath (relative to the config.yaml) +#### 2. Pass the filepath (relative to the config.yaml) Pass the filepath to the config.yaml @@ -631,18 +671,18 @@ general_settings: -## [BETA] Dynamo DB +### [BETA] Dynamo DB Only live in `v1.16.21.dev1`. -### Step 1. Save keys to env +#### Step 1. Save keys to env ```shell AWS_ACCESS_KEY_ID = "your-aws-access-key-id" AWS_SECRET_ACCESS_KEY = "your-aws-secret-access-key" ``` -### Step 2. Add details to config +#### Step 2. Add details to config ```yaml general_settings: @@ -657,7 +697,7 @@ general_settings: } ``` -### Step 3. Generate Key +#### Step 3. Generate Key ```bash curl --location 'http://0.0.0.0:8000/key/generate' \ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5c843177210..e018f7e078b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -92,6 +92,7 @@ def generate_feedback_box(): hash_token, html_form, _read_request_body, + _is_valid_team_configs, ) from litellm.proxy.secret_managers.google_kms import load_google_kms import pydantic @@ -2294,6 +2295,9 @@ async def chat_completion( pass else: team_id = team_config.pop("team_id", None) + _is_valid_team_configs( + team_id=team_id, team_config=team_config, request_data=data + ) data["metadata"]["team_id"] = team_id data = { **team_config, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bb2625dced8a..6506e6ab3c72 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1303,6 +1303,20 @@ async def _read_request_body(request): return request_data +def _is_valid_team_configs(team_id=None, team_config=None, request_data=None): + if team_id is None or team_config is None or request_data is None: + return + # check if valid model called for team + if "models" in team_config: + valid_models = team_config.pop("models") + model_in_request = request_data["model"] + if model_in_request not in valid_models: + raise Exception( + f"Invalid model for team {team_id}: {model_in_request}. Valid models for team are: {valid_models}\n" + ) + return + + # LiteLLM Admin UI - Non SSO Login html_form = """