phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Merge pull request #1926 from BerriAI/litellm_ui_further_improvements

Browse source 
Enable viewing key alias instead of hashed tokens
This commit is contained in:
Krish Dholakia 2024-02-10 20:52:23 -08:00 • committed by GitHub
commit a36aa6aa01
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 151 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

@ -187,8 +187,8 @@ class UpdateKeyRequest(GenerateKeyRequest):
metadata: Optional[dict] = None
class DeleteKeyRequest(LiteLLMBase):
keys: List
class KeyRequest(LiteLLMBase):
keys: List[str]
class NewUserRequest(GenerateKeyRequest):

69
litellm/proxy/proxy_server.py
View file

@ -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.
@ -2954,6 +2954,73 @@ async def delete_key_fn(data: DeleteKeyRequest):
)
@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.
**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
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)]
)

34
litellm/proxy/utils.py
View file

@ -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(t)
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

10
litellm/tests/test_key_generate_prisma.py
View file

@ -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)

30
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -202,6 +202,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`;

34
ui/litellm-dashboard/src/components/usage.tsx
View file

@ -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 {
@ -75,7 +75,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;
}
@ -164,18 +164,29 @@ const UsagePage: React.FC<UsagePageProps> = ({
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
@ -185,11 +196,6 @@ const UsagePage: React.FC<UsagePageProps> = ({
}
}, [accessToken, token, userRole, userID, startTime, endTime]);
topUsers.forEach((obj) => {
Object.values(obj).forEach((value) => {
console.log(value);
});
});
return (
<div style={{ width: "100%" }}>
<Grid numItems={2} className="gap-2 p-10 h-[75vh] w-full">
@ -217,7 +223,7 @@ const UsagePage: React.FC<UsagePageProps> = ({
index="key"
categories={["spend"]}
colors={["blue"]}
yAxisWidth={200}
yAxisWidth={80}
tickGap={5}
layout="vertical"
showXAxis={false}