scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 151

perf(proxy): index customer budget ownership lookups

Browse source
This commit is contained in:
lewis 2026-10-17 21:12:01 -04:00
parent 8408cd94b5
commit e664274d6a
9 changed files with 104 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000100_customer_budget_reference_index/migration.sql Normal file
View file

@ -0,0 +1 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_EndUserTable_budget_id_idx" ON "LiteLLM_EndUserTable" ("budget_id");

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000200_key_budget_reference_index/migration.sql Normal file
View file

@ -0,0 +1 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_VerificationToken_budget_id_idx" ON "LiteLLM_VerificationToken" ("budget_id");

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000300_key_customer_budget_reference_index/migration.sql Normal file
View file

@ -0,0 +1 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_VerificationToken_end_user_budget_id_idx" ON "LiteLLM_VerificationToken" USING GIN (("metadata"->'end_user_budget_id') jsonb_path_ops);

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -563,6 +563,7 @@ model LiteLLM_VerificationToken {
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
@@index([budget_reset_at, expires])
@@index([budget_id])
}
model LiteLLM_JWTKeyMapping {
@ -678,6 +679,7 @@ model LiteLLM_EndUserTable {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
@@index([budget_id])
}
// Budget and shared spend for a model access group. The groups themselves are not rows anywhere:

3
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -76,7 +76,8 @@ _CUSTOMER_SHARED_BUDGET_SQL: Final = """
SELECT "budget_id" FROM "LiteLLM_BudgetTable" WHERE "budget_id" = $1 AND EXISTS (
SELECT 1 FROM "LiteLLM_EndUserTable" WHERE "budget_id" = $1 AND "user_id" <> $2
UNION ALL SELECT 1 FROM "LiteLLM_VerificationToken" WHERE "budget_id" = $1
OR "metadata"->>'end_user_budget_id' = $1
OR (("metadata"->'end_user_budget_id') @> to_jsonb($1::text)
AND jsonb_typeof("metadata"->'end_user_budget_id') = 'string')
UNION ALL SELECT 1 FROM "LiteLLM_OrganizationTable" WHERE "budget_id" = $1
UNION ALL SELECT 1 FROM "LiteLLM_ProjectTable" WHERE "budget_id" = $1
UNION ALL SELECT 1 FROM "LiteLLM_TagTable" WHERE "budget_id" = $1

2
litellm/proxy/schema.prisma
View file

@ -563,6 +563,7 @@ model LiteLLM_VerificationToken {
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
@@index([budget_reset_at, expires])
@@index([budget_id])
}
model LiteLLM_JWTKeyMapping {
@ -678,6 +679,7 @@ model LiteLLM_EndUserTable {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
@@index([budget_id])
}
// Budget and shared spend for a model access group. The groups themselves are not rows anywhere:

2
schema.prisma
View file

@ -563,6 +563,7 @@ model LiteLLM_VerificationToken {
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
@@index([budget_reset_at, expires])
@@index([budget_id])
}
model LiteLLM_JWTKeyMapping {
@ -678,6 +679,7 @@ model LiteLLM_EndUserTable {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
@@index([budget_id])
}
// Budget and shared spend for a model access group. The groups themselves are not rows anywhere:

89
tests/integration/database/migrations/test_customer_budget_reference_lookup.py Normal file
View file

@ -0,0 +1,89 @@
import os
from pathlib import Path
from typing import Final
import psycopg
from psycopg import sql
from pydantic import BaseModel, Field, TypeAdapter
from litellm.proxy.management_endpoints.customer_endpoints import _CUSTOMER_SHARED_BUDGET_SQL
_REPOSITORY: Final = Path(__file__).resolve().parents[4]
_MIGRATIONS: Final = _REPOSITORY / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
_INDEX_MIGRATIONS: Final = (
"20261018000100_customer_budget_reference_index",
"20261018000200_key_budget_reference_index",
"20261018000300_key_customer_budget_reference_index",
)
_REFERENCE_TABLES: Final = (
"LiteLLM_OrganizationTable",
"LiteLLM_ProjectTable",
"LiteLLM_TagTable",
"LiteLLM_ModelAccessGroupBudgetTable",
"LiteLLM_TeamMembership",
"LiteLLM_OrganizationMembership",
"LiteLLM_TeamTable",
)
class _PlanNode(BaseModel):
rows: float = Field(default=0, alias="Actual Rows")
removed: float = Field(default=0, alias="Rows Removed by Filter")
loops: float = Field(default=0, alias="Actual Loops")
plans: tuple["_PlanNode", ...] = Field(default=(), alias="Plans")
def examined_rows(self) -> float:
return (self.rows + self.removed) * self.loops + sum(plan.examined_rows() for plan in self.plans)
class _QueryPlan(BaseModel):
plan: _PlanNode = Field(alias="Plan")
def test_customer_budget_guard_keeps_generic_plan_work_bounded() -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute('CREATE TABLE "LiteLLM_BudgetTable" (budget_id text PRIMARY KEY)')
connection.execute('CREATE TABLE "LiteLLM_EndUserTable" (user_id text PRIMARY KEY, budget_id text)')
connection.execute(
'CREATE TABLE "LiteLLM_VerificationToken" (token text PRIMARY KEY, budget_id text, metadata jsonb)'
)
for table in _REFERENCE_TABLES:
connection.execute(
sql.SQL("CREATE TABLE {} (budget_id text, metadata jsonb)").format(sql.Identifier(table))
)
connection.execute("INSERT INTO \"LiteLLM_BudgetTable\" SELECT 'budget-'||n FROM generate_series(0,500) n")
connection.execute(
"INSERT INTO \"LiteLLM_EndUserTable\" SELECT 'customer-'||n,'budget-'||n FROM generate_series(0,500) n"
)
key_count: Final = 50_000
connection.execute(
"INSERT INTO \"LiteLLM_VerificationToken\" SELECT 'key-'||n,NULL,"
"jsonb_build_object('end_user_budget_id','unrelated-budget') FROM generate_series(1,%s) n",
(key_count,),
)
for migration in _INDEX_MIGRATIONS:
connection.execute((_MIGRATIONS / migration / "migration.sql").read_text())
connection.execute("ANALYZE")
connection.execute("SET plan_cache_mode = force_generic_plan")
connection.execute("PREPARE guard(text,text) AS " + _CUSTOMER_SHARED_BUDGET_SQL)
plan_row: Final = connection.execute(
"EXPLAIN (ANALYZE, FORMAT JSON) EXECUTE guard('budget-0','customer-0')"
).fetchone()
assert plan_row is not None
plan: Final = TypeAdapter(tuple[_QueryPlan, ...]).validate_python(plan_row[0])[0].plan
assert plan.examined_rows() < key_count / 100
assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == []
for metadata in ('{"end_user_budget_id": ["budget-0"]}', '{"end_user_budget_id": null}'):
connection.execute(
"INSERT INTO \"LiteLLM_VerificationToken\" VALUES ('malformed',NULL,%s::jsonb)", (metadata,)
)
assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == []
connection.execute("DELETE FROM \"LiteLLM_VerificationToken\" WHERE token='malformed'")
connection.execute(
'INSERT INTO "LiteLLM_VerificationToken" VALUES '
"('metadata-budget',NULL,jsonb_build_object('end_user_budget_id','budget-0'))"
)
assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == [("budget-0",)]
connection.execute("DELETE FROM \"LiteLLM_VerificationToken\" WHERE token='metadata-budget'")
connection.execute("INSERT INTO \"LiteLLM_VerificationToken\" VALUES ('direct-budget','budget-0','{}')")
assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == [("budget-0",)]

5
tests/unit/proxy/management_endpoints/test_customer_endpoints.py
View file

@ -1318,7 +1318,10 @@ class _CustomerBatchDatabase:
[
dict(row)
for row in database.execute(
query, {str(index): value for index, value in enumerate(parameters, 1)}
query.replace("@> to_jsonb($1::text)", "= json_quote($1)")
.replace("jsonb_typeof", "json_type")
.replace("= 'string'", "= 'text'"),
{str(index): value for index, value in enumerate(parameters, 1)},
).fetchall()
]
)