mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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.
36 lines
964 B
Python
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"')
|