forked from phoenix/litellm-mirror
Merge pull request #1926 from BerriAI/litellm_ui_further_improvements
Enable viewing key alias instead of hashed tokens
This commit is contained in:
Krish Dholakia
GitHub
commit
a36aa6aa01
6 changed files with 151 additions and 30 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,6 +497,7 @@ 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:
|
||||
if isinstance(token, str):
|
||||
hashed_token = token
|
||||
if token.startswith("sk-"):
|
||||
hashed_token = self.hash_token(token=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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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`;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
).then(async (response) => {
|
||||
const topKeysResponse = await keyInfoCall(
|
||||
accessToken,
|
||||
getTopKeys(response)
|
||||
);
|
||||
|
||||
setTopKeys(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}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue