From 1c987a436eff64147672c071854cd2d7aefd228e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:34:33 -0800 Subject: [PATCH 01/11] (docs) virtual_keys --- docs/my-website/docs/proxy/virtual_keys.md | 35 +++++++++++++++++----- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index ee8a99042..3adbb07ff 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -1,4 +1,4 @@ -# Key Management +# Virtual Keys Track Spend, Set budgets and create virtual keys for the proxy Grant other's temporary access to your proxy, with keys that expire after a set duration. @@ -12,7 +12,7 @@ Grant other's temporary access to your proxy, with keys that expire after a set ::: -## Quick Start +## Setup Requirements: @@ -58,16 +58,37 @@ litellm --config /path/to/config.yaml curl 'http://0.0.0.0:8000/key/generate' \ --header 'Authorization: Bearer ' \ --header 'Content-Type: application/json' \ ---data-raw '{"models": ["gpt-3.5-turbo", "gpt-4", "claude-2"], "duration": "20m","metadata": {"user": "ishaan@berri.ai", "team": "core-infra"}}' +--data-raw '{"models": ["gpt-3.5-turbo", "gpt-4", "claude-2"], "duration": "20m","metadata": {"user": "ishaan@berri.ai"}}' ``` + +## /key/generate + +### Request +```shell +curl 'http://0.0.0.0:8000/key/generate' \ +--header 'Authorization: Bearer ' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "models": ["gpt-3.5-turbo", "gpt-4", "claude-2"], + "duration": "20m", + "metadata": {"user": "ishaan@berri.ai"}, + "team_id": "core-infra" +}' +``` + + +Request Params: + - `models`: *list or null (optional)* - Specify the models a token has access too. If null, then token has access to all models on server. - `duration`: *str or null (optional)* Specify the length of time the token is valid for. If null, default is set to 1 hour. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - `metadata`: *dict or null (optional)* Pass metadata for the created token. If null defaults to {} -Expected response: +- `team_id`: *str or null (optional)* Specify team_id for the associated key + +### Response ```python { @@ -76,7 +97,7 @@ Expected response: } ``` -## Keys that don't expire +### Keys that don't expire Just set duration to None. @@ -87,7 +108,7 @@ curl --location 'http://0.0.0.0:8000/key/generate' \ --data '{"models": ["azure-models"], "aliases": {"mistral-7b": "gpt-3.5-turbo"}, "duration": null}' ``` -## Upgrade/Downgrade Models +### Upgrade/Downgrade Models If a user is expected to use a given model (i.e. gpt3-5), and you want to: @@ -137,7 +158,7 @@ curl -X POST "https://0.0.0.0:8000/key/generate" \ - **How are routing between diff keys/api bases done?** litellm handles this by shuffling between different models in the model list with the same model_name. [**See Code**](https://github.com/BerriAI/litellm/blob/main/litellm/router.py) -## Grant Access to new model +### Grant Access to new model Use model access groups to give users access to select models, and add new ones to it over time (e.g. mistral, llama-2, etc.) From ea32a8757bd4f984ca0e3cd58c1d9d84f426cad0 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:34:51 -0800 Subject: [PATCH 02/11] (feat) set team_id on virtual_keys --- litellm/proxy/_types.py | 1 + litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + 3 files changed, 3 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ad90173a4..e258582ef 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -129,6 +129,7 @@ class GenerateKeyRequest(LiteLLMBase): config: Optional[dict] = {} spend: Optional[float] = 0 user_id: Optional[str] = None + team_id: Optional[str] = None max_parallel_requests: Optional[int] = None metadata: Optional[dict] = {} diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index aa45a8818..1ed76140e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -24,6 +24,7 @@ model LiteLLM_VerificationToken { aliases Json @default("{}") config Json @default("{}") user_id String? + team_id String? max_parallel_requests Int? metadata Json @default("{}") } diff --git a/schema.prisma b/schema.prisma index 704ada42c..3c2bf22cb 100644 --- a/schema.prisma +++ b/schema.prisma @@ -24,6 +24,7 @@ model LiteLLM_VerificationToken { aliases Json @default("{}") config Json @default("{}") user_id String? + team_id String? max_parallel_requests Int? metadata Json @default("{}") } From 42ad12b2bd6f0cb0cb19870ef1c2b0ab71eeb696 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:48:52 -0800 Subject: [PATCH 03/11] (fix) support team_id for /key/generate --- litellm/proxy/proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 10c968b1c..e77223861 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1037,6 +1037,7 @@ async def generate_key_helper_fn( max_budget: Optional[float] = None, token: Optional[str] = None, user_id: Optional[str] = None, + team_id: Optional[str] = None, user_email: Optional[str] = None, max_parallel_requests: Optional[int] = None, metadata: Optional[dict] = {}, From 90509a159a4518b92ac0f1c0beec0feb7688066b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:54:08 -0800 Subject: [PATCH 04/11] (fix) write team_id to key table --- litellm/proxy/proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e77223861..f95137630 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1098,6 +1098,7 @@ async def generate_key_helper_fn( "config": config_json, "spend": spend, "user_id": user_id, + "team_id": team_id, "max_parallel_requests": max_parallel_requests, "metadata": metadata_json, } From f405a827e37c5fce5c23094a9b83138d2b6c635d Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 13:55:32 -0800 Subject: [PATCH 05/11] (docs) virtual_keys --- docs/my-website/docs/proxy/virtual_keys.md | 32 ++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 3adbb07ff..49f4ed23c 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -186,6 +186,38 @@ curl --location 'http://localhost:8000/key/generate' \ "max_budget": 0,}' ``` + + + + ## Tracking Spend You can get spend for a key by using the `/key/info` endpoint. From 5beef6dbcd32271b92d8c12a4e200c90cfc73cde Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 14:33:13 -0800 Subject: [PATCH 06/11] (test) setting team_id --- litellm/tests/test_key_generate_prisma.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 5accd03c6..ae51e4e96 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -543,7 +543,8 @@ def test_generate_and_update_key(prisma_client): async def test(): await litellm.proxy.proxy_server.prisma_client.connect() request = NewUserRequest( - metadata={"team": "litellm-team3", "project": "litellm-project3"} + metadata={"team": "litellm-team3", "project": "litellm-project3"}, + team_id="litellm-core-infra@gmail.com", ) key = await new_user(request) print(key) @@ -560,6 +561,7 @@ def test_generate_and_update_key(prisma_client): "team": "litellm-team3", "project": "litellm-project3", } + assert result["info"].team_id == "litellm-core-infra@gmail.com" request = Request(scope={"type": "http"}) request._url = URL(url="/update/key") From 2b6972111e8b3256453b0acdec24c22232e16774 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 14:42:46 -0800 Subject: [PATCH 07/11] (feat) write team_id to User Table --- litellm/proxy/proxy_server.py | 3 +++ litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + 3 files changed, 5 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f95137630..80e89709d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1081,12 +1081,15 @@ async def generate_key_helper_fn( config_json = json.dumps(config) metadata_json = json.dumps(metadata) user_id = user_id or str(uuid.uuid4()) + if type(team_id) is not str: + team_id = str(team_id) try: # Create a new verification token (you may want to enhance this logic based on your needs) user_data = { "max_budget": max_budget, "user_email": user_email, "user_id": user_id, + "team_id": team_id, "spend": spend, "models": models, } diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1ed76140e..24f7f4f3d 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -9,6 +9,7 @@ generator client { model LiteLLM_UserTable { user_id String @unique + team_id String? max_budget Float? spend Float @default(0.0) user_email String? diff --git a/schema.prisma b/schema.prisma index 3c2bf22cb..ae5c53a4d 100644 --- a/schema.prisma +++ b/schema.prisma @@ -9,6 +9,7 @@ generator client { model LiteLLM_UserTable { user_id String @unique + team_id String? max_budget Float? spend Float @default(0.0) user_email String? From d5e720e16109a1db7708b92ac49f654397ae7fc6 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 14:45:49 -0800 Subject: [PATCH 08/11] (docs) /key/update --- docs/my-website/docs/proxy/virtual_keys.md | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 49f4ed23c..391f1b60d 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -186,9 +186,9 @@ curl --location 'http://localhost:8000/key/generate' \ "max_budget": 0,}' ``` - +## /key/update - +``` ## Tracking Spend From 0e3e8050d7160f3c39565982a91d77d3bcfa2e75 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 14:54:11 -0800 Subject: [PATCH 09/11] (docs) /key/info --- docs/my-website/docs/proxy/virtual_keys.md | 36 +++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 391f1b60d..e9ce49078 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -186,6 +186,40 @@ curl --location 'http://localhost:8000/key/generate' \ "max_budget": 0,}' ``` + +## /key/info + +### Request +```shell +curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \ +-H "Authorization: Bearer sk-1234" +``` + +Request Params: +- key: str - The key you want the info for + +### Response + +```json +{ + "key": "sk-02Wr4IAlN3NvPXvL5JVvDA", + "info": { + "token": "80321a12d03412c527f2bd9db5fabd746abead2e1d50b435a534432fbaca9ef5", + "spend": 0.0, + "expires": "2024-01-18T23:52:09.125000+00:00", + "models": ["azure-gpt-3.5", "azure-embedding-model"], + "aliases": {}, + "config": {}, + "user_id": "ishaan2@berri.ai", + "team_id": "None", + "max_parallel_requests": null, + "metadata": {} + } +} + + +``` + ## /key/update ### Request @@ -212,7 +246,7 @@ Request Params: ### Response -```python +```json { "key": "sk-kdEXbIqZRwEeEiHwdg7sFA", "models": ["gpt-3.5-turbo", "gpt-4", "claude-2"], From cdede8836fc11245103f384afa64fa1053c58d33 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 15:16:40 -0800 Subject: [PATCH 10/11] (docs)virtual keys --- docs/my-website/docs/proxy/virtual_keys.md | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index e9ce49078..1cb28a2e3 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -200,6 +200,7 @@ Request Params: ### Response +`token` is the hashed key (The DB stores the hashed key for security) ```json { "key": "sk-02Wr4IAlN3NvPXvL5JVvDA", @@ -257,6 +258,30 @@ Request Params: ``` + +## /key/delete + +### Request +```shell +curl 'http://0.0.0.0:8000/key/delete' \ +--header 'Authorization: Bearer ' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "keys": ["sk-kdEXbIqZRwEeEiHwdg7sFA"] +}' +``` + +Request Params: +- keys: List[str] - List of keys to delete + +### Response + +```json +{ + "deleted_keys": ["sk-kdEXbIqZRwEeEiHwdg7sFA"] +} +``` + ## Tracking Spend You can get spend for a key by using the `/key/info` endpoint. From 340706565fc4a02c3a48a0bc1d06d30c07417d6b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 15:23:05 -0800 Subject: [PATCH 11/11] (fix) add team_id to doc string --- litellm/proxy/proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 80e89709d..6a6a59ee1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2048,6 +2048,7 @@ async def generate_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). **(Default is set to 1 hour.)** + - team_id: Optional[str] - The team id of the user - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models - config: Optional[dict] - any key-specific configs, overrides config in config.yaml