scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #5536 from BerriAI/litellm_allow_internal_user_view_usage

Browse source 
[Fix-Proxy] allow internal user and internal viewer to view usage
This commit is contained in:
Ishaan Jaff 2024-09-05 16:46:12 -07:00 • committed by GitHub
commit fa9a1f4dbc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 876 additions and 167 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
.circleci/config.yml
View file

@ -111,6 +111,44 @@ jobs:
# Store test results
- store_test_results:
path: test-results
ui_endpoint_testing:
docker:
- image: cimg/python:3.11
working_directory: ~/project
steps:
- checkout
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-asyncio==0.21.1"
- save_cache:
paths:
- ./venv
key: v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Run prisma ./entrypoint.sh
command: |
set +e
chmod +x entrypoint.sh
./entrypoint.sh
set -e
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
pwd
ls
python -m pytest -vv tests/proxy_admin_ui_tests -x --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 120m
# Store test results
- store_test_results:
path: test-results
installing_litellm_on_python:
docker:
@ -251,7 +289,7 @@ jobs:
command: |
pwd
ls
python -m pytest -s -vv tests/ -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests
python -m pytest -s -vv tests/ -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests
no_output_timeout: 120m
# Store test results
@ -539,6 +577,12 @@ workflows:
only:
- main
- /litellm_.*/
- ui_endpoint_testing:
filters:
branches:
only:
- main
- /litellm_.*/
- build_and_test:
filters:
branches:

7
litellm/proxy/_types.py
View file

@ -344,6 +344,13 @@ class LiteLLMRoutes(enum.Enum):
"/key/update",
"/key/delete",
"/key/info",
"/global/spend/tags",
"/global/spend/keys",
"/global/spend/models",
"/global/spend/provider",
"/global/spend/end_users",
"/global/activity",
"/global/activity/model",
]
+ spend_tracking_routes
+ sso_only_routes

