Skip to content

Commit cdcc8ea

Browse files
Connect UI to "LiteLLM_DailyUserSpend" spend table - enables usage tab to work at 1m+ spend logs (#9603)
* feat(spend_management_endpoints.py): expose new endpoint for querying user's usage at 1m+ spend logs Allows user to view their spend at 1m+ spend logs * build(schema.prisma): add api_requests to dailyuserspend table * build(migration.sql): add migration file for new column to daily user spend table * build(prisma_client.py): add logic for copying over migration folder, if deploy/migrations present in expected location enables easier testing of prisma migration flow * build(ui/): initial commit successfully using the dailyuserspend table on the UI * refactor(internal_user_endpoints.py): refactor `/user/daily/activity` to give breakdowns by provider/model/key * feat: feature parity (cost page) with existing 'usage' page * build(ui/): add activity tab to new_usage.tsx gets to feature parity on 'All Up' page of 'usage.tsx' * fix(proxy/utils.py): count number of api requests in daily user spend table allows us to see activity by model on new usage tab * style(new_usage.tsx): fix y-axis to be in ascending order of date * fix: fix linting errors * fix: fix ruff check errors
1 parent b155f4f commit cdcc8ea

File tree

14 files changed

+909
-22
lines changed

14 files changed

+909
-22
lines changed

ci_cd/baseline_db.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import os
21
import subprocess
32
from pathlib import Path
43
from datetime import datetime
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
-- AlterTable
2+
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "api_requests" INTEGER NOT NULL DEFAULT 0;
3+

litellm/proxy/_types.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2718,3 +2718,4 @@ class DailyUserSpendTransaction(TypedDict):
27182718
prompt_tokens: int
27192719
completion_tokens: int
27202720
spend: float
2721+
api_requests: int

litellm/proxy/db/prisma_client.py

Lines changed: 83 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44

55
import asyncio
6+
import glob
67
import os
78
import random
89
import subprocess
@@ -178,6 +179,69 @@ def _create_baseline_migration(schema_path: str) -> bool:
178179
verbose_proxy_logger.warning(f"Error creating baseline migration: {e}")
179180
return False
180181

182+
@staticmethod
183+
def _copy_spend_tracking_migrations(prisma_dir: str) -> bool:
184+
import shutil
185+
from pathlib import Path
186+
187+
"""
188+
Check for and copy over spend tracking migrations if they exist in the deploy directory.
189+
Returns True if migrations were found and copied, False otherwise.
190+
"""
191+
try:
192+
# Get the current file's directory
193+
current_dir = Path(__file__).parent
194+
195+
# Check for migrations in the deploy directory (../../deploy/migrations)
196+
deploy_migrations_dir = (
197+
current_dir.parent.parent.parent / "deploy" / "migrations"
198+
)
199+
200+
# Local migrations directory
201+
local_migrations_dir = Path(prisma_dir + "/migrations")
202+
203+
if deploy_migrations_dir.exists():
204+
# Create local migrations directory if it doesn't exist
205+
local_migrations_dir.mkdir(parents=True, exist_ok=True)
206+
207+
# Copy all migration files
208+
# Copy entire migrations folder recursively
209+
shutil.copytree(
210+
deploy_migrations_dir, local_migrations_dir, dirs_exist_ok=True
211+
)
212+
213+
return True
214+
return False
215+
except Exception:
216+
return False
217+
218+
@staticmethod
219+
def _get_migration_names(migrations_dir: str) -> list:
220+
"""Get all migration directory names from the migrations folder"""
221+
migration_paths = glob.glob(f"{migrations_dir}/*/migration.sql")
222+
return [Path(p).parent.name for p in migration_paths]
223+
224+
@staticmethod
225+
def _resolve_all_migrations(migrations_dir: str):
226+
"""Mark all existing migrations as applied"""
227+
migration_names = PrismaManager._get_migration_names(migrations_dir)
228+
for migration_name in migration_names:
229+
try:
230+
verbose_proxy_logger.info(f"Resolving migration: {migration_name}")
231+
subprocess.run(
232+
["prisma", "migrate", "resolve", "--applied", migration_name],
233+
timeout=60,
234+
check=True,
235+
capture_output=True,
236+
text=True,
237+
)
238+
verbose_proxy_logger.debug(f"Resolved migration: {migration_name}")
239+
except subprocess.CalledProcessError as e:
240+
if "is already recorded as applied in the database." not in e.stderr:
241+
verbose_proxy_logger.warning(
242+
f"Failed to resolve migration {migration_name}: {e.stderr}"
243+
)
244+
181245
@staticmethod
182246
def setup_database(use_migrate: bool = False) -> bool:
183247
"""
@@ -194,8 +258,10 @@ def setup_database(use_migrate: bool = False) -> bool:
194258
os.chdir(prisma_dir)
195259
try:
196260
if use_migrate:
261+
PrismaManager._copy_spend_tracking_migrations(
262+
prisma_dir
263+
) # place a migration in the migrations directory
197264
verbose_proxy_logger.info("Running prisma migrate deploy")
198-
# First try to run migrate deploy directly
199265
try:
200266
subprocess.run(
201267
["prisma", "migrate", "deploy"],
@@ -205,25 +271,31 @@ def setup_database(use_migrate: bool = False) -> bool:
205271
text=True,
206272
)
207273
verbose_proxy_logger.info("prisma migrate deploy completed")
274+
275+
# Resolve all migrations in the migrations directory
276+
migrations_dir = os.path.join(prisma_dir, "migrations")
277+
PrismaManager._resolve_all_migrations(migrations_dir)
278+
208279
return True
209280
except subprocess.CalledProcessError as e:
210-
# Check if this is the non-empty schema error
281+
verbose_proxy_logger.warning(
282+
f"prisma db error: {e.stderr}, e: {e.stdout}"
283+
)
211284
if (
212285
"P3005" in e.stderr
213286
and "database schema is not empty" in e.stderr
214287
):
215-
# Create baseline migration
288+
verbose_proxy_logger.info("Creating baseline migration")
216289
if PrismaManager._create_baseline_migration(schema_path):
217-
# Try migrate deploy again after baseline
218-
subprocess.run(
219-
["prisma", "migrate", "deploy"],
220-
timeout=60,
221-
check=True,
290+
verbose_proxy_logger.info(
291+
"Resolving all migrations after baseline"
222292
)
293+
294+
# Resolve all migrations after baseline
295+
migrations_dir = os.path.join(prisma_dir, "migrations")
296+
PrismaManager._resolve_all_migrations(migrations_dir)
297+
223298
return True
224-
else:
225-
# If it's a different error, raise it
226-
raise e
227299
else:
228300
# Use prisma db push with increased timeout
229301
subprocess.run(

0 commit comments

Comments
 (0)