mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(proxy): bound entity activity summaries
This commit is contained in:
parent
cee0094d58
commit
16f686ca36
8 changed files with 397 additions and 9 deletions
23
docs/features/bounded-activity.md
Normal file
23
docs/features/bounded-activity.md
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
# Bounded Activity summaries
|
||||
|
||||
Entity Activity tabs use a bounded summary of existing daily activity records. The summary preserves stored daily and range totals while limiting the model breakdown returned to the browser.
|
||||
|
||||
Use `bounded=true` on `GET /user/daily/activity/aggregated` or `GET /team/daily/activity/aggregated`. Existing callers that omit this parameter keep the original response and behavior. Existing user, team, API-key and model filters and permissions still apply; ordinary team members retain their existing self-only view.
|
||||
|
||||
Provide `start_date` and `end_date` in `YYYY-MM-DD` format. The inclusive range may contain at most 366 days. The response contains:
|
||||
|
||||
- `results`: exact metrics for each recorded day, without per-day dimension dictionaries.
|
||||
- `metadata`: exact totals across the full matching range.
|
||||
- `model_breakdown`: null unless `include_model_breakdown=true`.
|
||||
|
||||
With a model breakdown, `model_limit` selects the highest-spend models across the full range. It defaults to 20 and accepts 1–100. `total_models` counts all matching model identities. The separate `other` object contains the combined metrics and count of every remaining model. It is null when no models remain and is never used as a model identity. A real model named `Other` remains a distinct model. Identifiers longer than 1,024 characters are included in the `other` aggregate, preserving their totals without truncating an identity or returning an oversized name.
|
||||
|
||||
For example, request a key's bounded summary with:
|
||||
|
||||
```text
|
||||
GET /user/daily/activity/aggregated?bounded=true&include_model_breakdown=true&model_limit=20&api_key=example-reporting-key-id&start_date=2026-11-01&end_date=2026-11-04
|
||||
```
|
||||
|
||||
Use a reporting key ID, not a secret API key. The model contributions plus `other` sum to the complete range totals. Derived PTU-hours and guardrail or model-group breakdowns are not part of this summary; use the existing Usage view for those breakdowns.
|
||||
|
||||
The SQL query limits result rows before the database driver materializes them. At the maximum range and model limit it returns at most 468 rows; without models it returns at most 367. Queries also have bounded concurrent admission, statement time, lock wait and working memory. A busy or over-budget request returns `503` with `Retry-After: 5`. Retry later or narrow the date range. These limits bound the summary response and query resources; query cost still depends on the number of matching stored activity rows.
|
||||
|
|
@ -29,7 +29,7 @@ from litellm.repositories.daily_activity_buckets import (
|
|||
model_group_bucket,
|
||||
raw_api_key_filter,
|
||||
)
|
||||
from litellm.repositories.daily_activity_repository import DailyActivityRepository
|
||||
from litellm.repositories.daily_activity_repository import DailyActivityQueryDatabase, DailyActivityRepository
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
BreakdownMetrics,
|
||||
DailySpendData,
|
||||
|
|
@ -572,8 +572,12 @@ class _ProxyDailyActivityReads:
|
|||
)
|
||||
|
||||
|
||||
def daily_activity_repository(prisma_client: PrismaClient) -> DailyActivityRepository:
|
||||
return DailyActivityRepository(prisma_client, proxy_reads=_ProxyDailyActivityReads(prisma_client))
|
||||
def daily_activity_repository(
|
||||
prisma_client: PrismaClient, *, query_database: DailyActivityQueryDatabase | None = None
|
||||
) -> DailyActivityRepository:
|
||||
return DailyActivityRepository(
|
||||
prisma_client, proxy_reads=_ProxyDailyActivityReads(prisma_client), query_database=query_database
|
||||
)
|
||||
|
||||
|
||||
def daily_activity_scope(
|
||||
|
|
|
|||
|
|
@ -39,10 +39,13 @@ from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, w
|
|||
from litellm.proxy.management_endpoints.team_endpoints import aggregated_date_range_error
|
||||
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||||
from litellm.proxy.utils import PrismaClient, get_prisma_client_or_throw
|
||||
from litellm.repositories.analytics_budget import AnalyticsCapacityError
|
||||
from litellm.repositories.bounded_daily_activity import BoundedDailyActivityRepository
|
||||
from litellm.repositories.daily_activity_buckets import api_key_bucket
|
||||
from litellm.repositories.daily_activity_repository import DailyActivityRepository
|
||||
from litellm.router import Router
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
BoundedDailyActivityResponse,
|
||||
CacheLeakageKeysResponse,
|
||||
DailyActivityKeyPageResponse,
|
||||
DailyActivityKeySearchResponse,
|
||||
|
|
@ -268,8 +271,13 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
|
|||
api_key_limit: Annotated[
|
||||
int, Query(ge=1, le=constants.USAGE_TOP_API_KEYS_MAX)
|
||||
] = constants.USAGE_TOP_API_KEYS_DEFAULT,
|
||||
) -> SpendAnalyticsPaginatedResponse:
|
||||
bounded: bool = False,
|
||||
include_model_breakdown: bool = False,
|
||||
model_limit: Annotated[int, Query(ge=1, le=100)] = 20,
|
||||
) -> SpendAnalyticsPaginatedResponse | BoundedDailyActivityResponse:
|
||||
try:
|
||||
if bounded and resolver.entity not in ("user", "team"):
|
||||
raise HTTPException(status_code=422, detail="Bounded Activity is supported for user and team scopes")
|
||||
resolved: ResolvedScope = await _resolved_scope(
|
||||
resolver,
|
||||
entity_query,
|
||||
|
|
@ -277,6 +285,13 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
|
|||
prisma_client,
|
||||
user_aggregated=resolver.entity == "user",
|
||||
)
|
||||
if bounded:
|
||||
try:
|
||||
return await BoundedDailyActivityRepository.from_prisma(prisma_client).aggregate(
|
||||
resolved.scope, model_limit=model_limit if include_model_breakdown else None
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
activity: Final = await get_daily_activity_aggregated(
|
||||
repository,
|
||||
resolved.scope,
|
||||
|
|
@ -290,6 +305,8 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
|
|||
return with_ptu_consumption(activity, llm_router, single_team_id(resolved.scope.entity_ids))
|
||||
except HTTPException:
|
||||
raise
|
||||
except AnalyticsCapacityError as exc:
|
||||
raise HTTPException(status_code=503, detail=str(exc), headers={"Retry-After": "5"}) from exc
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.exception("Daily activity aggregation failed: %s", exc)
|
||||
raise HTTPException(status_code=500, detail={"error": f"Failed to fetch analytics: {exc}"})
|
||||
|
|
@ -301,7 +318,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
|
|||
name=resolver.operation_names["aggregated"],
|
||||
tags=list(resolver.tags),
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=SpendAnalyticsPaginatedResponse,
|
||||
response_model=SpendAnalyticsPaginatedResponse | BoundedDailyActivityResponse,
|
||||
include_in_schema=prefix != "/end_user",
|
||||
)
|
||||
|
||||
|
|
|
|||
99
litellm/repositories/bounded_daily_activity.py
Normal file
99
litellm/repositories/bounded_daily_activity.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
"""Bounded daily totals and optional model contributions for Activity panels."""
|
||||
|
||||
from datetime import date
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast # noqa: TID251 # generated Prisma transaction interface
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm.repositories.analytics_budget import AnalyticsDatabase, BoundedAnalyticsDatabase
|
||||
from litellm.repositories.daily_activity_repository import DailyActivityQueryDatabase
|
||||
from litellm.repositories.daily_activity_sql import build_bounded_daily_activity_sql
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
BoundedDailyActivityResponse,
|
||||
DailyActivityModelBreakdown,
|
||||
DailyActivityModelContribution,
|
||||
DailyActivityOtherModels,
|
||||
DailySpendData,
|
||||
DailySpendMetadata,
|
||||
SpendMetrics,
|
||||
)
|
||||
from litellm.types.repositories.daily_activity import DailyActivityScope
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class _ActivityRow(BaseModel, frozen=True):
|
||||
kind: Literal["day", "total", "model", "other"]
|
||||
date: date | None
|
||||
model: str | None
|
||||
model_count: int
|
||||
total_api_keys: int | None
|
||||
metrics: SpendMetrics
|
||||
|
||||
|
||||
def _metadata(row: _ActivityRow) -> DailySpendMetadata:
|
||||
metrics: Final = row.metrics
|
||||
return DailySpendMetadata(
|
||||
total_spend=metrics.spend,
|
||||
total_flat_cost=metrics.flat_cost,
|
||||
total_prompt_tokens=metrics.prompt_tokens,
|
||||
total_completion_tokens=metrics.completion_tokens,
|
||||
total_tokens=metrics.total_tokens,
|
||||
total_api_requests=metrics.api_requests,
|
||||
total_successful_requests=metrics.successful_requests,
|
||||
total_failed_requests=metrics.failed_requests,
|
||||
total_cache_read_input_tokens=metrics.cache_read_input_tokens,
|
||||
total_cache_creation_input_tokens=metrics.cache_creation_input_tokens,
|
||||
total_compression_saved_tokens=metrics.compression_saved_tokens,
|
||||
total_compression_savings_spend=metrics.compression_savings_spend,
|
||||
total_prompt_caching_savings_spend=metrics.prompt_caching_savings_spend,
|
||||
total_gateway_injected_caching_savings_spend=metrics.gateway_injected_caching_savings_spend,
|
||||
total_autorouter_savings_spend=metrics.autorouter_savings_spend,
|
||||
total_response_time_ms=metrics.total_response_time_ms,
|
||||
total_timed_requests=metrics.timed_requests,
|
||||
total_api_keys=row.total_api_keys,
|
||||
)
|
||||
|
||||
|
||||
class BoundedDailyActivityRepository:
|
||||
def __init__(self, database: DailyActivityQueryDatabase) -> None:
|
||||
self.database = database
|
||||
|
||||
@classmethod
|
||||
def from_prisma(cls, prisma: "PrismaClient") -> "BoundedDailyActivityRepository":
|
||||
return cls(BoundedAnalyticsDatabase(cast(AnalyticsDatabase, prisma.db))) # cast-ok: generated Prisma tx
|
||||
|
||||
async def aggregate(
|
||||
self, scope: DailyActivityScope, *, model_limit: int | None = None
|
||||
) -> BoundedDailyActivityResponse:
|
||||
if model_limit is not None and not 1 <= model_limit <= 100:
|
||||
raise ValueError("model_limit must be between 1 and 100")
|
||||
query: Final = build_bounded_daily_activity_sql(scope, model_limit=model_limit)
|
||||
rows: Final = TypeAdapter(tuple[_ActivityRow, ...]).validate_python(
|
||||
await self.database.query_raw(query.sql, *query.params)
|
||||
)
|
||||
total: Final = next(row for row in rows if row.kind == "total")
|
||||
other: Final = next((row for row in rows if row.kind == "other"), None)
|
||||
days: Final = tuple(
|
||||
DailySpendData(date=row.date, metrics=row.metrics)
|
||||
for row in rows
|
||||
if row.kind == "day" and row.date is not None
|
||||
)
|
||||
models: Final = tuple(
|
||||
DailyActivityModelContribution(model=row.model, metrics=row.metrics)
|
||||
for row in rows
|
||||
if row.kind == "model" and row.model is not None
|
||||
)
|
||||
return BoundedDailyActivityResponse(
|
||||
results=tuple(sorted(days, key=lambda day: day.date, reverse=True)),
|
||||
metadata=_metadata(total),
|
||||
model_breakdown=DailyActivityModelBreakdown(
|
||||
limit=model_limit,
|
||||
total_models=total.model_count,
|
||||
models=tuple(sorted(models, key=lambda model: (-model.metrics.spend, model.model))),
|
||||
other=DailyActivityOtherModels(model_count=other.model_count, metrics=other.metrics) if other else None,
|
||||
)
|
||||
if model_limit is not None
|
||||
else None,
|
||||
)
|
||||
|
|
@ -71,6 +71,10 @@ class _QueryRaw(Protocol):
|
|||
async def __call__(self, query: str, *values: object) -> Sequence[Mapping[str, object]] | None: ...
|
||||
|
||||
|
||||
class DailyActivityQueryDatabase(Protocol):
|
||||
async def query_raw(self, query: str, *values: object) -> Sequence[Mapping[str, object]] | None: ...
|
||||
|
||||
|
||||
class _DailyActivityDatabase(Protocol):
|
||||
query_raw: _QueryRaw
|
||||
litellm_usertable: TableActions[_UserTableRow]
|
||||
|
|
@ -155,9 +159,16 @@ def _next_export_cursor(batch: tuple[ExportRow, ...], export_type: ExportType) -
|
|||
|
||||
|
||||
class DailyActivityRepository:
|
||||
def __init__(self, prisma_client: DailyActivityDatabase, *, proxy_reads: DailyActivityProxyReads) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
prisma_client: DailyActivityDatabase,
|
||||
*,
|
||||
proxy_reads: DailyActivityProxyReads,
|
||||
query_database: DailyActivityQueryDatabase | None = None,
|
||||
) -> None:
|
||||
self._prisma_client = prisma_client
|
||||
self._proxy_reads = proxy_reads
|
||||
self._query_database = query_database if query_database is not None else prisma_client.db
|
||||
|
||||
async def _query(self, query: SqlQuery) -> tuple[Mapping[str, object], ...]:
|
||||
from litellm.proxy.db.db_span import db_span
|
||||
|
|
@ -166,7 +177,7 @@ class DailyActivityRepository:
|
|||
first_line: Final = query.sql.lstrip().splitlines()[0].lstrip("(").strip()
|
||||
verbose_proxy_logger.debug("DailyActivityRepository query: %s", first_line)
|
||||
async with db_span("daily_activity_query", sql_relation(query.sql)):
|
||||
result: Sequence[Mapping[str, object]] | None = await self._prisma_client.db.query_raw(
|
||||
result: Final[Sequence[Mapping[str, object]] | None] = await self._query_database.query_raw(
|
||||
query.sql, *query.params
|
||||
)
|
||||
if result is None:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from itertools import count, islice
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
|
@ -159,6 +159,78 @@ def _top_api_keys_sql(pg_table: str, where_clause: str, *, sentinel_param: int,
|
|||
"""
|
||||
|
||||
|
||||
def _bounded_metrics_json(alias: str) -> str:
|
||||
return f"""(jsonb_strip_nulls(to_jsonb({alias}) - ARRAY[
|
||||
'date','model','group_level','model_count','total_api_keys','ptu_flat_cost'
|
||||
]) || jsonb_build_object(
|
||||
'flat_cost',COALESCE({alias}.ptu_flat_cost,0),
|
||||
'total_tokens',COALESCE({alias}.prompt_tokens,0)+COALESCE({alias}.completion_tokens,0)
|
||||
))"""
|
||||
|
||||
|
||||
def build_bounded_daily_activity_sql(scope: DailyActivityScope, *, model_limit: int | None) -> SqlQuery:
|
||||
days: Final = (date.fromisoformat(scope.end_date) - date.fromisoformat(scope.start_date)).days
|
||||
if not 0 <= days <= 365:
|
||||
raise ValueError("Choose an ordered date range of at most 366 days")
|
||||
if model_limit is not None and not 1 <= model_limit <= 100:
|
||||
raise ValueError("model_limit must be between 1 and 100")
|
||||
where, params = build_where_clause(scope)
|
||||
sentinel: Final = len(params) + 1
|
||||
limit_param: Final = len(params) + 2
|
||||
metrics: Final = _rollup_metric_select(scope.table)
|
||||
grouping: Final = (
|
||||
f"GROUPING SETS ((date),({_MODEL_EXPR}),())" if model_limit is not None else "GROUPING SETS ((date),())"
|
||||
)
|
||||
model: Final = _MODEL_EXPR if model_limit is not None else "NULL::text"
|
||||
level: Final = f"GROUPING(date,{_MODEL_EXPR})" if model_limit is not None else "GROUPING(date)"
|
||||
total_level: Final = 3 if model_limit is not None else 1
|
||||
day_level: Final = 1 if model_limit is not None else 0
|
||||
model_count: Final = "(SELECT count(*) FROM rollups WHERE group_level=2)" if model_limit is not None else "0"
|
||||
model_ctes: Final = (
|
||||
f""", top_models AS MATERIALIZED (
|
||||
SELECT * FROM rollups WHERE group_level=2 AND length(model)<=1024
|
||||
ORDER BY spend DESC NULLS LAST,model LIMIT ${limit_param}
|
||||
), other_models AS (
|
||||
SELECT count(*)::bigint AS model_count,{metrics}
|
||||
FROM rollups r WHERE group_level=2
|
||||
AND NOT EXISTS(SELECT 1 FROM top_models t WHERE t.model=r.model)
|
||||
HAVING count(*)>0
|
||||
)"""
|
||||
if model_limit is not None
|
||||
else ""
|
||||
)
|
||||
model_rows: Final = (
|
||||
f"""UNION ALL SELECT 'model',NULL::text,r.model,1::bigint,r.total_api_keys,{_bounded_metrics_json("r")}
|
||||
FROM top_models r
|
||||
UNION ALL SELECT 'other',NULL::text,NULL::text,r.model_count,NULL::bigint,{_bounded_metrics_json("r")}
|
||||
FROM other_models r"""
|
||||
if model_limit is not None
|
||||
else ""
|
||||
)
|
||||
sql: Final = f"""WITH rollups AS MATERIALIZED (
|
||||
SELECT date,{model} AS model,{level} AS group_level,
|
||||
count(DISTINCT api_key) FILTER(WHERE api_key<>${sentinel})::bigint AS total_api_keys,
|
||||
{metrics}
|
||||
FROM "{PRISMA_TO_PG_TABLE[scope.table]}" WHERE {where}
|
||||
GROUP BY {grouping}
|
||||
){model_ctes}
|
||||
SELECT 'day' AS kind,date,NULL::text AS model,0::bigint AS model_count,total_api_keys,
|
||||
{_bounded_metrics_json("r")} AS metrics FROM rollups r WHERE group_level={day_level}
|
||||
UNION ALL SELECT 'total',NULL::text,NULL::text,{model_count},total_api_keys,{_bounded_metrics_json("r")}
|
||||
FROM rollups r WHERE group_level={total_level}
|
||||
{model_rows}
|
||||
LIMIT ${limit_param + int(model_limit is not None)}"""
|
||||
return SqlQuery(
|
||||
sql=sql,
|
||||
params=(
|
||||
*params,
|
||||
PTU_SENTINEL_API_KEY,
|
||||
*((model_limit,) if model_limit is not None else ()),
|
||||
367 if model_limit is None else 368 + model_limit,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> SqlQuery:
|
||||
pg_table: Final = PRISMA_TO_PG_TABLE[scope.table]
|
||||
where_clause, where_params = build_where_clause(scope)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence
|
|||
from datetime import date
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import ConfigDict, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -264,3 +264,31 @@ class LiteLLM_DailyUserSpend(LiteLLMBaseModel):
|
|||
class GroupedData(TypedDict):
|
||||
metrics: SpendMetrics
|
||||
breakdown: BreakdownMetrics
|
||||
|
||||
|
||||
class DailyActivityModelContribution(BaseModel, frozen=True):
|
||||
model: str = Field(max_length=1024)
|
||||
metrics: SpendMetrics
|
||||
|
||||
|
||||
class DailyActivityOtherModels(BaseModel, frozen=True):
|
||||
model_count: int = Field(ge=1)
|
||||
metrics: SpendMetrics
|
||||
|
||||
|
||||
class DailyActivityModelBreakdown(BaseModel, frozen=True):
|
||||
limit: int = Field(ge=1, le=100)
|
||||
total_models: int = Field(ge=0)
|
||||
models: tuple[DailyActivityModelContribution, ...] = Field(max_length=100)
|
||||
other: DailyActivityOtherModels | None = None
|
||||
|
||||
|
||||
class BoundedDailyActivityResponse(BaseModel, frozen=True):
|
||||
results: tuple[DailySpendData, ...] = Field(max_length=366)
|
||||
metadata: DailySpendMetadata
|
||||
model_breakdown: DailyActivityModelBreakdown | None = Field(
|
||||
default=None,
|
||||
description="When requested, includes the highest-spend models across the full date range and an explicit "
|
||||
"aggregate of all remaining models. Daily totals always cover the complete activity scope. "
|
||||
"Other per-day dimension breakdowns are omitted.",
|
||||
)
|
||||
|
|
|
|||
134
tests/integration/database/test_bounded_daily_activity.py
Normal file
134
tests/integration/database/test_bounded_daily_activity.py
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
from collections.abc import Mapping
|
||||
from datetime import date, timedelta
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from integration._support.actors import ManagementDatabase
|
||||
from integration._support.client import Gateway
|
||||
from psycopg import sql
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.repositories.bounded_daily_activity import BoundedDailyActivityRepository
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import BoundedDailyActivityResponse
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_bounded_activity_routes_keep_exact_totals_and_scoped_permissions(
|
||||
gateway: Gateway, db: ManagementDatabase
|
||||
) -> None:
|
||||
repository: Final = BoundedDailyActivityRepository.from_prisma(db.prisma)
|
||||
identity: Final = uuid4().hex
|
||||
day: Final = date.today() - timedelta(days=3)
|
||||
with gateway.scenario() as scenario:
|
||||
owner: Final = scenario.user(user_role="internal_user")
|
||||
member: Final = scenario.user(user_role="internal_user")
|
||||
team: Final = scenario.team(
|
||||
members_with_roles=[{"user_id": owner, "role": "admin"}, {"user_id": member, "role": "user"}]
|
||||
)
|
||||
owner_key: Final = scenario.key(user_id=owner, team_id=team, key_type="management")
|
||||
member_key: Final = scenario.key(user_id=member, team_id=team, key_type="management")
|
||||
reporting_key: Final = sha256(owner_key.encode()).hexdigest()
|
||||
member_reporting_key: Final = sha256(member_key.encode()).hexdigest()
|
||||
for table, entity, column in (
|
||||
("LiteLLM_DailyUserSpend", owner, "user_id"),
|
||||
("LiteLLM_DailyTeamSpend", team, "team_id"),
|
||||
):
|
||||
await repository.database.query_raw(
|
||||
f'''INSERT INTO "{table}" (id,{column},date,api_key,model,custom_llm_provider,
|
||||
spend,api_requests,successful_requests,total_response_time_ms,timed_requests,updated_at)
|
||||
SELECT $1::text||'-'||m||'-'||d,$2::text,($3::date+d)::text,$4::text,
|
||||
CASE WHEN m=1 THEN 'Other' WHEN m=137 THEN repeat('x',1025)
|
||||
ELSE 'example-model-'||m END,'example-provider',
|
||||
CASE WHEN m=1 THEN 10000 ELSE m END,10,10,2000,10,CURRENT_TIMESTAMP
|
||||
FROM generate_series(1,137) m CROSS JOIN generate_series(0,2) d''',
|
||||
identity,
|
||||
entity,
|
||||
day.isoformat(),
|
||||
reporting_key,
|
||||
)
|
||||
await repository.database.query_raw(
|
||||
"""INSERT INTO "LiteLLM_DailyTeamSpend"
|
||||
(id,team_id,date,api_key,model,api_requests,successful_requests,updated_at)
|
||||
VALUES($1,$2,$3,$4,'member-model',7,7,CURRENT_TIMESTAMP)""",
|
||||
identity + "-member",
|
||||
team,
|
||||
day.isoformat(),
|
||||
member_reporting_key,
|
||||
)
|
||||
query: Final = {"start_date": day.isoformat(), "end_date": (day + timedelta(days=2)).isoformat()}
|
||||
try:
|
||||
path: Final = "/user/daily/activity/aggregated"
|
||||
legacy: Final = gateway.request("GET", path, params={**query, "api_key": reporting_key})
|
||||
assert legacy.status_code == 200 and "model_breakdown" not in legacy.json()
|
||||
result: Final = gateway.request(
|
||||
"GET",
|
||||
path,
|
||||
params={**query, "api_key": reporting_key, "bounded": "true", "include_model_breakdown": "true"},
|
||||
key=owner_key,
|
||||
)
|
||||
assert result.status_code == 200, result.text
|
||||
bounded: Final = BoundedDailyActivityResponse.model_validate(result.json())
|
||||
assert bounded.metadata.model_dump(exclude={"api_key_limit"}) == {
|
||||
key: value for key, value in legacy.json()["metadata"].items() if key != "api_key_limit"
|
||||
}
|
||||
models: Final = bounded.model_breakdown
|
||||
assert models is not None and models.total_models == 137 and len(models.models) == 20
|
||||
assert models.models[0].model == "Other" and models.other is not None and models.other.model_count == 117
|
||||
assert all(len(item.model) <= 1024 for item in models.models)
|
||||
assert (
|
||||
sum(item.metrics.spend for item in models.models) + models.other.metrics.spend
|
||||
== bounded.metadata.total_spend
|
||||
)
|
||||
assert len(result.content) < len(legacy.content) / 10
|
||||
assert bounded.metadata.total_api_requests == 4110 and len(bounded.results) == 3
|
||||
largest: Final = gateway.request(
|
||||
"GET",
|
||||
path,
|
||||
params={
|
||||
**query,
|
||||
"api_key": reporting_key,
|
||||
"bounded": "true",
|
||||
"include_model_breakdown": "true",
|
||||
"model_limit": "100",
|
||||
},
|
||||
)
|
||||
assert largest.status_code == 200 and len(largest.json()["model_breakdown"]["models"]) == 100
|
||||
assert (
|
||||
gateway.request("GET", path, params={**query, "bounded": "true", "model_limit": "101"}).status_code
|
||||
== 422
|
||||
)
|
||||
assert (
|
||||
gateway.request(
|
||||
"GET",
|
||||
path,
|
||||
params={**query, "bounded": "true", "end_date": (day + timedelta(days=366)).isoformat()},
|
||||
).status_code
|
||||
== 422
|
||||
)
|
||||
team_query: Final = {**query, "team_ids": team, "bounded": "true"}
|
||||
team_path: Final = "/team/daily/activity/aggregated"
|
||||
team_admin: Final = gateway.request("GET", team_path, params=team_query, key=owner_key)
|
||||
team_member: Final = gateway.request("GET", team_path, params=team_query, key=member_key)
|
||||
assert team_admin.status_code == team_member.status_code == 200
|
||||
assert team_admin.json()["metadata"]["total_api_requests"] == 4117
|
||||
assert team_member.json()["metadata"]["total_api_requests"] == 7
|
||||
assert team_admin.json()["model_breakdown"] is None
|
||||
schemas: Final = TypeAdapter(tuple[Mapping[str, str], ...]).validate_python(
|
||||
await repository.database.query_raw("SELECT current_schema() AS name")
|
||||
)
|
||||
assert gateway.database_url is not None
|
||||
with psycopg.connect(gateway.database_url) as blocker:
|
||||
blocker.execute(
|
||||
sql.SQL("LOCK TABLE {} IN ACCESS EXCLUSIVE MODE").format(
|
||||
sql.Identifier(schemas[0]["name"], "LiteLLM_DailyUserSpend")
|
||||
)
|
||||
)
|
||||
busy: Final = gateway.request("GET", path, params={**query, "bounded": "true"})
|
||||
assert busy.status_code == 503 and busy.headers["Retry-After"] == "5"
|
||||
assert gateway.request("GET", path, params={**query, "bounded": "true"}).status_code == 200
|
||||
finally:
|
||||
await repository.database.query_raw('DELETE FROM "LiteLLM_DailyUserSpend" WHERE user_id=$1', owner)
|
||||
await repository.database.query_raw('DELETE FROM "LiteLLM_DailyTeamSpend" WHERE team_id=$1', team)
|
||||
Loading…
Add table
Reference in a new issue