scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

refactor(proxy): inject gateway activity through a typed repository (#49116)

Browse source 
* refactor(proxy): inject gateway activity through a typed repository

* ci: keep dependency tests in their existing shard

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee 2026-10-30 16:25:11 -07:00 • committed by GitHub
parent fbc157437c
commit ef9385c712
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 215 additions and 126 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/models/gateway_activity.py Normal file
View file

@ -0,0 +1,31 @@
from typing import Literal, Protocol, TypeAlias
from pydantic import ConfigDict
from litellm.types.llms.base import LiteLLMBaseModel
class GatewayRouteActivity(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["route"]
date: str
category: str
route: str
successful_requests: int
failed_requests: int
class GatewayStatusCodeActivity(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["status"]
status_code: int
failed_requests: int
GatewayActivity: TypeAlias = GatewayRouteActivity | GatewayStatusCodeActivity
class GatewayActivityStore(Protocol):
async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]: ...

11
litellm/proxy/deps/repositories.py
View file

@ -3,9 +3,11 @@ from typing import Annotated, TypeAlias
from fastapi import Depends
from litellm.models.deployment_warning import DeploymentWarningIgnoreStore
from litellm.models.gateway_activity import GatewayActivityStore
from litellm.models.user_banner import UserBannerStore
from litellm.proxy.deps.database import DatabaseDep
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
@ -23,3 +25,12 @@ async def get_deployment_warning_ignore_repository(database: DatabaseDep) -> Dep
DeploymentWarningIgnoreRepositoryDep: TypeAlias = Annotated[
DeploymentWarningIgnoreStore | None, Depends(get_deployment_warning_ignore_repository)
]
async def get_gateway_activity_repository(database: DatabaseDep) -> GatewayActivityStore | None:
return None if database is None else GatewayActivityRepository(database.prisma)
GatewayActivityRepositoryDep: TypeAlias = Annotated[
GatewayActivityStore | None, Depends(get_gateway_activity_repository)
]

74
litellm/proxy/management_endpoints/gateway_request_endpoints.py
View file

@ -13,15 +13,15 @@ and the endpoint is restricted to proxy admin roles.
from collections.abc import Sequence
from datetime import datetime, timedelta, timezone
from typing import Annotated, Final, Literal
from typing import Annotated, Final
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import ConfigDict, Field, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.proxy.deps.repositories import GatewayActivityRepositoryDep
from litellm.types.proxy.gateway_requests import (
GatewayRequestActivityResponse,
GatewayRequestBreakdownEntry,
@ -33,50 +33,6 @@ router: Final = APIRouter()
_DEFAULT_LOOKBACK_DAYS: Final = 30
_ACTIVITY_SQL: Final = """
SELECT
'route' AS kind, date, category, route,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests,
NULL::integer AS status_code
FROM "LiteLLM_DailyGatewayRequests"
WHERE date >= $1 AND date <= $2
GROUP BY date, category, route
UNION ALL
SELECT
'status' AS kind, NULL, NULL, NULL,
0::bigint,
SUM(failed_requests)::bigint,
status_code
FROM "LiteLLM_DailyGatewayFailedRequests"
WHERE date >= $1 AND date <= $2
GROUP BY status_code
"""
class _AggregateRow(LiteLLMBaseModel):
"""Validates one query_raw row so the handler works with typed values, not Any."""
kind: Literal["route"]
date: str
category: str
route: str
successful_requests: int
failed_requests: int
class _StatusCodeAggregateRow(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["status"]
status_code: int
failed_requests: int
_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter(
tuple[Annotated[_AggregateRow | _StatusCodeAggregateRow, Field(discriminator="kind")], ...]
)
def _default_range() -> tuple[str, str]:
end: Final = datetime.now(timezone.utc)
@ -84,7 +40,7 @@ def _default_range() -> tuple[str, str]:
return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")
def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEntry, ...]:
def _fold_by_date(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestDailyEntry, ...]:
dates: Final = sorted(frozenset(row.date for row in rows))
return tuple(
GatewayRequestDailyEntry(
@ -96,7 +52,7 @@ def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEnt
)
def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdownEntry, ...]:
def _fold_by_route(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestBreakdownEntry, ...]:
pairs: Final = sorted(frozenset((row.category, row.route) for row in rows))
entries: Final = tuple(
GatewayRequestBreakdownEntry(
@ -113,7 +69,7 @@ def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdo
def _fold_by_status_code(
rows: Sequence[_StatusCodeAggregateRow],
rows: Sequence[GatewayStatusCodeActivity],
) -> tuple[GatewayRequestStatusCodeEntry, ...]:
return tuple(
GatewayRequestStatusCodeEntry(status_code=row.status_code, failed_requests=row.failed_requests)
@ -128,6 +84,7 @@ def _fold_by_status_code(
)
async def get_gateway_daily_activity(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
activity: GatewayActivityRepositoryDep,
start_date: str | None = Query(default=None, description="Start date in YYYY-MM-DD format"),
end_date: str | None = Query(default=None, description="End date in YYYY-MM-DD format"),
) -> GatewayRequestActivityResponse:
@ -137,8 +94,6 @@ async def get_gateway_daily_activity(
Deployment-wide: the underlying table has no per-key or per-user dimension,
so this is admin-only.
"""
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
@ -148,22 +103,15 @@ async def get_gateway_daily_activity(
detail="Only proxy admin roles can view gateway request counts across the deployment",
)
if prisma_client is None:
if activity is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
default_start, default_end = _default_range()
selected_start: Final = start_date or default_start
selected_end: Final = end_date or default_end
raw_rows: Final = await prisma_client.db.query_raw( # pyright: ignore[reportAny] # untyped prisma client
_ACTIVITY_SQL,
selected_start,
selected_end,
)
activity_rows: Final = _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ())
rows: Final[tuple[_AggregateRow, ...]] = tuple(row for row in activity_rows if isinstance(row, _AggregateRow))
status_code_rows: Final[tuple[_StatusCodeAggregateRow, ...]] = tuple(
row for row in activity_rows if isinstance(row, _StatusCodeAggregateRow)
)
activity_rows: Final = await activity.get_activity(selected_start, selected_end)
rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayRouteActivity))
status_code_rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayStatusCodeActivity))
verbose_proxy_logger.debug("/gateway/daily/activity - aggregated %d rows", len(rows))
return GatewayRequestActivityResponse(

41
litellm/repositories/gateway_activity_repository.py Normal file
View file

@ -0,0 +1,41 @@
from typing import TYPE_CHECKING, Annotated, Final
from pydantic import Field, TypeAdapter
from litellm.models.gateway_activity import GatewayActivity
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_ACTIVITY_SQL: Final = """
SELECT
'route' AS kind, date, category, route,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests,
NULL::integer AS status_code
FROM "LiteLLM_DailyGatewayRequests"
WHERE date >= $1 AND date <= $2
GROUP BY date, category, route
UNION ALL
SELECT
'status' AS kind, NULL, NULL, NULL,
0::bigint,
SUM(failed_requests)::bigint,
status_code
FROM "LiteLLM_DailyGatewayFailedRequests"
WHERE date >= $1 AND date <= $2
GROUP BY status_code
"""
_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter(tuple[Annotated[GatewayActivity, Field(discriminator="kind")], ...])
class GatewayActivityRepository:
def __init__(self, prisma_client: "PrismaClient") -> None:
self._prisma_client = prisma_client
async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]:
raw_rows: Final = await self._prisma_client.db.query_raw( # pyright: ignore[reportAny] # untyped Prisma client
_ACTIVITY_SQL, start_date, end_date
)
return _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ())

29
tests/unit/proxy/deps/test_repositories.py
View file

@ -8,9 +8,14 @@ from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.models.deployment_warning import AdminIgnore
from litellm.models.gateway_activity import GatewayActivity
from litellm.models.user_banner import UserBanner
from litellm.proxy.deps.database import get_database
from litellm.proxy.deps.repositories import DeploymentWarningIgnoreRepositoryDep, UserBannerRepositoryDep
from litellm.proxy.deps.repositories import (
DeploymentWarningIgnoreRepositoryDep,
GatewayActivityRepositoryDep,
UserBannerRepositoryDep,
)
from litellm.repositories import database as database_runtime
from tests._support.native_database import native_database
@ -32,6 +37,10 @@ def _app(database: database_runtime.Database | None) -> FastAPI:
async def ignores(ignores: DeploymentWarningIgnoreRepositoryDep) -> Mapping[str, AdminIgnore] | None:
return None if ignores is None else await ignores.get_ignores()
@app.get("/activity")
async def activity(activity: GatewayActivityRepositoryDep) -> tuple[GatewayActivity, ...] | None:
return None if activity is None else await activity.get_activity("2026-01-01", "2026-01-31")
return app
@ -83,3 +92,21 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne
assert client.get("/banner").json() == UserBanner(message="stored").model_dump()
monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
assert client.get("/banner").json() == UserBanner().model_dump()
def test_gateway_activity_reads_the_lifespan_reader_not_the_writer() -> None:
prisma: Final = MagicMock()
row: Final = {"kind": "status", "status_code": 429, "failed_requests": 3}
prisma.db.query_raw = AsyncMock(return_value=[row])
prisma.writer_db.query_raw = AsyncMock(return_value=[{**row, "failed_requests": 99}])
with TestClient(_app(database_runtime.Database(prisma=prisma, native=NATIVE))) as client:
response: Final = client.get("/activity")
assert response.status_code == 200
assert response.json() == [row]
def test_gateway_activity_without_database_is_unavailable() -> None:
with TestClient(_app(None)) as client:
response: Final = client.get("/activity")
assert response.status_code == 200
assert response.json() is None

108
tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py
View file

@ -1,47 +1,25 @@
import os
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, call, patch
from tests._master_key import MASTER_KEY
# Patching ``litellm.proxy.proxy_server.prisma_client`` imports that module, whose
# module-level setup reads DATABASE_URL and LITELLM_MASTER_KEY. Tier-zero runners
# set neither, so pin throwaways first, as test_component_allowlists.py does. The
# prior values are restored below so a non-postgres URL cannot leak into sibling
# tests sharing the xdist worker and make them treat a phantom database as live.
_THROWAWAY_ENV = {
"DATABASE_URL": "sqlite:///:memory:",
"LITELLM_MASTER_KEY": MASTER_KEY,
}
_PRE_EXISTING_ENV = {key: os.environ.get(key) for key in _THROWAWAY_ENV}
for _key, _value in _THROWAWAY_ENV.items():
os.environ.setdefault(_key, _value)
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import get_gateway_activity_repository
from litellm.proxy.management_endpoints.gateway_request_endpoints import (
_ACTIVITY_SQL,
_AggregateRow,
_default_range,
_fold_by_date,
_fold_by_route,
_fold_by_status_code,
_StatusCodeAggregateRow,
get_gateway_daily_activity,
router,
)
for _key, _previous in _PRE_EXISTING_ENV.items():
if _previous is None:
os.environ.pop(_key, None)
else:
os.environ[_key] = _previous
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
# The handler stamps "today" from the wall clock, so any assertion that names a
# date has to pin it. Recomputing the expected range in the assertion instead
@ -67,8 +45,8 @@ def _row(
route: str = "/chat/completions",
successful: int = 0,
failed: int = 0,
) -> _AggregateRow:
return _AggregateRow(
) -> GatewayRouteActivity:
return GatewayRouteActivity(
kind="route",
date=date,
category=category,
@ -86,9 +64,7 @@ def _prisma_returning(rows: Sequence[Mapping[str, str | int | None]] | None) ->
client: Final = MagicMock()
client.db = MagicMock()
async def query_raw(query: str, *_args: str) -> Sequence[Mapping[str, str | int | None]] | None:
if query != _ACTIVITY_SQL:
raise AssertionError(f"Unexpected gateway activity query: {query}")
async def query_raw(_query: str, *_args: str) -> Sequence[Mapping[str, str | int | None]] | None:
return rows
client.db.query_raw = AsyncMock(side_effect=query_raw)
@ -171,9 +147,9 @@ class TestFoldByStatusCode:
def test_orders_by_failed_count_then_status_code(self):
folded: Final = _fold_by_status_code(
(
_StatusCodeAggregateRow(kind="status", status_code=503, failed_requests=1),
_StatusCodeAggregateRow(kind="status", status_code=500, failed_requests=2),
_StatusCodeAggregateRow(kind="status", status_code=429, failed_requests=2),
GatewayStatusCodeActivity(kind="status", status_code=503, failed_requests=1),
GatewayStatusCodeActivity(kind="status", status_code=500, failed_requests=2),
GatewayStatusCodeActivity(kind="status", status_code=429, failed_requests=2),
)
)
assert [(entry.status_code, entry.failed_requests) for entry in folded] == [(429, 2), (500, 2), (503, 1)]
@ -191,11 +167,11 @@ class TestGatewayDailyActivityEndpoint:
],
)
async def test_refuses_every_non_admin_role(self, role):
with patch("litellm.proxy.proxy_server.prisma_client", _prisma_returning([])):
with pytest.raises(HTTPException) as exc:
await get_gateway_daily_activity(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role),
)
with pytest.raises(HTTPException) as exc:
await get_gateway_daily_activity(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role),
activity=None,
)
assert exc.value.status_code == 403
@pytest.mark.asyncio
@ -204,17 +180,18 @@ class TestGatewayDailyActivityEndpoint:
[LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY],
)
async def test_serves_both_admin_roles(self, role):
with patch("litellm.proxy.proxy_server.prisma_client", _prisma_returning([])):
response = await get_gateway_daily_activity(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role),
)
response = await get_gateway_daily_activity(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role),
activity=GatewayActivityRepository(_prisma_returning([])),
start_date="2026-01-01",
end_date="2026-01-31",
)
assert response.total_successful_requests == 0
@pytest.mark.asyncio
async def test_reports_db_not_connected_rather_than_crashing(self):
with patch("litellm.proxy.proxy_server.prisma_client", None):
with pytest.raises(HTTPException) as exc:
await get_gateway_daily_activity(user_api_key_dict=_admin())
with pytest.raises(HTTPException) as exc:
await get_gateway_daily_activity(user_api_key_dict=_admin(), activity=None)
assert exc.value.status_code == 500
@pytest.mark.asyncio
@ -267,8 +244,12 @@ class TestGatewayDailyActivityEndpoint:
},
)
prisma: Final = _prisma_returning(rows)
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
response: Final = await get_gateway_daily_activity(user_api_key_dict=_admin())
response: Final = await get_gateway_daily_activity(
user_api_key_dict=_admin(),
activity=GatewayActivityRepository(prisma),
start_date="2026-08-03",
end_date="2026-08-04",
)
assert response.total_successful_requests == 16
assert response.total_failed_requests == 4
@ -277,13 +258,16 @@ class TestGatewayDailyActivityEndpoint:
assert [(entry.status_code, entry.failed_requests) for entry in response.by_status_code] == [(429, 3), (500, 1)]
assert [entry.date for entry in response.by_date] == ["2026-08-03", "2026-08-04"]
assert [entry.route for entry in response.by_route] == ["/chat/completions", "/embeddings"]
assert prisma.db.query_raw.call_count == 1
@pytest.mark.asyncio
async def test_a_null_result_set_is_not_an_error(self):
client = _prisma_returning(None)
with patch("litellm.proxy.proxy_server.prisma_client", client):
response = await get_gateway_daily_activity(user_api_key_dict=_admin())
response = await get_gateway_daily_activity(
user_api_key_dict=_admin(),
activity=GatewayActivityRepository(client),
start_date="2026-01-01",
end_date="2026-01-31",
)
assert response.total_successful_requests == 0
assert response.by_date == ()
assert response.by_route == ()
@ -305,23 +289,23 @@ class TestGatewayDailyActivityRoute:
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = _admin
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
response = TestClient(app).get(
"/gateway/daily/activity",
params={"start_date": "2026-01-01", "end_date": "2026-01-31"},
)
app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma)
response = TestClient(app).get(
"/gateway/daily/activity",
params={"start_date": "2026-01-01", "end_date": "2026-01-31"},
)
assert response.status_code == 200
assert prisma.db.query_raw.call_args_list == [call(_ACTIVITY_SQL, "2026-01-01", "2026-01-31")]
assert prisma.db.query_raw.call_args_list[0].args[1:] == ("2026-01-01", "2026-01-31")
def test_omitted_dates_fall_back_to_the_default_window(self, frozen_clock):
prisma = _prisma_returning([])
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = _admin
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
response = TestClient(app).get("/gateway/daily/activity")
app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma)
response = TestClient(app).get("/gateway/daily/activity")
assert response.status_code == 200
assert prisma.db.query_raw.call_args_list == [call(_ACTIVITY_SQL, *_FROZEN_RANGE)]
assert prisma.db.query_raw.call_args_list[0].args[1:] == _FROZEN_RANGE
def test_serialized_response_carries_the_documented_shape(self):
prisma: Final = _prisma_returning(
@ -349,8 +333,8 @@ class TestGatewayDailyActivityRoute:
app: Final = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = _admin
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
body: Final = TestClient(app).get("/gateway/daily/activity").json()
app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma)
body: Final = TestClient(app).get("/gateway/daily/activity").json()
assert body == {
"total_successful_requests": 7,

47
tests/unit/repositories/test_gateway_activity_repository.py Normal file
View file

@ -0,0 +1,47 @@
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from pydantic import ValidationError
from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity
from litellm.repositories.gateway_activity_repository import GatewayActivityRepository
@pytest.mark.asyncio
async def test_gateway_activity_decodes_both_row_kinds_with_nullable_union_columns() -> None:
prisma: Final = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=[
{
"kind": "route", "date": "2026-01-02", "category": "llm", "route": "/chat/completions",
"successful_requests": 7, "failed_requests": 2, "status_code": None,
},
{
"kind": "status", "date": None, "category": None, "route": None,
"successful_requests": 0, "failed_requests": 2, "status_code": 429,
},
])
result: Final = await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31")
assert result == (
GatewayRouteActivity(
kind="route", date="2026-01-02", category="llm", route="/chat/completions",
successful_requests=7, failed_requests=2,
),
GatewayStatusCodeActivity(kind="status", status_code=429, failed_requests=2),
)
@pytest.mark.asyncio
async def test_gateway_activity_rejects_malformed_database_rows() -> None:
prisma: Final = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=[{"kind": "status", "status_code": None, "failed_requests": 2}])
with pytest.raises(ValidationError):
await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31")
@pytest.mark.asyncio
async def test_gateway_activity_does_not_hide_database_failure() -> None:
prisma: Final = MagicMock()
prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("reader unavailable"))
with pytest.raises(RuntimeError, match="reader unavailable"):
await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31")