From 6b83001459de6d7700daa1ea01053404952d5103 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 10 Feb 2024 14:24:12 -0800 Subject: [PATCH 1/4] feat(proxy_server.py): support key lists for /key/info --- litellm/proxy/_types.py | 2 +- litellm/proxy/proxy_server.py | 40 +++++++++++++------ litellm/proxy/utils.py | 34 ++++++++++++---- ui/litellm-dashboard/src/components/usage.tsx | 6 +++ 4 files changed, 61 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ca5d4b05b..9b7ad75de 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -187,7 +187,7 @@ class UpdateKeyRequest(GenerateKeyRequest): metadata: Optional[dict] = None -class DeleteKeyRequest(LiteLLMBase): +class KeyRequest(LiteLLMBase): keys: List diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33eaae472..53162d29c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2894,7 +2894,7 @@ async def update_key_fn(request: Request, data: UpdateKeyRequest): @router.post( "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) -async def delete_key_fn(data: DeleteKeyRequest): +async def delete_key_fn(data: KeyRequest): """ Delete a key from the key management system. @@ -2958,6 +2958,7 @@ async def delete_key_fn(data: DeleteKeyRequest): "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) async def info_key_fn( + data: Optional[KeyRequest] = None, key: Optional[str] = fastapi.Query( default=None, description="Key in the request parameters" ), @@ -2966,16 +2967,23 @@ async def info_key_fn( """ Retrieve information about a key. Parameters: - key: Optional[str] = Query parameter representing the key in the request + key: Optional[str | list] = Query parameter representing the key(s) in the request user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key Returns: Dict containing the key and its associated information Example Curl: + **1 key** ``` curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \ -H "Authorization: Bearer sk-1234" ``` + **multiple keys** + ``` + curl -X GET "http://0.0.0.0:8000/key/info" \ +-H "Authorization: Bearer sk-1234" \ +-d {"keys": ["sk-1", "sk-2", "sk-3"]} + ``` Example Curl - if no key is passed, it will use the Key Passed in Authorization Header ``` @@ -2989,16 +2997,24 @@ async def info_key_fn( raise Exception( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - if key == None: - key = user_api_key_dict.api_key - key_info = await prisma_client.get_data(token=key) - ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## - try: - key_info = key_info.model_dump() # noqa - except: - # if using pydantic v1 - key_info = key_info.dict() - key_info.pop("token") + if key is None and data is None: + key_info = await prisma_client.get_data(token=user_api_key_dict.api_key) + + if key is not None: # single key + key_info = await prisma_client.get_data(token=key) + elif data is not None: # list of keys + key_info = await prisma_client.get_data( + token=data.keys, table_name="key", query_type="find_all" + ) + filtered_key_info = [] + for k in key_info: + try: + k = k.model_dump() # noqa + except: + # if using pydantic v1 + k = k.dict() + filtered_key_info.append(k) + return {"key": data.keys, "info": filtered_key_info} return {"key": key, "info": key_info} except Exception as e: if isinstance(e, HTTPException): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 35b647257..41646149b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -483,7 +483,7 @@ class PrismaClient: ) async def get_data( self, - token: Optional[str] = None, + token: Optional[Union[str, list]] = None, user_id: Optional[str] = None, user_id_list: Optional[list] = None, key_val: Optional[dict] = None, @@ -497,12 +497,13 @@ class PrismaClient: if token is not None or (table_name is not None and table_name == "key"): # check if plain text or hash if token is not None: - hashed_token = token - if token.startswith("sk-"): - hashed_token = self.hash_token(token=token) - verbose_proxy_logger.debug( - f"PrismaClient: find_unique for token: {hashed_token}" - ) + if isinstance(token, str): + hashed_token = token + if token.startswith("sk-"): + hashed_token = self.hash_token(token=token) + verbose_proxy_logger.debug( + f"PrismaClient: find_unique for token: {hashed_token}" + ) if query_type == "find_unique": response = await self.db.litellm_verificationtoken.find_unique( where={"token": hashed_token} @@ -540,8 +541,25 @@ class PrismaClient: if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() elif query_type == "find_all": + where_filter: dict = {} + if token is not None: + where_filter["token"] = {} + if isinstance(token, str): + if token.startswith("sk-"): + token = self.hash_token(token=token) + where_filter["token"]["in"] = [token] + elif isinstance(token, list): + hashed_tokens = [] + for t in token: + assert isinstance(t, str) + if t.startswith("sk-"): + new_token = self.hash_token(token=t) + hashed_tokens.append(new_token) + else: + hashed_tokens.append(new_token) + where_filter["token"]["in"] = hashed_tokens response = await self.db.litellm_verificationtoken.find_many( - order={"spend": "desc"}, + order={"spend": "desc"}, where=where_filter # type: ignore ) if response is not None: return response diff --git a/ui/litellm-dashboard/src/components/usage.tsx b/ui/litellm-dashboard/src/components/usage.tsx index 4b80332a7..7b113e812 100644 --- a/ui/litellm-dashboard/src/components/usage.tsx +++ b/ui/litellm-dashboard/src/components/usage.tsx @@ -187,6 +187,12 @@ const UsagePage: React.FC = ({ }; fetchData(); } + if (topKeys.length > 0) { + /** + * get the key alias / secret names for the created keys + * Use that for the key name instead of token hash for easier association + */ + } }, [accessToken, token, userRole, userID, startTime, endTime]); topUsers.forEach((obj) => { From 34c118e5e94a8e11b9532e25af24df98fc892231 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 10 Feb 2024 15:45:42 -0800 Subject: [PATCH 2/4] feat(ui): show key alias instead of hashed token --- litellm/proxy/_types.py | 2 +- litellm/proxy/proxy_server.py | 81 +++++++++++++++---- litellm/proxy/utils.py | 2 +- .../src/components/networking.tsx | 30 +++++++ ui/litellm-dashboard/src/components/usage.tsx | 40 ++++----- 5 files changed, 118 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b7ad75de..0d7355dad 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -188,7 +188,7 @@ class UpdateKeyRequest(GenerateKeyRequest): class KeyRequest(LiteLLMBase): - keys: List + keys: List[str] class NewUserRequest(GenerateKeyRequest): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 53162d29c..52420611e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2954,11 +2954,75 @@ async def delete_key_fn(data: KeyRequest): ) +@router.post( + "/v2/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)] +) +async def info_key_fn_v2( + data: Optional[KeyRequest] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Retrieve information about a list of keys. + Parameters: + keys: Optional[list] = body parameter representing the key(s) in the request + user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + Returns: + Dict containing the key and its associated information + + Example Curl: + ``` + curl -X GET "http://0.0.0.0:8000/key/info" \ +-H "Authorization: Bearer sk-1234" \ +-d {"keys": ["sk-1", "sk-2", "sk-3"]} + ``` + """ + global prisma_client + try: + if prisma_client is None: + raise Exception( + f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + ) + if data is None: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": "Malformed request. No keys passed in."}, + ) + + key_info = await prisma_client.get_data( + token=data.keys, table_name="key", query_type="find_all" + ) + filtered_key_info = [] + for k in key_info: + try: + k = k.model_dump() # noqa + except: + # if using pydantic v1 + k = k.dict() + filtered_key_info.append(k) + return {"key": data.keys, "info": filtered_key_info} + + except Exception as e: + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Authentication Error({str(e)})"), + type="auth_error", + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type="auth_error", + param=getattr(e, "param", "None"), + code=status.HTTP_400_BAD_REQUEST, + ) + + @router.get( "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) async def info_key_fn( - data: Optional[KeyRequest] = None, key: Optional[str] = fastapi.Query( default=None, description="Key in the request parameters" ), @@ -2997,24 +3061,11 @@ async def info_key_fn( raise Exception( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - if key is None and data is None: + if key is None: key_info = await prisma_client.get_data(token=user_api_key_dict.api_key) if key is not None: # single key key_info = await prisma_client.get_data(token=key) - elif data is not None: # list of keys - key_info = await prisma_client.get_data( - token=data.keys, table_name="key", query_type="find_all" - ) - filtered_key_info = [] - for k in key_info: - try: - k = k.model_dump() # noqa - except: - # if using pydantic v1 - k = k.dict() - filtered_key_info.append(k) - return {"key": data.keys, "info": filtered_key_info} return {"key": key, "info": key_info} except Exception as e: if isinstance(e, HTTPException): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 41646149b..bb2625dce 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -556,7 +556,7 @@ class PrismaClient: new_token = self.hash_token(token=t) hashed_tokens.append(new_token) else: - hashed_tokens.append(new_token) + hashed_tokens.append(t) where_filter["token"]["in"] = hashed_tokens response = await self.db.litellm_verificationtoken.find_many( order={"spend": "desc"}, where=where_filter # type: ignore diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 7026f394a..06a436053 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -200,6 +200,36 @@ export const userSpendLogsCall = async ( } }; +export const keyInfoCall = async (accessToken: String, keys: String[]) => { + try { + let url = proxyBaseUrl ? `${proxyBaseUrl}/v2/key/info` : `/key/info`; + + const response = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + keys: keys, + }), + }); + + if (!response.ok) { + const errorData = await response.text(); + message.error(errorData); + throw new Error("Network response was not ok"); + } + + const data = await response.json(); + console.log(data); + return data; + } catch (error) { + console.error("Failed to create key:", error); + throw error; + } +}; + export const spendUsersCall = async (accessToken: String, userID: String) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/spend/users` : `/spend/users`; diff --git a/ui/litellm-dashboard/src/components/usage.tsx b/ui/litellm-dashboard/src/components/usage.tsx index 7b113e812..76c8d7fdd 100644 --- a/ui/litellm-dashboard/src/components/usage.tsx +++ b/ui/litellm-dashboard/src/components/usage.tsx @@ -2,7 +2,7 @@ import { BarChart, Card, Title } from "@tremor/react"; import React, { useState, useEffect } from "react"; import { Grid, Col, Text, LineChart } from "@tremor/react"; -import { userSpendLogsCall } from "./networking"; +import { userSpendLogsCall, keyInfoCall } from "./networking"; import { start } from "repl"; interface UsagePageProps { @@ -79,7 +79,7 @@ function getTopKeys(data: Array<{ [key: string]: unknown }>): any[] { spendKeys.sort((a, b) => Number(b.spend) - Number(a.spend)); - const topKeys = spendKeys.slice(0, 5); + const topKeys = spendKeys.slice(0, 5).map((k) => k.key); console.log(`topKeys: ${Object.keys(topKeys[0])}`); return topKeys; } @@ -168,18 +168,29 @@ const UsagePage: React.FC = ({ if (accessToken && token && userRole && userID) { const fetchData = async () => { try { - const response = await userSpendLogsCall( + await userSpendLogsCall( accessToken, token, userRole, userID, startTime, endTime - ); - - setTopKeys(getTopKeys(response)); - setTopUsers(getTopUsers(response)); - setKeySpendData(response); + ).then(async (response) => { + const topKeysResponse = await keyInfoCall( + accessToken, + getTopKeys(response) + ); + const filtered_keys = topKeysResponse["info"].map((k: any) => ({ + key: (k["key_name"] || k["key_alias"] || k["token"]).substring( + 0, + 7 + ), + spend: k["spend"], + })); + setTopKeys(filtered_keys); + setTopUsers(getTopUsers(response)); + setKeySpendData(response); + }); } catch (error) { console.error("There was an error fetching the data", error); // Optionally, update your UI to reflect the error state here as well @@ -187,19 +198,8 @@ const UsagePage: React.FC = ({ }; fetchData(); } - if (topKeys.length > 0) { - /** - * get the key alias / secret names for the created keys - * Use that for the key name instead of token hash for easier association - */ - } }, [accessToken, token, userRole, userID, startTime, endTime]); - topUsers.forEach((obj) => { - Object.values(obj).forEach((value) => { - console.log(value); - }); - }); return (
@@ -227,7 +227,7 @@ const UsagePage: React.FC = ({ index="key" categories={["spend"]} colors={["blue"]} - yAxisWidth={200} + yAxisWidth={80} tickGap={5} layout="vertical" showXAxis={false} From f4dc643b1c57eb2b1546ff864891c544a16c4706 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 10 Feb 2024 16:16:12 -0800 Subject: [PATCH 3/4] test(test_key_generate_prisma.py): fix test import --- litellm/tests/test_key_generate_prisma.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index d4f405b7b..89d4f4d3e 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -59,7 +59,7 @@ from litellm.proxy._types import ( NewUserRequest, GenerateKeyRequest, DynamoDBArgs, - DeleteKeyRequest, + KeyRequest, UpdateKeyRequest, GenerateKeyRequest, ) @@ -740,7 +740,7 @@ def test_delete_key(prisma_client): generated_key = key.key bearer_token = "Bearer " + generated_key - delete_key_request = DeleteKeyRequest(keys=[generated_key]) + delete_key_request = KeyRequest(keys=[generated_key]) # delete the key result_delete_key = await delete_key_fn(data=delete_key_request) @@ -778,7 +778,7 @@ def test_delete_key_auth(prisma_client): generated_key = key.key bearer_token = "Bearer " + generated_key - delete_key_request = DeleteKeyRequest(keys=[generated_key]) + delete_key_request = KeyRequest(keys=[generated_key]) # delete the key result_delete_key = await delete_key_fn(data=delete_key_request) @@ -839,7 +839,7 @@ def test_generate_and_call_key_info(prisma_client): } # cleanup - delete key - delete_key_request = DeleteKeyRequest(keys=[generated_key]) + delete_key_request = KeyRequest(keys=[generated_key]) # delete the key await delete_key_fn(data=delete_key_request) @@ -908,7 +908,7 @@ def test_generate_and_update_key(prisma_client): assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"] # cleanup - delete key - delete_key_request = DeleteKeyRequest(keys=[generated_key]) + delete_key_request = KeyRequest(keys=[generated_key]) # delete the key await delete_key_fn(data=delete_key_request) From 7e2e545bed4456da64f8cc7d7c99f62b103cf14e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 10 Feb 2024 17:31:23 -0800 Subject: [PATCH 4/4] fix(proxy_server.py): don't change old /key/info endpoint --- litellm/proxy/proxy_server.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 52420611e..c5c843177 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2963,6 +2963,8 @@ async def info_key_fn_v2( ): """ Retrieve information about a list of keys. + + **New endpoint**. Currently admin only. Parameters: keys: Optional[list] = body parameter representing the key(s) in the request user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key @@ -3031,23 +3033,16 @@ async def info_key_fn( """ Retrieve information about a key. Parameters: - key: Optional[str | list] = Query parameter representing the key(s) in the request + key: Optional[str] = Query parameter representing the key in the request user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key Returns: Dict containing the key and its associated information Example Curl: - **1 key** ``` curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \ -H "Authorization: Bearer sk-1234" ``` - **multiple keys** - ``` - curl -X GET "http://0.0.0.0:8000/key/info" \ --H "Authorization: Bearer sk-1234" \ --d {"keys": ["sk-1", "sk-2", "sk-3"]} - ``` Example Curl - if no key is passed, it will use the Key Passed in Authorization Header ``` @@ -3061,11 +3056,16 @@ async def info_key_fn( raise Exception( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - if key is None: - key_info = await prisma_client.get_data(token=user_api_key_dict.api_key) - - if key is not None: # single key - key_info = await prisma_client.get_data(token=key) + if key == None: + key = user_api_key_dict.api_key + key_info = await prisma_client.get_data(token=key) + ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## + try: + key_info = key_info.model_dump() # noqa + except: + # if using pydantic v1 + key_info = key_info.dict() + key_info.pop("token") return {"key": key, "info": key_info} except Exception as e: if isinstance(e, HTTPException):