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

Refactor 'check_view_exists' logic (#5659)

Browse source 
* fix(proxy/utils.py): comment out auto-upsert logic in check_view_exists

Prevents proxy from failing on startup due to faulty logic

* fix(db/migration_scripts/create_views.py): fix 'DailyTagSpend' quotation on check

* fix(create_views.py): mongly global spend time period should be 30d not 20d

* fix(schema.prisma): index on startTime and endUser for efficient UI querying
This commit is contained in:
Krish Dholakia 2024-09-12 13:39:50 -07:00 • committed by GitHub
parent 5c1a70be21
commit c76d2c6ade
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 202 additions and 209 deletions
Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/db/create_views.py → db_scripts/create_views.py
View file

@ -1,7 +1,5 @@
""" """
python script to pre-create all views required by LiteLLM Proxy Server python script to pre-create all views required by LiteLLM Proxy Server
""" """
import asyncio import asyncio
@ -27,30 +25,6 @@ async def check_view_exists():
# connect to dB # connect to dB
await db.connect() await db.connect()
# Check to see if all of the necessary views exist and if they do, simply return
# This is more efficient because it lets us check for all views in one
# query instead of multiple queries.
try:
ret = await db.query_raw(
"""
SELECT SUM(1) FROM pg_views
WHERE schemaname = 'public' AND viewname IN (
'LiteLLM_VerificationTokenView',
'MonthlyGlobalSpend',
'Last30dKeysBySpend',
'Last30dModelsBySpend',
'MonthlyGlobalSpendPerKey',
'MonthlyGlobalSpendPerUserPerKey',
'Last30dTopEndUsersSpend'
)
"""
)
if ret[0]["sum"] == 8:
print("All necessary views exist!") # noqa
return
except Exception:
pass
try: try:
# Try to select one row from the view # Try to select one row from the view
await db.query_raw("""SELECT 1 FROM "LiteLLM_VerificationTokenView" LIMIT 1""") await db.query_raw("""SELECT 1 FROM "LiteLLM_VerificationTokenView" LIMIT 1""")
@ -180,7 +154,7 @@ async def check_view_exists():
FROM FROM
"LiteLLM_SpendLogs" "LiteLLM_SpendLogs"
WHERE WHERE
"startTime" >= (CURRENT_DATE - INTERVAL '20 days') "startTime" >= (CURRENT_DATE - INTERVAL '30 days')
GROUP BY GROUP BY
DATE("startTime"), DATE("startTime"),
"user", "user",
@ -191,7 +165,7 @@ async def check_view_exists():
print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa
try: try:
await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""") await db.query_raw("""SELECT 1 FROM DailyTagSpend LIMIT 1""")
print("DailyTagSpend Exists!") # noqa print("DailyTagSpend Exists!") # noqa
except Exception as e: except Exception as e:
sql_query = """ sql_query = """

4
litellm/proxy/db/dynamo_db.py
View file

@ -1,3 +1,7 @@
"""
Deprecated. Only PostgresSQL is supported.
"""
import json import json
from datetime import datetime from datetime import datetime
from typing import Any, List, Literal, Optional, Union from typing import Any, List, Literal, Optional, Union

4
litellm/proxy/schema.prisma
View file

@ -168,7 +168,7 @@ model LiteLLM_Config {
param_value Json? param_value Json?
} }
// View spend, model, hashed api_key per request // View spend, model, api_key per request
model LiteLLM_SpendLogs { model LiteLLM_SpendLogs {
request_id String @id request_id String @id
call_type String call_type String
@ -192,6 +192,8 @@ model LiteLLM_SpendLogs {
team_id String? team_id String?
end_user String? end_user String?
requester_ip_address String? requester_ip_address String?
@@index([startTime])
@@index([end_user])
} }
// View spend, model, api_key per request // View spend, model, api_key per request

371
litellm/proxy/utils.py
View file

@ -960,207 +960,218 @@ class PrismaClient:
# This is more efficient because it lets us check for all views in one # This is more efficient because it lets us check for all views in one
# query instead of multiple queries. # query instead of multiple queries.
try: try:
expected_views = [
"LiteLLM_VerificationTokenView",
"MonthlyGlobalSpend",
"Last30dKeysBySpend",
"Last30dModelsBySpend",
"MonthlyGlobalSpendPerKey",
"MonthlyGlobalSpendPerUserPerKey",
"Last30dTopEndUsersSpend",
]
expected_views_str = ", ".join(f"'{view}'" for view in expected_views)
ret = await self.db.query_raw( ret = await self.db.query_raw(
""" f"""
SELECT SUM(1) FROM pg_views SELECT SUM(1) FROM pg_views
WHERE schemaname = 'public' AND viewname IN ( WHERE schemaname = 'public' AND viewname IN (
'LiteLLM_VerificationTokenView', {expected_views_str}
'MonthlyGlobalSpend',
'Last30dKeysBySpend',
'Last30dModelsBySpend',
'MonthlyGlobalSpendPerKey',
'MonthlyGlobalSpendPerUserPerKey',
'Last30dTopEndUsersSpend'
) )
""" """
) )
if ret[0]["sum"] == 8: expected_total_views = len(expected_views)
print("All necessary views exist!") # noqa if ret[0]["sum"] == expected_total_views:
verbose_proxy_logger.info("All necessary views exist!")
return return
except Exception: else:
pass raise Exception(
"Not all views exist in db. Got={}. Expected={}. Run 'create_views.py' in litellm/db_scripts to create missing views.".format(
try: ret[0]["sum"], expected_total_views
# Try to select one row from the view )
await self.db.query_raw( )
"""SELECT 1 FROM "LiteLLM_VerificationTokenView" LIMIT 1"""
)
print("LiteLLM_VerificationTokenView Exists!") # noqa
except Exception as e: except Exception as e:
# If an error occurs, the view does not exist, so create it raise e
value = await self.health_check()
await self.db.execute_raw(
"""
CREATE VIEW "LiteLLM_VerificationTokenView" AS
SELECT
v.*,
t.spend AS team_spend,
t.max_budget AS team_max_budget,
t.tpm_limit AS team_tpm_limit,
t.rpm_limit AS team_rpm_limit
FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
"""
)
print("LiteLLM_VerificationTokenView Created!") # noqa # try:
# # Try to select one row from the view
# await self.db.query_raw(
# """SELECT 1 FROM "LiteLLM_VerificationTokenView" LIMIT 1"""
# )
# print("LiteLLM_VerificationTokenView Exists!") # noqa
# except Exception as e:
# # If an error occurs, the view does not exist, so create it
# value = await self.health_check()
# await self.db.execute_raw(
# """
# CREATE VIEW "LiteLLM_VerificationTokenView" AS
# SELECT
# v.*,
# t.spend AS team_spend,
# t.max_budget AS team_max_budget,
# t.tpm_limit AS team_tpm_limit,
# t.rpm_limit AS team_rpm_limit
# FROM "LiteLLM_VerificationToken" v
# LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
# """
# )
try: # print("LiteLLM_VerificationTokenView Created!") # noqa
await self.db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""")
print("MonthlyGlobalSpend Exists!") # noqa
except Exception as e:
sql_query = """
CREATE OR REPLACE VIEW "MonthlyGlobalSpend" AS
SELECT
DATE("startTime") AS date,
SUM("spend") AS spend
FROM
"LiteLLM_SpendLogs"
WHERE
"startTime" >= (CURRENT_DATE - INTERVAL '30 days')
GROUP BY
DATE("startTime");
"""
await self.db.execute_raw(query=sql_query)
print("MonthlyGlobalSpend Created!") # noqa # try:
# await self.db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""")
# print("MonthlyGlobalSpend Exists!") # noqa
# except Exception as e:
# sql_query = """
# CREATE OR REPLACE VIEW "MonthlyGlobalSpend" AS
# SELECT
# DATE("startTime") AS date,
# SUM("spend") AS spend
# FROM
# "LiteLLM_SpendLogs"
# WHERE
# "startTime" >= (CURRENT_DATE - INTERVAL '30 days')
# GROUP BY
# DATE("startTime");
# """
# await self.db.execute_raw(query=sql_query)
try: # print("MonthlyGlobalSpend Created!") # noqa
await self.db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""")
print("Last30dKeysBySpend Exists!") # noqa
except Exception as e:
sql_query = """
CREATE OR REPLACE VIEW "Last30dKeysBySpend" AS
SELECT
L."api_key",
V."key_alias",
V."key_name",
SUM(L."spend") AS total_spend
FROM
"LiteLLM_SpendLogs" L
LEFT JOIN
"LiteLLM_VerificationToken" V
ON
L."api_key" = V."token"
WHERE
L."startTime" >= (CURRENT_DATE - INTERVAL '30 days')
GROUP BY
L."api_key", V."key_alias", V."key_name"
ORDER BY
total_spend DESC;
"""
await self.db.execute_raw(query=sql_query)
print("Last30dKeysBySpend Created!") # noqa # try:
# await self.db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""")
# print("Last30dKeysBySpend Exists!") # noqa
# except Exception as e:
# sql_query = """
# CREATE OR REPLACE VIEW "Last30dKeysBySpend" AS
# SELECT
# L."api_key",
# V."key_alias",
# V."key_name",
# SUM(L."spend") AS total_spend
# FROM
# "LiteLLM_SpendLogs" L
# LEFT JOIN
# "LiteLLM_VerificationToken" V
# ON
# L."api_key" = V."token"
# WHERE
# L."startTime" >= (CURRENT_DATE - INTERVAL '30 days')
# GROUP BY
# L."api_key", V."key_alias", V."key_name"
# ORDER BY
# total_spend DESC;
# """
# await self.db.execute_raw(query=sql_query)
try: # print("Last30dKeysBySpend Created!") # noqa
await self.db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""")
print("Last30dModelsBySpend Exists!") # noqa
except Exception as e:
sql_query = """
CREATE OR REPLACE VIEW "Last30dModelsBySpend" AS
SELECT
"model",
SUM("spend") AS total_spend
FROM
"LiteLLM_SpendLogs"
WHERE
"startTime" >= (CURRENT_DATE - INTERVAL '30 days')
AND "model" != ''
GROUP BY
"model"
ORDER BY
total_spend DESC;
"""
await self.db.execute_raw(query=sql_query)
print("Last30dModelsBySpend Created!") # noqa # try:
try: # await self.db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""")
await self.db.query_raw( # print("Last30dModelsBySpend Exists!") # noqa
"""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""" # except Exception as e:
) # sql_query = """
print("MonthlyGlobalSpendPerKey Exists!") # noqa # CREATE OR REPLACE VIEW "Last30dModelsBySpend" AS
except Exception as e: # SELECT
sql_query = """ # "model",
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerKey" AS # SUM("spend") AS total_spend
SELECT # FROM
DATE("startTime") AS date, # "LiteLLM_SpendLogs"
SUM("spend") AS spend, # WHERE
api_key as api_key # "startTime" >= (CURRENT_DATE - INTERVAL '30 days')
FROM # AND "model" != ''
"LiteLLM_SpendLogs" # GROUP BY
WHERE # "model"
"startTime" >= (CURRENT_DATE - INTERVAL '30 days') # ORDER BY
GROUP BY # total_spend DESC;
DATE("startTime"), # """
api_key; # await self.db.execute_raw(query=sql_query)
"""
await self.db.execute_raw(query=sql_query)
print("MonthlyGlobalSpendPerKey Created!") # noqa # print("Last30dModelsBySpend Created!") # noqa
try: # try:
await self.db.query_raw( # await self.db.query_raw(
"""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""" # """SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1"""
) # )
print("MonthlyGlobalSpendPerUserPerKey Exists!") # noqa # print("MonthlyGlobalSpendPerKey Exists!") # noqa
except Exception as e: # except Exception as e:
sql_query = """ # sql_query = """
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerUserPerKey" AS # CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerKey" AS
SELECT # SELECT
DATE("startTime") AS date, # DATE("startTime") AS date,
SUM("spend") AS spend, # SUM("spend") AS spend,
api_key as api_key, # api_key as api_key
"user" as "user" # FROM
FROM # "LiteLLM_SpendLogs"
"LiteLLM_SpendLogs" # WHERE
WHERE # "startTime" >= (CURRENT_DATE - INTERVAL '30 days')
"startTime" >= (CURRENT_DATE - INTERVAL '20 days') # GROUP BY
GROUP BY # DATE("startTime"),
DATE("startTime"), # api_key;
"user", # """
api_key; # await self.db.execute_raw(query=sql_query)
"""
await self.db.execute_raw(query=sql_query)
print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa # 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)
try: # print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa
await self.db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""")
print("DailyTagSpend Exists!") # noqa
except Exception as e:
sql_query = """
CREATE OR REPLACE VIEW DailyTagSpend AS
SELECT
jsonb_array_elements_text(request_tags) AS individual_request_tag,
DATE(s."startTime") AS spend_date,
COUNT(*) AS log_count,
SUM(spend) AS total_spend
FROM "LiteLLM_SpendLogs" s
GROUP BY individual_request_tag, DATE(s."startTime");
"""
await self.db.execute_raw(query=sql_query)
print("DailyTagSpend Created!") # noqa # try:
# await self.db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""")
# print("DailyTagSpend Exists!") # noqa
# except Exception as e:
# sql_query = """
# CREATE OR REPLACE VIEW DailyTagSpend AS
# SELECT
# jsonb_array_elements_text(request_tags) AS individual_request_tag,
# DATE(s."startTime") AS spend_date,
# COUNT(*) AS log_count,
# SUM(spend) AS total_spend
# FROM "LiteLLM_SpendLogs" s
# GROUP BY individual_request_tag, DATE(s."startTime");
# """
# await self.db.execute_raw(query=sql_query)
try: # print("DailyTagSpend Created!") # noqa
await self.db.query_raw(
"""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1"""
)
print("Last30dTopEndUsersSpend Exists!") # noqa
except Exception as e:
sql_query = """
CREATE VIEW "Last30dTopEndUsersSpend" AS
SELECT end_user, COUNT(*) AS total_events, SUM(spend) AS total_spend
FROM "LiteLLM_SpendLogs"
WHERE end_user <> '' AND end_user <> user
AND "startTime" >= CURRENT_DATE - INTERVAL '30 days'
GROUP BY end_user
ORDER BY total_spend DESC
LIMIT 100;
"""
await self.db.execute_raw(query=sql_query)
print("Last30dTopEndUsersSpend Created!") # noqa # try:
# await self.db.query_raw(
# """SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1"""
# )
# print("Last30dTopEndUsersSpend Exists!") # noqa
# except Exception as e:
# sql_query = """
# CREATE VIEW "Last30dTopEndUsersSpend" AS
# SELECT end_user, COUNT(*) AS total_events, SUM(spend) AS total_spend
# FROM "LiteLLM_SpendLogs"
# WHERE end_user <> '' AND end_user <> user
# AND "startTime" >= CURRENT_DATE - INTERVAL '30 days'
# GROUP BY end_user
# ORDER BY total_spend DESC
# LIMIT 100;
# """
# await self.db.execute_raw(query=sql_query)
# print("Last30dTopEndUsersSpend Created!") # noqa
return return

2
schema.prisma
View file

@ -192,6 +192,8 @@ model LiteLLM_SpendLogs {
team_id String? team_id String?
end_user String? end_user String?
requester_ip_address String? requester_ip_address String?
@@index([startTime])
@@index([end_user])
} }
// View spend, model, api_key per request // View spend, model, api_key per request