343
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -177,6 +177,35 @@ async def view_spend_tags(
)
async def get_global_activity_internal_user(
user_api_key_dict: UserAPIKeyAuth, start_date: datetime, end_date: datetime
):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
user_id = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(status_code=500, detail={"error": "No user_id found"})
sql_query = """
SELECT
date_trunc('day', "startTime") AS date,
COUNT(*) AS api_requests,
SUM(total_tokens) AS total_tokens
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day'
AND "user" = $3
GROUP BY date_trunc('day', "startTime")
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date, end_date, user_id
)
return db_response
@router.get(
"/global/activity",
tags=["Budget & Spend Tracking"],
@ -195,6 +224,7 @@ async def get_global_activity(
default=None,
description="Time till which to view spend",
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get number of API Requests, total tokens through proxy
@ -236,18 +266,27 @@ async def get_global_activity(
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
sql_query = """
SELECT
date_trunc('day', "startTime") AS date,
COUNT(*) AS api_requests,
SUM(total_tokens) AS total_tokens
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day'
GROUP BY date_trunc('day', "startTime")
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
db_response = await get_global_activity_internal_user(
user_api_key_dict, start_date_obj, end_date_obj
)
else:
sql_query = """
SELECT
date_trunc('day', "startTime") AS date,
COUNT(*) AS api_requests,
SUM(total_tokens) AS total_tokens
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day'
GROUP BY date_trunc('day', "startTime")
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
if db_response is None:
return []
@ -282,6 +321,36 @@ async def get_global_activity(
)
async def get_global_activity_model_internal_user(
user_api_key_dict: UserAPIKeyAuth, start_date: datetime, end_date: datetime
):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
user_id = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(status_code=500, detail={"error": "No user_id found"})
sql_query = """
SELECT
model_group,
date_trunc('day', "startTime") AS date,
COUNT(*) AS api_requests,
SUM(total_tokens) AS total_tokens
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day'
AND "user" = $3
GROUP BY model_group, date_trunc('day', "startTime")
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date, end_date, user_id
)
return db_response
@router.get(
"/global/activity/model",
tags=["Budget & Spend Tracking"],
@ -300,6 +369,7 @@ async def get_global_activity_model(
default=None,
description="Time till which to view spend",
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get number of API Requests, total tokens through proxy - Grouped by MODEL
@ -364,19 +434,28 @@ async def get_global_activity_model(
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
sql_query = """
SELECT
model_group,
date_trunc('day', "startTime") AS date,
COUNT(*) AS api_requests,
SUM(total_tokens) AS total_tokens
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day'
GROUP BY model_group, date_trunc('day', "startTime")
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
db_response = await get_global_activity_model_internal_user(
user_api_key_dict, start_date_obj, end_date_obj
)
else:
sql_query = """
SELECT
model_group,
date_trunc('day', "startTime") AS date,
COUNT(*) AS api_requests,
SUM(total_tokens) AS total_tokens
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day'
GROUP BY model_group, date_trunc('day', "startTime")
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
if db_response is None:
return []
@ -711,6 +790,7 @@ async def get_global_spend_provider(
default=None,
description="Time till which to view spend",
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get breakdown of spend per provider
@ -748,19 +828,42 @@ async def get_global_spend_provider(
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
sql_query = """
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
user_id = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(
status_code=400, detail={"error": "No user_id found"}
)
SELECT
model_id,
SUM(spend) AS spend
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date AND length(model_id) > 0
GROUP BY model_id
"""
sql_query = """
SELECT
model_id,
SUM(spend) AS spend
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date
AND length(model_id) > 0
AND "user" = $3
GROUP BY model_id
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj, user_id
)
else:
sql_query = """
SELECT
model_id,
SUM(spend) AS spend
FROM "LiteLLM_SpendLogs"
WHERE "startTime" BETWEEN $1::date AND $2::date AND length(model_id) > 0
GROUP BY model_id
"""
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
db_response = await prisma_client.db.query_raw(
sql_query, start_date_obj, end_date_obj
)
if db_response is None:
return []
@ -1515,6 +1618,7 @@ async def view_spend_logs(
default=None,
description="Time till which to view key spend",
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
View all spend logs, if request_id is provided, only logs for that request_id will be returned
@ -1545,6 +1649,12 @@ async def view_spend_logs(
"""
from litellm.proxy.proxy_server import prisma_client
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
user_id = user_api_key_dict.user_id
try:
verbose_proxy_logger.debug("inside view_spend_logs")
if prisma_client is None:
@ -1733,6 +1843,45 @@ async def global_spend_reset():
}
async def global_spend_for_internal_user(
api_key: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise ProxyException(
message="Prisma Client is not initialized",
type="internal_error",
param="None",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
try:
user_id = user_api_key_dict.user_id
if user_id is None:
raise ValueError(f"/global/spend/logs Error: User ID is None")
if api_key is not None:
sql_query = """
SELECT * FROM "MonthlyGlobalSpendPerUserPerKey"
WHERE "api_key" = $1 AND "user" = $2
ORDER BY "date";
"""
response = await prisma_client.db.query_raw(sql_query, api_key, user_id)
return response
sql_query = """SELECT * FROM "MonthlyGlobalSpendPerUserPerKey" WHERE "user" = $1 ORDER BY "date";"""
response = await prisma_client.db.query_raw(sql_query, user_id)
return response
except Exception as e:
verbose_proxy_logger.error(f"/global/spend/logs Error: {str(e)}")
raise e
@router.get(
"/global/spend/logs",
tags=["Budget & Spend Tracking"],
@ -1743,7 +1892,8 @@ async def global_spend_logs(
api_key: str = fastapi.Query(
default=None,
description="API Key to get global spend (spend per day for last 30d). Admin-only endpoint",
)
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
[BETA] This is a beta endpoint. It will change.
@ -1764,6 +1914,17 @@ async def global_spend_logs(
param="None",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
response = await global_spend_for_internal_user(
api_key=api_key, user_api_key_dict=user_api_key_dict
)
return response
if api_key is None:
sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";"""
@ -1784,6 +1945,7 @@ async def global_spend_logs(
except Exception as e:
error_trace = traceback.format_exc()
error_str = str(e) + "\n" + error_trace
verbose_proxy_logger.error(f"/global/spend/logs Error: {error_str}")
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"/global/spend/logs Error({error_str})"),
@ -1851,6 +2013,52 @@ async def global_spend():
)
async def global_spend_key_internal_user(
user_api_key_dict: UserAPIKeyAuth, limit: int = 10
):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
user_id = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(status_code=500, detail={"error": "No user_id found"})
sql_query = """
WITH top_api_keys AS (
SELECT
api_key,
SUM(spend) as total_spend
FROM
"LiteLLM_SpendLogs"
WHERE
"user" = $1
GROUP BY
api_key
ORDER BY
total_spend DESC
LIMIT $2 -- Adjust this number to get more or fewer top keys
)
SELECT
t.api_key,
t.total_spend,
v.key_alias,
v.key_name
FROM
top_api_keys t
LEFT JOIN
"LiteLLM_VerificationToken" v ON t.api_key = v.token
ORDER BY
t.total_spend DESC;
"""
response = await prisma_client.db.query_raw(sql_query, user_id, limit)
return response
@router.get(
"/global/spend/keys",
tags=["Budget & Spend Tracking"],
@ -1861,7 +2069,8 @@ async def global_spend_keys(
limit: int = fastapi.Query(
default=None,
description="Number of keys to get. Will return Top 'n' keys.",
)
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
[BETA] This is a beta endpoint. It will change.
@ -1870,6 +2079,15 @@ async def global_spend_keys(
"""
from litellm.proxy.proxy_server import prisma_client
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
response = await global_spend_key_internal_user(
user_api_key_dict=user_api_key_dict
)
return response
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
sql_query = f"""SELECT * FROM "Last30dKeysBySpend" LIMIT {limit};"""
@ -2057,6 +2275,39 @@ LIMIT 100
return response
async def global_spend_models_internal_user(
user_api_key_dict: UserAPIKeyAuth, limit: int = 10
):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
user_id = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(status_code=500, detail={"error": "No user_id found"})
sql_query = """
SELECT
model,
SUM(spend) as total_spend,
SUM(total_tokens) as total_tokens
FROM
"LiteLLM_SpendLogs"
WHERE
"user" = $1
GROUP BY
model
ORDER BY
total_spend DESC
LIMIT $2;
"""
response = await prisma_client.db.query_raw(sql_query, user_id, limit)
return response
@router.get(
"/global/spend/models",
tags=["Budget & Spend Tracking"],
@ -2065,17 +2316,27 @@ LIMIT 100
)
async def global_spend_models(
limit: int = fastapi.Query(
default=None,
default=10,
description="Number of models to get. Will return Top 'n' models.",
)
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
[BETA] This is a beta endpoint. It will change.
Use this to get the top 'n' keys with the highest spend, ordered by spend.
Use this to get the top 'n' models with the highest spend, ordered by spend.
"""
from litellm.proxy.proxy_server import prisma_client
if (
user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
):
response = await global_spend_models_internal_user(
user_api_key_dict=user_api_key_dict, limit=limit
)
return response
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})

28
litellm/proxy/utils.py
View file

@ -969,11 +969,12 @@ class PrismaClient:
'Last30dKeysBySpend',
'Last30dModelsBySpend',
'MonthlyGlobalSpendPerKey',
'MonthlyGlobalSpendPerUserPerKey',
'Last30dTopEndUsersSpend'
)
"""
)
if ret[0]["sum"] == 6:
if ret[0]["sum"] == 7:
print("All necessary views exist!") # noqa
return
except Exception:
@ -1097,6 +1098,31 @@ class PrismaClient:
await self.db.execute_raw(query=sql_query)
print("MonthlyGlobalSpendPerKey Created!") # noqa
try:
await self.db.query_raw(
"""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1"""
)
print("MonthlyGlobalSpendPerUserPerKey Exists!") # noqa
except Exception as e:
sql_query = """
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerUserPerKey" AS
SELECT
DATE("startTime") AS date,
SUM("spend") AS spend,
api_key as api_key,
"user" as "user"
FROM
"LiteLLM_SpendLogs"
WHERE
"startTime" >= (CURRENT_DATE - INTERVAL '20 days')
GROUP BY
DATE("startTime"),
"user",
api_key;
"""
await self.db.execute_raw(query=sql_query)
print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa
try:
await self.db.query_raw(

5
litellm/tests/test_completion.py
View file

@ -11,7 +11,8 @@ import os
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system-path
) # Adds the parent directory to the systempath
import os
from unittest.mock import AsyncMock, MagicMock, patch
@ -23,7 +24,7 @@ from litellm import RateLimitError, Timeout, completion, completion_cost, embedd
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.prompt_templates.factory import anthropic_messages_pt
# litellm.num_retries=3
# litellm.num_retries = 3
litellm.cache = None
litellm.success_callback = []
user_message = "Write a short poem about the sky"

21
litellm/tests/test_key_generate_prisma.py
View file

@ -537,6 +537,7 @@ def test_call_with_user_over_budget(prisma_client):
asyncio.run(test())
except Exception as e:
print("got an errror=", e)
error_detail = e.message
assert "ExceededBudget:" in error_detail
assert isinstance(e, ProxyException)
@ -1492,7 +1493,10 @@ def test_call_with_key_over_budget(prisma_client):
proxy_logging_obj=proxy_logging_obj,
)
# test spend_log was written and we can read it
spend_logs = await view_spend_logs(request_id=request_id)
spend_logs = await view_spend_logs(
request_id=request_id,
user_api_key_dict=UserAPIKeyAuth(api_key=generated_key),
)
print("read spend logs", spend_logs)
assert len(spend_logs) == 1
@ -1607,7 +1611,10 @@ def test_call_with_key_over_budget_no_cache(prisma_client):
proxy_logging_obj=proxy_logging_obj,
)
# test spend_log was written and we can read it
spend_logs = await view_spend_logs(request_id=request_id)
spend_logs = await view_spend_logs(
request_id=request_id,
user_api_key_dict=UserAPIKeyAuth(api_key=generated_key),
)
print("read spend logs", spend_logs)
assert len(spend_logs) == 1
@ -1727,7 +1734,10 @@ def test_call_with_key_over_model_budget(prisma_client):
proxy_logging_obj=proxy_logging_obj,
)
# test spend_log was written and we can read it
spend_logs = await view_spend_logs(request_id=request_id)
spend_logs = await view_spend_logs(
request_id=request_id,
user_api_key_dict=UserAPIKeyAuth(api_key=generated_key),
)
print("read spend logs", spend_logs)
assert len(spend_logs) == 1
@ -2296,7 +2306,10 @@ async def test_proxy_load_test_db(prisma_client):
await asyncio.sleep(120)
try:
# call spend logs
spend_logs = await view_spend_logs(api_key=generated_key)
spend_logs = await view_spend_logs(
api_key=generated_key,
user_api_key_dict=UserAPIKeyAuth(api_key=generated_key),
)
print(f"len responses: {len(spend_logs)}")
assert len(spend_logs) == n

321
tests/proxy_admin_ui_tests/test_usage_endpoints.py Normal file
View file

@ -0,0 +1,321 @@
"""
Tests the following endpoints used by the UI
/global/spend/logs
/global/spend/keys
/global/spend/models
/global/activity
/global/activity/model
For all tests - test the following:
- Response is valid
- Response for Admin User is different from response from Internal User
"""
import os
import sys
import traceback
import uuid
from datetime import datetime
from dotenv import load_dotenv
from fastapi import Request
from fastapi.routing import APIRoute
load_dotenv()
import io
import os
import time
# this file is to test litellm/proxy
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import logging
import pytest
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.management_endpoints.internal_user_endpoints import (
new_user,
user_info,
user_update,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_key_fn,
generate_key_fn,
generate_key_helper_fn,
info_key_fn,
regenerate_key_fn,
update_key_fn,
)
from litellm.proxy.management_endpoints.team_endpoints import (
new_team,
team_info,
update_team,
)
from litellm.proxy.proxy_server import (
LitellmUserRoles,
audio_transcriptions,
chat_completion,
completion,
embeddings,
image_generation,
model_list,
moderations,
new_end_user,
user_api_key_auth,
)
from litellm.proxy.spend_tracking.spend_management_endpoints import (
global_spend,
global_spend_logs,
global_spend_models,
global_spend_keys,
spend_key_fn,
spend_user_fn,
view_spend_logs,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend
verbose_proxy_logger.setLevel(level=logging.DEBUG)
from starlette.datastructures import URL
from litellm.caching import DualCache
from litellm.proxy._types import (
DynamoDBArgs,
GenerateKeyRequest,
KeyRequest,
LiteLLM_UpperboundKeyGenerateParams,
NewCustomerRequest,
NewTeamRequest,
NewUserRequest,
ProxyErrorTypes,
ProxyException,
UpdateKeyRequest,
UpdateTeamRequest,
UpdateUserRequest,
UserAPIKeyAuth,
)
from litellm.proxy.utils import DBClient
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
@pytest.fixture
def prisma_client():
from litellm.proxy.proxy_cli import append_query_params
### add connection pool + pool timeout args
params = {"connection_limit": 100, "pool_timeout": 60}
database_url = os.getenv("DATABASE_URL")
modified_url = append_query_params(database_url, params)
os.environ["DATABASE_URL"] = modified_url
# Assuming DBClient is a class that needs to be instantiated
prisma_client = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
# Reset litellm.proxy.proxy_server.prisma_client to None
litellm.proxy.proxy_server.custom_db_client = None
litellm.proxy.proxy_server.litellm_proxy_budget_name = (
f"litellm-proxy-budget-{time.time()}"
)
litellm.proxy.proxy_server.user_custom_key_generate = None
return prisma_client
@pytest.mark.asyncio()
async def test_view_daily_spend_ui(prisma_client):
print("prisma client=", prisma_client)
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
from litellm.proxy.proxy_server import user_api_key_cache
spend_logs_for_admin = await global_spend_logs(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234",
user_role=LitellmUserRoles.PROXY_ADMIN,
),
api_key=None,
)
print("spend_logs_for_admin=", spend_logs_for_admin)
spend_logs_for_internal_user = await global_spend_logs(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234"
),
api_key=None,
)
print("spend_logs_for_internal_user=", spend_logs_for_internal_user)
# Calculate total spend for admin
admin_total_spend = sum(log.get("spend", 0) for log in spend_logs_for_admin)
# Calculate total spend for internal user (0 in this case, but we'll keep it generic)
internal_user_total_spend = sum(
log.get("spend", 0) for log in spend_logs_for_internal_user
)
print("total_spend_for_admin=", admin_total_spend)
print("total_spend_for_internal_user=", internal_user_total_spend)
assert (
admin_total_spend > internal_user_total_spend
), "Admin should have more spend than internal user"
@pytest.mark.asyncio
async def test_global_spend_models(prisma_client):
print("prisma client=", prisma_client)
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Test for admin user
models_spend_for_admin = await global_spend_models(
limit=10,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234",
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
print("models_spend_for_admin=", models_spend_for_admin)
# Test for internal user
models_spend_for_internal_user = await global_spend_models(
limit=10,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234"
),
)
print("models_spend_for_internal_user=", models_spend_for_internal_user)
# Assertions
assert isinstance(models_spend_for_admin, list), "Admin response should be a list"
assert isinstance(
models_spend_for_internal_user, list
), "Internal user response should be a list"
# Check if the response has the expected shape for both admin and internal user
expected_keys = ["model", "total_spend"]
if len(models_spend_for_admin) > 0:
assert all(
key in models_spend_for_admin[0] for key in expected_keys
), f"Admin response should contain keys: {expected_keys}"
assert isinstance(
models_spend_for_admin[0]["model"], str
), "Model should be a string"
assert isinstance(
models_spend_for_admin[0]["total_spend"], (int, float)
), "Total spend should be a number"
if len(models_spend_for_internal_user) > 0:
assert all(
key in models_spend_for_internal_user[0] for key in expected_keys
), f"Internal user response should contain keys: {expected_keys}"
assert isinstance(
models_spend_for_internal_user[0]["model"], str
), "Model should be a string"
assert isinstance(
models_spend_for_internal_user[0]["total_spend"], (int, float)
), "Total spend should be a number"
# Check if the lists are sorted by total_spend in descending order
if len(models_spend_for_admin) > 1:
assert all(
models_spend_for_admin[i]["total_spend"]
>= models_spend_for_admin[i + 1]["total_spend"]
for i in range(len(models_spend_for_admin) - 1)
), "Admin response should be sorted by total_spend in descending order"
if len(models_spend_for_internal_user) > 1:
assert all(
models_spend_for_internal_user[i]["total_spend"]
>= models_spend_for_internal_user[i + 1]["total_spend"]
for i in range(len(models_spend_for_internal_user) - 1)
), "Internal user response should be sorted by total_spend in descending order"
# Check if admin has access to more or equal models compared to internal user
assert len(models_spend_for_admin) >= len(
models_spend_for_internal_user
), "Admin should have access to at least as many models as internal user"
# Check if the response contains expected fields
if len(models_spend_for_admin) > 0:
assert all(
key in models_spend_for_admin[0] for key in ["model", "total_spend"]
), "Admin response should contain model, total_spend, and total_tokens"
if len(models_spend_for_internal_user) > 0:
assert all(
key in models_spend_for_internal_user[0] for key in ["model", "total_spend"]
), "Internal user response should contain model, total_spend, and total_tokens"
@pytest.mark.asyncio
async def test_global_spend_keys(prisma_client):
print("prisma client=", prisma_client)
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Test for admin user
keys_spend_for_admin = await global_spend_keys(
limit=10,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234",
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
print("keys_spend_for_admin=", keys_spend_for_admin)
# Test for internal user
keys_spend_for_internal_user = await global_spend_keys(
limit=10,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234"
),
)
print("keys_spend_for_internal_user=", keys_spend_for_internal_user)
# Assertions
assert isinstance(keys_spend_for_admin, list), "Admin response should be a list"
assert isinstance(
keys_spend_for_internal_user, list
), "Internal user response should be a list"
# Check if admin has access to more or equal keys compared to internal user
assert len(keys_spend_for_admin) >= len(
keys_spend_for_internal_user
), "Admin should have access to at least as many keys as internal user"
# Check if the response contains expected fields
if len(keys_spend_for_admin) > 0:
assert all(
key in keys_spend_for_admin[0]
for key in ["api_key", "total_spend", "key_alias", "key_name"]
), "Admin response should contain api_key, total_spend, key_alias, and key_name"
if len(keys_spend_for_internal_user) > 0:
assert all(
key in keys_spend_for_internal_user[0]
for key in ["api_key", "total_spend", "key_alias", "key_name"]
), "Internal user response should contain api_key, total_spend, key_alias, and key_name"

4
ui/litellm-dashboard/src/components/leftnav.tsx
View file

@ -12,6 +12,8 @@ interface SidebarProps {
defaultSelectedKey: string[] | null;
}
const rolesAllowedToSeeUsage = ["Admin", "Admin Viewer", "Internal User", "Internal Viewer"];
const Sidebar: React.FC<SidebarProps> = ({
setPage,
userRole,
@ -62,7 +64,7 @@ const Sidebar: React.FC<SidebarProps> = ({
<Text>Models</Text>
</Menu.Item>
) : null}
{userRole == "Admin" ? (
{rolesAllowedToSeeUsage.includes(userRole) ? (
<Menu.Item key="4" onClick={() => setPage("usage")}>
<Text>Usage</Text>
</Menu.Item>

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

@ -132,6 +132,13 @@ type DataDict = { [key: string]: unknown };
type UserData = { user_id: string; spend: number };
const isAdminOrAdminViewer = (role: string | null): boolean => {
if (role === null) return false;
return role === 'Admin' || role === 'Admin Viewer';
};
const UsagePage: React.FC<UsagePageProps> = ({
accessToken,
token,
@ -256,125 +263,144 @@ const UsagePage: React.FC<UsagePageProps> = ({
const valueFormatter = (number: number) =>
`$ ${new Intl.NumberFormat("us").format(number).toString()}`;
const fetchAndSetData = async (
fetchFunction: () => Promise<any>,
setStateFunction: React.Dispatch<React.SetStateAction<any>>,
errorMessage: string
) => {
try {
const data = await fetchFunction();
setStateFunction(data);
} catch (error) {
console.error(errorMessage, error);
// Optionally, update UI to reflect error state for this specific data
}
};
const fetchOverallSpend = () => fetchAndSetData(
() => accessToken ? adminSpendLogsCall(accessToken) : Promise.reject("No access token"),
setKeySpendData,
"Error fetching overall spend"
);
const fetchProviderSpend = () => fetchAndSetData(
() => accessToken && token ? adminspendByProvider(accessToken, token, startTime, endTime) : Promise.reject("No access token or token"),
setSpendByProvider,
"Error fetching provider spend"
);
const fetchTopKeys = async () => {
if (!accessToken) return;
await fetchAndSetData(
async () => {
const top_keys = await adminTopKeysCall(accessToken);
return top_keys.map((k: any) => ({
key: (k["key_alias"] || k["key_name"] || k["api_key"]).substring(0, 10),
spend: k["total_spend"],
}));
},
setTopKeys,
"Error fetching top keys"
);
};
const fetchTopModels = async () => {
if (!accessToken) return;
await fetchAndSetData(
async () => {
const top_models = await adminTopModelsCall(accessToken);
return top_models.map((k: any) => ({
key: k["model"],
spend: k["total_spend"],
}));
},
setTopModels,
"Error fetching top models"
);
};
const fetchTeamSpend = async () => {
if (!accessToken) return;
await fetchAndSetData(
async () => {
const teamSpend = await teamSpendLogsCall(accessToken);
setTeamSpendData(teamSpend.daily_spend);
setUniqueTeamIds(teamSpend.teams);
return teamSpend.total_spend_per_team.map((tspt: any) => ({
name: tspt["team_id"] || "",
value: (tspt["total_spend"] || 0).toFixed(2),
}));
},
setTotalSpendPerTeam,
"Error fetching team spend"
);
};
const fetchTagNames = () => {
if (!accessToken) return;
fetchAndSetData(
async () => {
const all_tag_names = await allTagNamesCall(accessToken);
return all_tag_names.tag_names;
},
setAllTagNames,
"Error fetching tag names"
);
};
const fetchTopTags = () => {
if (!accessToken) return;
fetchAndSetData(
() => tagsSpendLogsCall(accessToken, dateValue.from?.toISOString(), dateValue.to?.toISOString(), undefined),
(data) => setTopTagsData(data.spend_per_tag),
"Error fetching top tags"
);
};
const fetchTopEndUsers = () => {
if (!accessToken) return;
fetchAndSetData(
() => adminTopEndUsersCall(accessToken, null, undefined, undefined),
setTopUsers,
"Error fetching top end users"
);
};
const fetchGlobalActivity = () => {
if (!accessToken) return;
fetchAndSetData(
() => adminGlobalActivity(accessToken, startTime, endTime),
setGlobalActivity,
"Error fetching global activity"
);
};
const fetchGlobalActivityPerModel = () => {
if (!accessToken) return;
fetchAndSetData(
() => adminGlobalActivityPerModel(accessToken, startTime, endTime),
setGlobalActivityPerModel,
"Error fetching global activity per model"
);
};
useEffect(() => {
if (accessToken && token && userRole && userID) {
const fetchData = async () => {
try {
/**
* If user is Admin - query the global views endpoints
* If user is App Owner - use the normal spend logs call
*/
console.log(`user role: ${userRole}`);
if (userRole == "Admin" || userRole == "Admin Viewer") {
const overall_spend = await adminSpendLogsCall(accessToken);
setKeySpendData(overall_spend);
const provider_spend = await adminspendByProvider(accessToken, token, startTime, endTime);
console.log("provider_spend", provider_spend);
setSpendByProvider(provider_spend);
const top_keys = await adminTopKeysCall(accessToken);
const filtered_keys = top_keys.map((k: any) => ({
key: (k["key_alias"] || k["key_name"] || k["api_key"]).substring(
0,
10
),
spend: k["total_spend"],
}));
setTopKeys(filtered_keys);
const top_models = await adminTopModelsCall(accessToken);
const filtered_models = top_models.map((k: any) => ({
key: k["model"],
spend: k["total_spend"],
}));
setTopModels(filtered_models);
fetchOverallSpend();
fetchProviderSpend();
fetchTopKeys();
fetchTopModels();
fetchGlobalActivity();
fetchGlobalActivityPerModel();
const teamSpend = await teamSpendLogsCall(accessToken);
console.log("teamSpend", teamSpend);
setTeamSpendData(teamSpend.daily_spend);
setUniqueTeamIds(teamSpend.teams)
let total_spend_per_team = teamSpend.total_spend_per_team;
// in total_spend_per_team, replace null team_id with "" and replace null total_spend with 0
total_spend_per_team = total_spend_per_team.map((tspt: any) => {
tspt["name"] = tspt["team_id"] || "";
tspt["value"] = tspt["total_spend"] || 0;
// round the value to 2 decimal places
tspt["value"] = tspt["value"].toFixed(2);
return tspt;
})
setTotalSpendPerTeam(total_spend_per_team);
// all_tag_names -> used for dropdown
const all_tag_names = await allTagNamesCall(accessToken);
setAllTagNames(all_tag_names.tag_names);
//get top tags
const top_tags = await tagsSpendLogsCall(accessToken, dateValue.from?.toISOString(), dateValue.to?.toISOString(), undefined);
setTopTagsData(top_tags.spend_per_tag);
// get spend per end-user
let spend_user_call = await adminTopEndUsersCall(accessToken, null, undefined, undefined);
setTopUsers(spend_user_call);
console.log("spend/user result", spend_user_call);
let global_activity_response = await adminGlobalActivity(accessToken, startTime, endTime);
setGlobalActivity(global_activity_response)
let global_activity_per_model = await adminGlobalActivityPerModel(accessToken, startTime, endTime);
console.log("global activity per model", global_activity_per_model);
setGlobalActivityPerModel(global_activity_per_model)
} else if (userRole == "App Owner") {
await userSpendLogsCall(
accessToken,
token,
userRole,
userID,
startTime,
endTime
).then(async (response) => {
console.log("result from spend logs call", response);
if ("daily_spend" in response) {
// this is from clickhouse analytics
//
let daily_spend = response["daily_spend"];
console.log("daily spend", daily_spend);
setKeySpendData(daily_spend);
let topApiKeys = response.top_api_keys;
setTopKeys(topApiKeys);
} else {
const topKeysResponse = await keyInfoCall(
accessToken,
getTopKeys(response)
);
const filtered_keys = topKeysResponse["info"].map((k: any) => ({
key: (
k["key_name"] ||
k["key_alias"]
).substring(0, 10),
spend: k["spend"],
}));
setTopKeys(filtered_keys);
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
}
};
fetchData();
if (isAdminOrAdminViewer(userRole)) {
fetchTeamSpend();
fetchTagNames();
fetchTopTags();
fetchTopEndUsers();
}
}
}, [accessToken, token, userRole, userID, startTime, endTime]);
@ -385,9 +411,17 @@ const UsagePage: React.FC<UsagePageProps> = ({
<TabGroup>
<TabList className="mt-2">
<Tab>All Up</Tab>
<Tab>Team Based Usage</Tab>
<Tab>Customer Usage</Tab>
<Tab>Tag Based Usage</Tab>
{isAdminOrAdminViewer(userRole) ? (
<>
<Tab>Team Based Usage</Tab>
<Tab>Customer Usage</Tab>
<Tab>Tag Based Usage</Tab>
</>
) : (
<><div></div>
</>
)}
</TabList>
<TabPanels>
<TabPanel>