litellm/tests/proxy_behavior/spend/conftest.py
Tin Chi Lo 2e2b2624ed feat(spend): add a per-session auto-router benchmarks rollup
Fold each auto-routed turn into a per-(api_key, session, auto-router) row when it
happens, and serve the benchmarks dashboard by summing those rows. Nothing in the
feature reads LiteLLM_SpendLogs.

The row carries what each tier the session used left in the prompt cache, so a turn's
bucket is a question about one model's own record and the upsert answers it against the
row it is already writing. Absent from that record means a first visit; present means
warm or expired, on one comparison of the idle gap against the TTL the entry was written
with. Nothing is read before the write, so there is no state to load or validate, and the
statement is atomic.
2026-08-04 18:50:07 -07:00

36 lines
964 B
Python

"""A real Postgres connection for the rollup's SQL.
The auto-router rollup classifies each turn inside its upsert, against the row's own
stored state, so the classification only exists when a real database evaluates it.
"""
from dataclasses import dataclass
import pytest_asyncio
from prisma import Prisma
@dataclass(frozen=True)
class PrismaClientShim:
"""What the rollup writer needs from litellm's PrismaClient: a connected `db`."""
db: Prisma
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def prisma_db():
db = Prisma()
await db.connect()
try:
yield db
finally:
await db.disconnect()
@pytest_asyncio.fixture(loop_scope="session")
async def rollup_client(prisma_db):
await prisma_db.execute_raw('DELETE FROM "LiteLLM_AutoRouterSession"')
try:
yield PrismaClientShim(db=prisma_db)
finally:
await prisma_db.execute_raw('DELETE FROM "LiteLLM_AutoRouterSession"')