mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge dc54c09246 into 58226c585c
This commit is contained in:
commit
7f76294cd2
25 changed files with 1615 additions and 36 deletions
|
|
@ -15,9 +15,16 @@ from litellm.proxy.control_plane.common.exceptions import (
|
|||
install_exception_handlers,
|
||||
)
|
||||
from litellm.proxy.control_plane.container import Container
|
||||
from litellm.proxy.control_plane.dataplanes.dependencies import get_dataplane_client, get_registry
|
||||
from litellm.proxy.control_plane.dataplanes.dependencies import (
|
||||
get_control_plane_settings,
|
||||
get_dataplane_client,
|
||||
get_fanout,
|
||||
get_registry,
|
||||
)
|
||||
from litellm.proxy.control_plane.dataplanes.router import router as dataplanes_router
|
||||
from litellm.proxy.control_plane.fleet.router import router as fleet_router
|
||||
from litellm.proxy.control_plane.health.router import router as health_router
|
||||
from litellm.proxy.control_plane.usage.router import router as usage_router
|
||||
from litellm.proxy.middleware.request_size_limit_middleware import RequestSizeLimitMiddleware
|
||||
|
||||
TITLE: Final = "LiteLLM Control Plane"
|
||||
|
|
@ -82,9 +89,13 @@ def create_app(container: Container, license_check: License) -> FastAPI:
|
|||
app.dependency_overrides[get_token_verifier] = lambda: container.token_verifier
|
||||
app.dependency_overrides[get_login_service] = lambda: container.login_service
|
||||
app.dependency_overrides[get_registry] = lambda: container.registry
|
||||
app.dependency_overrides[get_control_plane_settings] = lambda: container.settings
|
||||
app.dependency_overrides[get_dataplane_client] = lambda: container.dataplane_client
|
||||
app.dependency_overrides[get_fanout] = lambda: container.fanout
|
||||
install_exception_handlers(app)
|
||||
app.include_router(health_router)
|
||||
app.include_router(auth_router, prefix=container.settings.api_prefix)
|
||||
app.include_router(dataplanes_router, prefix=container.settings.api_prefix)
|
||||
app.include_router(usage_router, prefix=container.settings.api_prefix)
|
||||
app.include_router(fleet_router, prefix=container.settings.api_prefix)
|
||||
return app
|
||||
|
|
|
|||
|
|
@ -25,6 +25,10 @@ def error_body(slug: str, title: str, status: int, detail: str) -> ErrorBody:
|
|||
return ErrorBody(type=f"{ERROR_TYPE_BASE}{slug}", title=title, status=status, detail=detail)
|
||||
|
||||
|
||||
def bad_request(detail: str) -> ErrorBody:
|
||||
return error_body("bad-request", "Bad Request", HTTPStatus.BAD_REQUEST, detail)
|
||||
|
||||
|
||||
def unauthorized(detail: str) -> ErrorBody:
|
||||
return error_body("unauthorized", "Unauthorized", 401, detail)
|
||||
|
||||
|
|
@ -41,6 +45,14 @@ def too_many_requests(detail: str) -> ErrorBody:
|
|||
return error_body("too-many-requests", "Too Many Requests", 429, detail)
|
||||
|
||||
|
||||
def bad_gateway(detail: str) -> ErrorBody:
|
||||
return error_body("bad-gateway", "Bad Gateway", 502, detail)
|
||||
|
||||
|
||||
def service_unavailable(detail: str) -> ErrorBody:
|
||||
return error_body("service-unavailable", "Service Unavailable", 503, detail)
|
||||
|
||||
|
||||
class ControlPlaneException(Exception):
|
||||
"""An error raised from a dependency or handler, rendered by the app-level handler below."""
|
||||
|
||||
|
|
@ -50,10 +62,6 @@ class ControlPlaneException(Exception):
|
|||
super().__init__(body.detail)
|
||||
|
||||
|
||||
def bad_gateway(detail: str) -> ErrorBody:
|
||||
return error_body("bad-gateway", "Bad Gateway", 502, detail)
|
||||
|
||||
|
||||
def _render(body: ErrorBody) -> JSONResponse:
|
||||
response: Final = error_response(body)
|
||||
if body.status == 401:
|
||||
|
|
|
|||
90
litellm/proxy/control_plane/common/fanout.py
Normal file
90
litellm/proxy/control_plane/common/fanout.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
"""Bounded scatter/gather over dataplanes. A dataplane that fails is a Missed value in the result, never an exception."""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Generic, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.proxy.control_plane.dataplanes.client import CallResult, DataplaneReply, DataplaneUnreachable
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import CredentialUnavailable
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
|
||||
T: Final = TypeVar("T")
|
||||
M: Final = TypeVar("M")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Failure:
|
||||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Missed:
|
||||
dataplane: Dataplane
|
||||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Reached(Generic[T]):
|
||||
dataplane: Dataplane
|
||||
value: T
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Gathered(Generic[T]):
|
||||
reached: tuple[Reached[T], ...]
|
||||
missed: tuple[Missed, ...]
|
||||
|
||||
@property
|
||||
def partial(self) -> bool:
|
||||
return bool(self.missed)
|
||||
|
||||
@property
|
||||
def nothing_answered(self) -> bool:
|
||||
return bool(self.missed) and not self.reached
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Fanout:
|
||||
concurrency: int
|
||||
_gate: asyncio.Semaphore = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "_gate", asyncio.Semaphore(self.concurrency))
|
||||
|
||||
async def gather(
|
||||
self, dataplanes: tuple[Dataplane, ...], fetch: Callable[[Dataplane], Awaitable[T | Failure]]
|
||||
) -> Gathered[T]:
|
||||
async def bounded(dataplane: Dataplane) -> Reached[T] | Missed:
|
||||
async with self._gate:
|
||||
outcome: Final = await fetch(dataplane)
|
||||
if isinstance(outcome, Failure):
|
||||
return Missed(dataplane, outcome.reason)
|
||||
return Reached(dataplane, outcome)
|
||||
|
||||
outcomes: Final = await asyncio.gather(*(bounded(dataplane) for dataplane in dataplanes))
|
||||
return Gathered(
|
||||
reached=tuple(outcome for outcome in outcomes if isinstance(outcome, Reached)),
|
||||
missed=tuple(outcome for outcome in outcomes if isinstance(outcome, Missed)),
|
||||
)
|
||||
|
||||
|
||||
def parse_reply(result: CallResult, shape: TypeAdapter[M]) -> M | Failure:
|
||||
match result:
|
||||
case DataplaneReply(status=status, body=body) if status < 400:
|
||||
try:
|
||||
return shape.validate_python(body)
|
||||
except ValidationError as error:
|
||||
return Failure(f"unexpected response shape: {error.error_count()} validation error(s)")
|
||||
case DataplaneReply(status=status, credential=credential):
|
||||
return Failure(f"HTTP {status} with the {credential} credential")
|
||||
case DataplaneUnreachable(detail=detail, attempts=attempts):
|
||||
return Failure(f"unreachable after {attempts} attempt(s): {detail}")
|
||||
case CredentialUnavailable(reason=reason):
|
||||
return Failure(f"no usable credential: {reason}")
|
||||
|
||||
|
||||
def describe_missed(missed: tuple[Missed, ...]) -> str:
|
||||
return ", ".join(f"{miss.dataplane.id}: {miss.reason}" for miss in missed)
|
||||
|
|
@ -24,6 +24,7 @@ from litellm.proxy.control_plane.auth.rbac import (
|
|||
build_policy,
|
||||
parse_role_definition,
|
||||
)
|
||||
from litellm.proxy.control_plane.common.fanout import Fanout
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, HttpTransport, RetryPolicy
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import TokenExchanger, http_token_exchanger
|
||||
from litellm.proxy.control_plane.dataplanes.registry import DataplaneRegistry, build_registry
|
||||
|
|
@ -39,6 +40,7 @@ class Container:
|
|||
login_service: LoginService
|
||||
registry: DataplaneRegistry
|
||||
dataplane_client: DataplaneClient
|
||||
fanout: Fanout
|
||||
|
||||
|
||||
class InvalidRolePolicy(ValueError):
|
||||
|
|
@ -143,4 +145,5 @@ def build_container(
|
|||
timeout_seconds=settings.dataplane_timeout_seconds,
|
||||
retry=RetryPolicy(attempts=settings.dataplane_retry_attempts),
|
||||
),
|
||||
fanout=Fanout(concurrency=settings.fanout_concurrency),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,20 @@
|
|||
from typing import Annotated, Final
|
||||
|
||||
from fastapi import Depends
|
||||
from fastapi import Depends, Query
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
|
||||
from litellm.proxy.control_plane.auth.dependencies import bearer_scheme, get_principal
|
||||
from litellm.proxy.control_plane.auth.rbac import Principal
|
||||
from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, not_found, unauthorized
|
||||
from litellm.proxy.control_plane.common.fanout import Fanout
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry, visible_to
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
|
||||
|
||||
def get_control_plane_settings() -> ControlPlaneSettings:
|
||||
raise RuntimeError("get_control_plane_settings is not bound, build the app with create_app")
|
||||
|
||||
|
||||
def get_registry() -> DataplaneRegistry:
|
||||
|
|
@ -19,6 +25,10 @@ def get_dataplane_client() -> DataplaneClient:
|
|||
raise RuntimeError("get_dataplane_client is not bound, build the app with create_app")
|
||||
|
||||
|
||||
def get_fanout() -> Fanout:
|
||||
raise RuntimeError("get_fanout is not bound, build the app with create_app")
|
||||
|
||||
|
||||
async def get_caller(
|
||||
credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)],
|
||||
principal: Annotated[Principal, Depends(get_principal)],
|
||||
|
|
@ -39,3 +49,13 @@ async def dataplane_access(
|
|||
if not principal.dataplanes.allows(dataplane_id):
|
||||
raise ControlPlaneException(forbidden(f"Not bound to dataplane {dataplane_id!r}."))
|
||||
return dataplane
|
||||
|
||||
|
||||
async def selected_dataplanes(
|
||||
principal: Annotated[Principal, Depends(get_principal)],
|
||||
registry: Annotated[DataplaneRegistry, Depends(get_registry)],
|
||||
dataplane: Annotated[str | None, Query(description="Restrict a fleet-wide call to one dataplane id.")] = None,
|
||||
) -> tuple[Dataplane, ...]:
|
||||
if dataplane is None:
|
||||
return visible_to(registry, principal)
|
||||
return (await dataplane_access(dataplane, principal, registry),)
|
||||
|
|
|
|||
35
litellm/proxy/control_plane/dataplanes/identity.py
Normal file
35
litellm/proxy/control_plane/dataplanes/identity.py
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
USER_INFO_PATH: Final = "/user/info"
|
||||
|
||||
|
||||
class UserInfo(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
user_id: str | None = None
|
||||
user_email: str | None = None
|
||||
user_role: str | None = None
|
||||
|
||||
|
||||
class UserInfoBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
user_id: str | None = None
|
||||
user_info: UserInfo | None = None
|
||||
|
||||
|
||||
def parse_user_info(body: object) -> UserInfoBody | None:
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
try:
|
||||
return UserInfoBody.model_validate(body)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def resolved_user_id(info: UserInfoBody) -> str | None:
|
||||
if info.user_id is not None:
|
||||
return info.user_id
|
||||
return None if info.user_info is None else info.user_info.user_id
|
||||
|
|
@ -2,7 +2,6 @@ from http import HTTPStatus
|
|||
from typing import Annotated, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, Security
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm.proxy.control_plane.auth.dependencies import require_scopes
|
||||
from litellm.proxy.control_plane.auth.rbac import Principal, Scope
|
||||
|
|
@ -15,6 +14,7 @@ from litellm.proxy.control_plane.dataplanes.dependencies import (
|
|||
get_dataplane_client,
|
||||
get_registry,
|
||||
)
|
||||
from litellm.proxy.control_plane.dataplanes.identity import USER_INFO_PATH, UserInfo, parse_user_info
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry, visible_to
|
||||
from litellm.types.proxy.control_plane.dataplanes import (
|
||||
DataplaneHealthResponse,
|
||||
|
|
@ -26,22 +26,6 @@ from litellm.types.proxy.control_plane.dataplanes import (
|
|||
router: Final = APIRouter(prefix="/dataplanes", tags=["dataplanes"])
|
||||
|
||||
LIVENESS_PATH: Final = "/health/liveliness"
|
||||
USER_INFO_PATH: Final = "/user/info"
|
||||
|
||||
|
||||
class _UserInfo(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
user_id: str | None = None
|
||||
user_email: str | None = None
|
||||
user_role: str | None = None
|
||||
|
||||
|
||||
class _UserInfoBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
user_id: str | None = None
|
||||
user_info: _UserInfo | None = None
|
||||
|
||||
|
||||
def _to_response(dataplane: Dataplane) -> DataplaneResponse:
|
||||
|
|
@ -96,10 +80,10 @@ async def dataplane_whoami(
|
|||
result: Final = await client.get(dataplane, USER_INFO_PATH, caller)
|
||||
match result:
|
||||
case DataplaneReply(status=status, body=body, credential=credential) if status < 400:
|
||||
info: Final = _identity_body(body)
|
||||
info: Final = parse_user_info(body)
|
||||
if info is None:
|
||||
raise ControlPlaneException(bad_gateway("Dataplane returned a malformed /user/info body."))
|
||||
user: Final = info.user_info or _UserInfo()
|
||||
user: Final = info.user_info or UserInfo()
|
||||
return DataplaneIdentityResponse(
|
||||
id=dataplane.id,
|
||||
credential=credential,
|
||||
|
|
@ -119,12 +103,3 @@ async def dataplane_whoami(
|
|||
raise ControlPlaneException(bad_gateway(f"Dataplane unreachable after {attempts} attempt(s): {detail}"))
|
||||
case CredentialUnavailable(reason=reason):
|
||||
raise ControlPlaneException(bad_gateway(f"No usable credential for this dataplane: {reason}"))
|
||||
|
||||
|
||||
def _identity_body(body: object) -> _UserInfoBody | None:
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
try:
|
||||
return _UserInfoBody.model_validate(body)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
|
|
|||
0
litellm/proxy/control_plane/fleet/__init__.py
Normal file
0
litellm/proxy/control_plane/fleet/__init__.py
Normal file
74
litellm/proxy/control_plane/fleet/router.py
Normal file
74
litellm/proxy/control_plane/fleet/router.py
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
from typing import Annotated, Final, TypeAlias, TypeVar
|
||||
|
||||
from fastapi import APIRouter, Depends, Security
|
||||
|
||||
from litellm.proxy.control_plane.auth.dependencies import require_scopes
|
||||
from litellm.proxy.control_plane.auth.rbac import Principal, Scope
|
||||
from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, service_unavailable
|
||||
from litellm.proxy.control_plane.common.fanout import Fanout, Gathered, describe_missed
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller
|
||||
from litellm.proxy.control_plane.dataplanes.dependencies import (
|
||||
get_caller,
|
||||
get_dataplane_client,
|
||||
get_fanout,
|
||||
selected_dataplanes,
|
||||
)
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
from litellm.proxy.control_plane.fleet.service import fetch_models, fetch_teams, group_models
|
||||
from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane
|
||||
from litellm.types.proxy.control_plane.fleet import FleetModel, FleetModelsResponse, FleetTeam, FleetTeamsResponse
|
||||
|
||||
T: Final = TypeVar("T")
|
||||
|
||||
router: Final = APIRouter(prefix="/fleet", tags=["fleet"])
|
||||
|
||||
FleetRead: TypeAlias = Annotated[Principal, Security(require_scopes, scopes=[Scope.FLEET_READ.value])]
|
||||
Selected: TypeAlias = Annotated[tuple[Dataplane, ...], Depends(selected_dataplanes)]
|
||||
|
||||
|
||||
def _or_503(gathered: Gathered[T]) -> Gathered[T]:
|
||||
if gathered.nothing_answered:
|
||||
raise ControlPlaneException(service_unavailable(f"No dataplane answered: {describe_missed(gathered.missed)}"))
|
||||
return gathered
|
||||
|
||||
|
||||
def _flatten(gathered: Gathered[tuple[T, ...]]) -> tuple[T, ...]:
|
||||
return tuple(item for reached in gathered.reached for item in reached.value) # comprehension-ok: flatten
|
||||
|
||||
|
||||
def _unreachable(gathered: Gathered[T]) -> tuple[UnreachableDataplane, ...]:
|
||||
return tuple(UnreachableDataplane(id=miss.dataplane.id, reason=miss.reason) for miss in gathered.missed)
|
||||
|
||||
|
||||
@router.get("/models", response_model=FleetModelsResponse)
|
||||
async def fleet_models(
|
||||
_principal: FleetRead,
|
||||
dataplanes: Selected,
|
||||
caller: Annotated[Caller, Depends(get_caller)],
|
||||
client: Annotated[DataplaneClient, Depends(get_dataplane_client)],
|
||||
fanout: Annotated[Fanout, Depends(get_fanout)],
|
||||
) -> FleetModelsResponse:
|
||||
"""Every deployment across the fleet, plus which dataplanes serve each public model name."""
|
||||
gathered: Final[Gathered[tuple[FleetModel, ...]]] = _or_503(
|
||||
await fanout.gather(dataplanes, lambda dataplane: fetch_models(client, dataplane, caller))
|
||||
)
|
||||
items: Final = _flatten(gathered)
|
||||
return FleetModelsResponse(
|
||||
items=items, groups=group_models(items), partial=gathered.partial, unreachable=_unreachable(gathered)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/teams", response_model=FleetTeamsResponse)
|
||||
async def fleet_teams(
|
||||
_principal: FleetRead,
|
||||
dataplanes: Selected,
|
||||
caller: Annotated[Caller, Depends(get_caller)],
|
||||
client: Annotated[DataplaneClient, Depends(get_dataplane_client)],
|
||||
fanout: Annotated[Fanout, Depends(get_fanout)],
|
||||
) -> FleetTeamsResponse:
|
||||
"""Every team the caller may see on each dataplane."""
|
||||
gathered: Final[Gathered[tuple[FleetTeam, ...]]] = _or_503(
|
||||
await fanout.gather(dataplanes, lambda dataplane: fetch_teams(client, dataplane, caller))
|
||||
)
|
||||
return FleetTeamsResponse(items=_flatten(gathered), partial=gathered.partial, unreachable=_unreachable(gathered))
|
||||
137
litellm/proxy/control_plane/fleet/service.py
Normal file
137
litellm/proxy/control_plane/fleet/service.py
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
"""Fleet inventory: the union of what every dataplane serves, read through the caller's own identity."""
|
||||
|
||||
from http import HTTPStatus
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.proxy.control_plane.common.fanout import Failure, parse_reply
|
||||
from litellm.proxy.control_plane.dataplanes.client import CallResult, DataplaneClient, DataplaneReply
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller
|
||||
from litellm.proxy.control_plane.dataplanes.identity import USER_INFO_PATH, parse_user_info, resolved_user_id
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
from litellm.types.proxy.control_plane.fleet import FleetModel, FleetModelGroup, FleetTeam
|
||||
|
||||
MODEL_INFO_PATH: Final = "/model/info"
|
||||
TEAM_LIST_PATH: Final = "/team/list"
|
||||
|
||||
|
||||
class _Lenient(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
|
||||
class _LitellmParams(_Lenient):
|
||||
model: str | None = None
|
||||
custom_llm_provider: str | None = None
|
||||
|
||||
|
||||
class _ModelInfo(_Lenient):
|
||||
id: str | None = None
|
||||
litellm_provider: str | None = None
|
||||
mode: str | None = None
|
||||
|
||||
|
||||
class _ModelEntry(_Lenient):
|
||||
model_name: str
|
||||
litellm_params: _LitellmParams = _LitellmParams()
|
||||
model_info: _ModelInfo = _ModelInfo()
|
||||
|
||||
|
||||
class _ModelInfoBody(_Lenient):
|
||||
data: tuple[_ModelEntry, ...] | _ModelEntry = ()
|
||||
|
||||
|
||||
class _TeamEntry(_Lenient):
|
||||
team_id: str
|
||||
team_alias: str | None = None
|
||||
organization_id: str | None = None
|
||||
spend: float | None = None
|
||||
max_budget: float | None = None
|
||||
members_with_roles: tuple[object, ...] = ()
|
||||
models: tuple[str, ...] = ()
|
||||
blocked: bool = False
|
||||
|
||||
|
||||
_MODELS: Final = TypeAdapter(_ModelInfoBody)
|
||||
_TEAMS: Final = TypeAdapter(tuple[_TeamEntry, ...])
|
||||
|
||||
|
||||
async def fetch_models(
|
||||
client: DataplaneClient, dataplane: Dataplane, caller: Caller
|
||||
) -> tuple[FleetModel, ...] | Failure:
|
||||
body: Final = parse_reply(await client.get(dataplane, MODEL_INFO_PATH, caller), _MODELS)
|
||||
if isinstance(body, Failure):
|
||||
return body
|
||||
entries: Final = body.data if isinstance(body.data, tuple) else (body.data,)
|
||||
return tuple(
|
||||
FleetModel(
|
||||
dataplane_id=dataplane.id,
|
||||
model_name=entry.model_name,
|
||||
litellm_model=entry.litellm_params.model,
|
||||
provider=entry.litellm_params.custom_llm_provider or entry.model_info.litellm_provider,
|
||||
model_id=entry.model_info.id,
|
||||
mode=entry.model_info.mode,
|
||||
)
|
||||
for entry in entries
|
||||
)
|
||||
|
||||
|
||||
NOT_ALLOWED_TO_LIST_ALL: Final = frozenset({HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN})
|
||||
|
||||
|
||||
def _refused(result: CallResult) -> bool:
|
||||
return isinstance(result, DataplaneReply) and result.status in NOT_ALLOWED_TO_LIST_ALL
|
||||
|
||||
|
||||
async def _list_teams(client: DataplaneClient, dataplane: Dataplane, caller: Caller) -> CallResult:
|
||||
"""Every team when the dataplane lets the caller see them all, else the teams the caller belongs to."""
|
||||
everything: Final = await client.get(dataplane, TEAM_LIST_PATH, caller)
|
||||
if not _refused(everything):
|
||||
return everything
|
||||
identity: Final = await client.get(dataplane, USER_INFO_PATH, caller)
|
||||
match identity:
|
||||
case DataplaneReply(status=status, body=body) if status < HTTPStatus.BAD_REQUEST:
|
||||
info: Final = parse_user_info(body)
|
||||
user_id: Final = None if info is None else resolved_user_id(info)
|
||||
if user_id is None:
|
||||
return everything
|
||||
own: Final = MappingProxyType({"user_id": user_id})
|
||||
return await client.get(dataplane, TEAM_LIST_PATH, caller, params=tuple(own.items()))
|
||||
case DataplaneReply(status=HTTPStatus.NOT_FOUND, credential=credential):
|
||||
return DataplaneReply(HTTPStatus.OK, (), credential)
|
||||
case _:
|
||||
return everything
|
||||
|
||||
|
||||
async def fetch_teams(client: DataplaneClient, dataplane: Dataplane, caller: Caller) -> tuple[FleetTeam, ...] | Failure:
|
||||
teams: Final = parse_reply(await _list_teams(client, dataplane, caller), _TEAMS)
|
||||
if isinstance(teams, Failure):
|
||||
return teams
|
||||
return tuple(
|
||||
FleetTeam(
|
||||
dataplane_id=dataplane.id,
|
||||
team_id=team.team_id,
|
||||
team_alias=team.team_alias,
|
||||
organization_id=team.organization_id,
|
||||
spend=team.spend or 0.0,
|
||||
max_budget=team.max_budget,
|
||||
members=len(team.members_with_roles),
|
||||
models=team.models,
|
||||
blocked=team.blocked,
|
||||
)
|
||||
for team in teams
|
||||
)
|
||||
|
||||
|
||||
def group_models(models: tuple[FleetModel, ...]) -> tuple[FleetModelGroup, ...]:
|
||||
by_name: Final = sorted(models, key=_model_name)
|
||||
return tuple(
|
||||
FleetModelGroup(model_name=name, dataplanes=tuple(sorted({model.dataplane_id for model in served})))
|
||||
for name, served in groupby(by_name, key=_model_name)
|
||||
)
|
||||
|
||||
|
||||
def _model_name(model: FleetModel) -> str:
|
||||
return model.model_name
|
||||
|
|
@ -110,7 +110,8 @@ class ControlPlaneSettings(BaseSettings):
|
|||
dataplanes_file: str | None = None
|
||||
dataplane_timeout_seconds: float = DEFAULT_DATAPLANE_TIMEOUT_SECONDS
|
||||
dataplane_retry_attempts: int = 3
|
||||
fanout_concurrency: int = 16
|
||||
fanout_concurrency: int = Field(default=16, ge=1)
|
||||
usage_max_window_days: int = Field(default=366, ge=1)
|
||||
api_prefix: str = "/api/v1"
|
||||
|
||||
@field_validator("master_key")
|
||||
|
|
|
|||
0
litellm/proxy/control_plane/usage/__init__.py
Normal file
0
litellm/proxy/control_plane/usage/__init__.py
Normal file
170
litellm/proxy/control_plane/usage/router.py
Normal file
170
litellm/proxy/control_plane/usage/router.py
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
from datetime import UTC, date, datetime, timedelta
|
||||
from typing import Annotated, Final, TypeAlias
|
||||
|
||||
from fastapi import APIRouter, Depends, Query, Security
|
||||
|
||||
from litellm.proxy.control_plane.auth.dependencies import require_scopes
|
||||
from litellm.proxy.control_plane.auth.rbac import Principal, Scope
|
||||
from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, bad_request, service_unavailable
|
||||
from litellm.proxy.control_plane.common.fanout import Fanout, Gathered, describe_missed
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller
|
||||
from litellm.proxy.control_plane.dataplanes.dependencies import (
|
||||
get_caller,
|
||||
get_control_plane_settings,
|
||||
get_dataplane_client,
|
||||
get_fanout,
|
||||
selected_dataplanes,
|
||||
)
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
from litellm.proxy.control_plane.usage.service import (
|
||||
DataplaneActivity,
|
||||
DataplaneDays,
|
||||
UsageWindow,
|
||||
daily_rows,
|
||||
dataplane_totals,
|
||||
fetch_daily_activity,
|
||||
ranked,
|
||||
totals,
|
||||
)
|
||||
from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane
|
||||
from litellm.types.proxy.control_plane.usage import (
|
||||
DailyUsageResponse,
|
||||
DailyUsageRow,
|
||||
DataplaneUsage,
|
||||
TopUsageResponse,
|
||||
UsageGroupBy,
|
||||
UsageRank,
|
||||
UsageSummaryResponse,
|
||||
UsageWindowResponse,
|
||||
)
|
||||
|
||||
DEFAULT_WINDOW_DAYS: Final = 30
|
||||
|
||||
router: Final = APIRouter(prefix="/usage", tags=["usage"])
|
||||
|
||||
UsageRead: TypeAlias = Annotated[Principal, Security(require_scopes, scopes=[Scope.USAGE_READ.value])]
|
||||
Selected: TypeAlias = Annotated[tuple[Dataplane, ...], Depends(selected_dataplanes)]
|
||||
GroupBy: TypeAlias = Annotated[UsageGroupBy, Query(description="Dimension to split the fleet's usage by.")]
|
||||
|
||||
|
||||
def usage_window(
|
||||
settings: Annotated[ControlPlaneSettings, Depends(get_control_plane_settings)],
|
||||
start_date: Annotated[date | None, Query(description="First day, inclusive. Defaults to 30 days ago.")] = None,
|
||||
end_date: Annotated[date | None, Query(description="Last day, inclusive. Defaults to today (UTC).")] = None,
|
||||
) -> UsageWindow:
|
||||
end: Final = end_date if end_date is not None else datetime.now(UTC).date()
|
||||
default_window_days: Final = min(DEFAULT_WINDOW_DAYS, settings.usage_max_window_days)
|
||||
start: Final = start_date if start_date is not None else end - timedelta(days=default_window_days - 1)
|
||||
if start > end:
|
||||
raise ControlPlaneException(bad_request(f"start_date {start} is after end_date {end}."))
|
||||
if (end - start).days + 1 > settings.usage_max_window_days:
|
||||
raise ControlPlaneException(bad_request(f"Usage windows may not exceed {settings.usage_max_window_days} days."))
|
||||
return UsageWindow(start=start, end=end)
|
||||
|
||||
|
||||
Window: TypeAlias = Annotated[UsageWindow, Depends(usage_window)]
|
||||
|
||||
|
||||
async def _gather_days(
|
||||
fanout: Fanout, client: DataplaneClient, caller: Caller, dataplanes: tuple[Dataplane, ...], window: UsageWindow
|
||||
) -> Gathered[DataplaneActivity]:
|
||||
gathered: Final = await fanout.gather(
|
||||
dataplanes, lambda dataplane: fetch_daily_activity(client, dataplane, caller, window)
|
||||
)
|
||||
if gathered.nothing_answered:
|
||||
raise ControlPlaneException(service_unavailable(f"No dataplane answered: {describe_missed(gathered.missed)}"))
|
||||
return gathered
|
||||
|
||||
|
||||
def _fleet(gathered: Gathered[DataplaneActivity]) -> tuple[DataplaneDays, ...]:
|
||||
return tuple(DataplaneDays(reached.dataplane.id, reached.value.days) for reached in gathered.reached)
|
||||
|
||||
|
||||
def _unreachable(gathered: Gathered[DataplaneActivity]) -> tuple[UnreachableDataplane, ...]:
|
||||
return tuple(UnreachableDataplane(id=miss.dataplane.id, reason=miss.reason) for miss in gathered.missed)
|
||||
|
||||
|
||||
def _truncated_dataplanes(gathered: Gathered[DataplaneActivity]) -> tuple[str, ...]:
|
||||
return tuple(reached.dataplane.id for reached in gathered.reached if reached.value.keys_truncated)
|
||||
|
||||
|
||||
def _window(window: UsageWindow) -> UsageWindowResponse:
|
||||
return UsageWindowResponse(start_date=window.start, end_date=window.end)
|
||||
|
||||
|
||||
@router.get("/summary", response_model=UsageSummaryResponse)
|
||||
async def usage_summary(
|
||||
_principal: UsageRead,
|
||||
dataplanes: Selected,
|
||||
window: Window,
|
||||
caller: Annotated[Caller, Depends(get_caller)],
|
||||
client: Annotated[DataplaneClient, Depends(get_dataplane_client)],
|
||||
fanout: Annotated[Fanout, Depends(get_fanout)],
|
||||
) -> UsageSummaryResponse:
|
||||
"""Spend and token totals across every dataplane the caller is bound to."""
|
||||
gathered: Final = await _gather_days(fanout, client, caller, dataplanes, window)
|
||||
fleet: Final = _fleet(gathered)
|
||||
truncated: Final = _truncated_dataplanes(gathered)
|
||||
return UsageSummaryResponse(
|
||||
window=_window(window),
|
||||
totals=totals(fleet),
|
||||
by_dataplane=tuple(DataplaneUsage(id=days.dataplane_id, metrics=dataplane_totals(days)) for days in fleet),
|
||||
partial=gathered.partial or bool(truncated),
|
||||
unreachable=_unreachable(gathered),
|
||||
truncated_dataplanes=truncated,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/daily", response_model=DailyUsageResponse)
|
||||
async def usage_daily(
|
||||
_principal: UsageRead,
|
||||
dataplanes: Selected,
|
||||
window: Window,
|
||||
caller: Annotated[Caller, Depends(get_caller)],
|
||||
client: Annotated[DataplaneClient, Depends(get_dataplane_client)],
|
||||
fanout: Annotated[Fanout, Depends(get_fanout)],
|
||||
group_by: GroupBy = "dataplane",
|
||||
) -> DailyUsageResponse:
|
||||
"""Per-day usage across the fleet, split by dataplane, model, provider, key, team or endpoint."""
|
||||
gathered: Final = await _gather_days(fanout, client, caller, dataplanes, window)
|
||||
truncated: Final = _truncated_dataplanes(gathered)
|
||||
return DailyUsageResponse(
|
||||
window=_window(window),
|
||||
group_by=group_by,
|
||||
rows=tuple(
|
||||
DailyUsageRow(date=day, group=group, metrics=metrics)
|
||||
for (day, group), metrics in daily_rows(_fleet(gathered), group_by)
|
||||
),
|
||||
partial=gathered.partial or bool(truncated),
|
||||
unreachable=_unreachable(gathered),
|
||||
truncated_dataplanes=truncated,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/top", response_model=TopUsageResponse)
|
||||
async def usage_top(
|
||||
_principal: UsageRead,
|
||||
dataplanes: Selected,
|
||||
window: Window,
|
||||
caller: Annotated[Caller, Depends(get_caller)],
|
||||
client: Annotated[DataplaneClient, Depends(get_dataplane_client)],
|
||||
fanout: Annotated[Fanout, Depends(get_fanout)],
|
||||
group_by: GroupBy = "model",
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 10,
|
||||
) -> TopUsageResponse:
|
||||
"""The biggest spenders in the window along one dimension, summed across the fleet."""
|
||||
gathered: Final = await _gather_days(fanout, client, caller, dataplanes, window)
|
||||
truncated: Final = _truncated_dataplanes(gathered)
|
||||
return TopUsageResponse(
|
||||
window=_window(window),
|
||||
group_by=group_by,
|
||||
limit=limit,
|
||||
items=tuple(
|
||||
UsageRank(group=group, metrics=metrics) for group, metrics in ranked(_fleet(gathered), group_by, limit)
|
||||
),
|
||||
partial=gathered.partial or bool(truncated),
|
||||
unreachable=_unreachable(gathered),
|
||||
truncated_dataplanes=truncated,
|
||||
)
|
||||
163
litellm/proxy/control_plane/usage/service.py
Normal file
163
litellm/proxy/control_plane/usage/service.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
"""Fleet-wide usage: fetch each dataplane's daily activity, then pure reducers over SpendMetrics."""
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import date
|
||||
from functools import reduce
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_MAX
|
||||
from litellm.proxy.control_plane.common.fanout import Failure, parse_reply
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
from litellm.types.proxy.control_plane.usage import UsageGroupBy
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
DailySpendData,
|
||||
SpendAnalyticsPaginatedResponse,
|
||||
SpendMetrics,
|
||||
)
|
||||
|
||||
DAILY_ACTIVITY_PATH: Final = "/user/daily/activity/aggregated"
|
||||
NO_TEAM: Final = "(no team)"
|
||||
|
||||
_DAILY: Final = TypeAdapter(SpendAnalyticsPaginatedResponse)
|
||||
_NUMBERS: Final = TypeAdapter(Mapping[str, float])
|
||||
|
||||
DatedGroup: TypeAlias = tuple[date, str]
|
||||
Grouped: TypeAlias = tuple[tuple[DatedGroup, SpendMetrics], ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageWindow:
|
||||
start: date
|
||||
end: date
|
||||
|
||||
def params(self) -> Mapping[str, str]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"start_date": self.start.isoformat(),
|
||||
"end_date": self.end.isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DataplaneDays:
|
||||
dataplane_id: str
|
||||
days: tuple[DailySpendData, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DataplaneActivity:
|
||||
days: tuple[DailySpendData, ...]
|
||||
keys_truncated: bool
|
||||
|
||||
|
||||
async def fetch_daily_activity(
|
||||
client: DataplaneClient, dataplane: Dataplane, caller: Caller, window: UsageWindow
|
||||
) -> DataplaneActivity | Failure:
|
||||
parsed: Final = parse_reply(
|
||||
await client.get(
|
||||
dataplane,
|
||||
DAILY_ACTIVITY_PATH,
|
||||
caller,
|
||||
params=tuple(window.params().items()) + (("api_key_limit", str(USAGE_TOP_API_KEYS_MAX)),),
|
||||
),
|
||||
_DAILY,
|
||||
)
|
||||
if isinstance(parsed, Failure):
|
||||
return parsed
|
||||
metadata: Final = parsed.metadata
|
||||
keys_truncated: Final = (
|
||||
metadata.api_key_limit is not None
|
||||
and metadata.total_api_keys is not None
|
||||
and metadata.total_api_keys > metadata.api_key_limit
|
||||
)
|
||||
return DataplaneActivity(days=tuple(parsed.results), keys_truncated=keys_truncated)
|
||||
|
||||
|
||||
def add_metrics(left: SpendMetrics, right: SpendMetrics) -> SpendMetrics:
|
||||
first: Final = _NUMBERS.validate_python(left.model_dump())
|
||||
second: Final = _NUMBERS.validate_python(right.model_dump())
|
||||
return SpendMetrics.model_validate({name: value + second[name] for name, value in first.items()})
|
||||
|
||||
|
||||
def sum_metrics(metrics: Iterable[SpendMetrics]) -> SpendMetrics:
|
||||
return reduce(add_metrics, metrics, SpendMetrics())
|
||||
|
||||
|
||||
def _key_label(api_key: str, alias: str | None) -> str:
|
||||
return alias if alias else api_key
|
||||
|
||||
|
||||
def group_day(day: DailySpendData, group_by: UsageGroupBy, dataplane_id: str) -> Mapping[str, SpendMetrics]:
|
||||
breakdown: Final = day.breakdown
|
||||
match group_by:
|
||||
case "dataplane":
|
||||
return MappingProxyType({dataplane_id: day.metrics})
|
||||
case "model":
|
||||
return MappingProxyType({name: entry.metrics for name, entry in breakdown.models.items()})
|
||||
case "model_group":
|
||||
return MappingProxyType({name: entry.metrics for name, entry in breakdown.model_groups.items()})
|
||||
case "provider":
|
||||
return MappingProxyType({name: entry.metrics for name, entry in breakdown.providers.items()})
|
||||
case "endpoint":
|
||||
return MappingProxyType({name: entry.metrics for name, entry in breakdown.endpoints.items()})
|
||||
case "api_key":
|
||||
by_label: Final = tuple(
|
||||
(_key_label(key, entry.metadata.key_alias), entry.metrics) for key, entry in breakdown.api_keys.items()
|
||||
)
|
||||
return MappingProxyType({label: metrics for (_, label), metrics in _merge_by_key(by_label, day.date)})
|
||||
case "team":
|
||||
by_team: Final = tuple(
|
||||
(entry.metadata.team_id or NO_TEAM, entry.metrics) for entry in breakdown.api_keys.values()
|
||||
)
|
||||
return MappingProxyType({team: metrics for (_, team), metrics in _merge_by_key(by_team, day.date)})
|
||||
|
||||
|
||||
def _merge_by_key(entries: Iterable[tuple[str, SpendMetrics]], day: date) -> Grouped:
|
||||
return merge_grouped(((day, group), metrics) for group, metrics in entries)
|
||||
|
||||
|
||||
def merge_grouped(entries: Iterable[tuple[DatedGroup, SpendMetrics]]) -> Grouped:
|
||||
ordered: Final = sorted(entries, key=_dated_group)
|
||||
return tuple(
|
||||
(key, sum_metrics(metrics for _, metrics in bucket)) for key, bucket in groupby(ordered, key=_dated_group)
|
||||
)
|
||||
|
||||
|
||||
def _dated_group(entry: tuple[DatedGroup, SpendMetrics]) -> DatedGroup:
|
||||
return entry[0]
|
||||
|
||||
|
||||
def _dated_groups(days: DataplaneDays, group_by: UsageGroupBy) -> Iterable[tuple[DatedGroup, SpendMetrics]]:
|
||||
for day in days.days:
|
||||
for group, metrics in group_day(day, group_by, days.dataplane_id).items():
|
||||
yield (day.date, group), metrics
|
||||
|
||||
|
||||
def daily_rows(fleet: Iterable[DataplaneDays], group_by: UsageGroupBy) -> Grouped:
|
||||
return merge_grouped(
|
||||
entry for days in fleet for entry in _dated_groups(days, group_by)
|
||||
) # comprehension-ok: flatten
|
||||
|
||||
|
||||
def ranked(fleet: Iterable[DataplaneDays], group_by: UsageGroupBy, limit: int) -> tuple[tuple[str, SpendMetrics], ...]:
|
||||
per_group: Final = merge_grouped(
|
||||
((date.min, group), metrics) for (_, group), metrics in daily_rows(fleet, group_by)
|
||||
)
|
||||
by_spend: Final = sorted(per_group, key=lambda item: item[1].spend, reverse=True)
|
||||
return tuple((group, metrics) for (_, group), metrics in by_spend[:limit])
|
||||
|
||||
|
||||
def totals(fleet: Iterable[DataplaneDays]) -> SpendMetrics:
|
||||
return sum_metrics(day.metrics for days in fleet for day in days.days) # comprehension-ok: flatten
|
||||
|
||||
|
||||
def dataplane_totals(days: DataplaneDays) -> SpendMetrics:
|
||||
return sum_metrics(day.metrics for day in days.days)
|
||||
|
|
@ -42,3 +42,12 @@ class DataplaneIdentityResponse(LiteLLMBaseModel):
|
|||
user_id: str | None
|
||||
user_email: str | None
|
||||
user_role: str | None
|
||||
|
||||
|
||||
class UnreachableDataplane(LiteLLMBaseModel):
|
||||
"""A dataplane that did not answer a fleet-wide call, named so the caller knows what the numbers miss."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
reason: str
|
||||
|
|
|
|||
55
litellm/types/proxy/control_plane/fleet.py
Normal file
55
litellm/types/proxy/control_plane/fleet.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
from pydantic import ConfigDict
|
||||
|
||||
from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane
|
||||
from litellm.types.utils import LiteLLMBaseModel
|
||||
|
||||
|
||||
class FleetModel(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
dataplane_id: str
|
||||
model_name: str
|
||||
litellm_model: str | None
|
||||
provider: str | None
|
||||
model_id: str | None
|
||||
mode: str | None
|
||||
|
||||
|
||||
class FleetModelGroup(LiteLLMBaseModel):
|
||||
"""One public model name and every dataplane that serves it."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
model_name: str
|
||||
dataplanes: tuple[str, ...]
|
||||
|
||||
|
||||
class FleetModelsResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
items: tuple[FleetModel, ...]
|
||||
groups: tuple[FleetModelGroup, ...]
|
||||
partial: bool
|
||||
unreachable: tuple[UnreachableDataplane, ...]
|
||||
|
||||
|
||||
class FleetTeam(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
dataplane_id: str
|
||||
team_id: str
|
||||
team_alias: str | None
|
||||
organization_id: str | None
|
||||
spend: float
|
||||
max_budget: float | None
|
||||
members: int
|
||||
models: tuple[str, ...]
|
||||
blocked: bool
|
||||
|
||||
|
||||
class FleetTeamsResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
items: tuple[FleetTeam, ...]
|
||||
partial: bool
|
||||
unreachable: tuple[UnreachableDataplane, ...]
|
||||
75
litellm/types/proxy/control_plane/usage.py
Normal file
75
litellm/types/proxy/control_plane/usage.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
from datetime import date
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from pydantic import ConfigDict
|
||||
|
||||
from litellm.types.proxy.control_plane.dataplanes import UnreachableDataplane
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import SpendMetrics
|
||||
from litellm.types.utils import LiteLLMBaseModel
|
||||
|
||||
UsageGroupBy: TypeAlias = Literal["dataplane", "model", "model_group", "provider", "api_key", "team", "endpoint"]
|
||||
|
||||
|
||||
class UsageWindowResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
start_date: date
|
||||
end_date: date
|
||||
|
||||
|
||||
class DataplaneUsage(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
metrics: SpendMetrics
|
||||
|
||||
|
||||
class UsageSummaryResponse(LiteLLMBaseModel):
|
||||
"""Fleet-wide totals for the window, with the per-dataplane split and whoever did not answer."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
window: UsageWindowResponse
|
||||
totals: SpendMetrics
|
||||
by_dataplane: tuple[DataplaneUsage, ...]
|
||||
partial: bool
|
||||
unreachable: tuple[UnreachableDataplane, ...]
|
||||
truncated_dataplanes: tuple[str, ...]
|
||||
|
||||
|
||||
class DailyUsageRow(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
date: date
|
||||
group: str
|
||||
metrics: SpendMetrics
|
||||
|
||||
|
||||
class DailyUsageResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
window: UsageWindowResponse
|
||||
group_by: UsageGroupBy
|
||||
rows: tuple[DailyUsageRow, ...]
|
||||
partial: bool
|
||||
unreachable: tuple[UnreachableDataplane, ...]
|
||||
truncated_dataplanes: tuple[str, ...]
|
||||
|
||||
|
||||
class UsageRank(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
group: str
|
||||
metrics: SpendMetrics
|
||||
|
||||
|
||||
class TopUsageResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
window: UsageWindowResponse
|
||||
group_by: UsageGroupBy
|
||||
limit: int
|
||||
items: tuple[UsageRank, ...]
|
||||
partial: bool
|
||||
unreachable: tuple[UnreachableDataplane, ...]
|
||||
truncated_dataplanes: tuple[str, ...]
|
||||
104
tests/unit/proxy/control_plane/common/test_fanout.py
Normal file
104
tests/unit/proxy/control_plane/common/test_fanout.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.proxy.control_plane.common.fanout import Failure, Fanout, Missed, parse_reply
|
||||
from litellm.proxy.control_plane.dataplanes.client import CallResult, DataplaneReply, DataplaneUnreachable
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller, CredentialUnavailable, IssuedCredential
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
|
||||
|
||||
class _NoCredential:
|
||||
async def headers(self, caller: Caller) -> IssuedCredential:
|
||||
raise AssertionError("fanout never issues credentials itself")
|
||||
|
||||
|
||||
def _dataplane(dataplane_id: str) -> Dataplane:
|
||||
return Dataplane(dataplane_id, dataplane_id, f"https://{dataplane_id}.test", _NoCredential(), ())
|
||||
|
||||
|
||||
@dataclass
|
||||
class _InFlight: # mutable-ok: test probe counting concurrent fetches
|
||||
now: int = 0
|
||||
peak: int = 0
|
||||
seen: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
async def test_gather_bounds_concurrency_and_keeps_dataplane_order() -> None:
|
||||
probe: Final = _InFlight()
|
||||
|
||||
async def fetch(dataplane: Dataplane) -> str | Failure:
|
||||
probe.now += 1
|
||||
probe.peak = max(probe.peak, probe.now)
|
||||
await asyncio.sleep(0)
|
||||
probe.now -= 1
|
||||
probe.seen.append(dataplane.id)
|
||||
return dataplane.id.upper()
|
||||
|
||||
dataplanes: Final = tuple(_dataplane(name) for name in ("a", "b", "c", "d", "e"))
|
||||
gathered: Final = await Fanout(concurrency=2).gather(dataplanes, fetch)
|
||||
assert probe.peak == 2, f"semaphore of 2 let {probe.peak} fetches run at once"
|
||||
assert [reached.value for reached in gathered.reached] == ["A", "B", "C", "D", "E"]
|
||||
assert gathered.missed == () and not gathered.partial
|
||||
|
||||
|
||||
async def test_concurrent_gathers_share_the_configured_concurrency_limit() -> None:
|
||||
probe: Final = _InFlight()
|
||||
fanout: Final = Fanout(concurrency=2)
|
||||
dataplanes: Final = tuple(_dataplane(name) for name in ("a", "b", "c"))
|
||||
|
||||
async def fetch(dataplane: Dataplane) -> str:
|
||||
probe.now += 1
|
||||
probe.peak = max(probe.peak, probe.now)
|
||||
await asyncio.sleep(0)
|
||||
probe.now -= 1
|
||||
return dataplane.id
|
||||
|
||||
await asyncio.gather(fanout.gather(dataplanes, fetch), fanout.gather(dataplanes, fetch))
|
||||
assert probe.peak == 2, f"two concurrent requests ran {probe.peak} fetches at once"
|
||||
|
||||
|
||||
async def test_a_failing_dataplane_is_a_missed_value_and_marks_the_result_partial() -> None:
|
||||
async def fetch(dataplane: Dataplane) -> int | Failure:
|
||||
return Failure("boom") if dataplane.id == "b" else 1
|
||||
|
||||
gathered: Final = await Fanout(concurrency=4).gather(tuple(_dataplane(name) for name in "abc"), fetch)
|
||||
assert [reached.dataplane.id for reached in gathered.reached] == ["a", "c"]
|
||||
assert gathered.missed == (Missed(gathered.missed[0].dataplane, "boom"),)
|
||||
assert gathered.missed[0].dataplane.id == "b"
|
||||
assert gathered.partial and not gathered.nothing_answered
|
||||
|
||||
|
||||
async def test_nothing_answered_only_when_every_dataplane_missed() -> None:
|
||||
async def fetch(dataplane: Dataplane) -> int | Failure:
|
||||
return Failure("down")
|
||||
|
||||
all_down: Final = await Fanout(concurrency=1).gather((_dataplane("a"),), fetch)
|
||||
none_asked: Final = await Fanout(concurrency=1).gather((), fetch)
|
||||
assert all_down.nothing_answered and not none_asked.nothing_answered
|
||||
|
||||
|
||||
class _Shape(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
n: int
|
||||
|
||||
|
||||
_SHAPE: Final = TypeAdapter(_Shape)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("result", "expected"),
|
||||
[
|
||||
(DataplaneReply(200, {"n": 3}, "service_token"), _Shape(n=3)),
|
||||
(DataplaneReply(200, {"n": "x"}, "service_token"), Failure("unexpected response shape: 1 validation error(s)")),
|
||||
(DataplaneReply(403, {"error": "no"}, "oidc_exchange"), Failure("HTTP 403 with the oidc_exchange credential")),
|
||||
(DataplaneUnreachable("ConnectError", 3), Failure("unreachable after 3 attempt(s): ConnectError")),
|
||||
(CredentialUnavailable("service_token", "no token"), Failure("no usable credential: no token")),
|
||||
],
|
||||
)
|
||||
def test_parse_reply_turns_every_bad_outcome_into_a_failure(result: CallResult, expected: _Shape | Failure) -> None:
|
||||
assert parse_reply(result, _SHAPE) == expected
|
||||
0
tests/unit/proxy/control_plane/fleet/__init__.py
Normal file
0
tests/unit/proxy/control_plane/fleet/__init__.py
Normal file
239
tests/unit/proxy/control_plane/fleet/test_router.py
Normal file
239
tests/unit/proxy/control_plane/fleet/test_router.py
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig
|
||||
from litellm.proxy.control_plane.auth.rbac import (
|
||||
USAGE_VIEWER_ROLE,
|
||||
VIEWER_ROLE,
|
||||
build_principal,
|
||||
default_policy,
|
||||
)
|
||||
from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, RetryPolicy
|
||||
from litellm.proxy.control_plane.dataplanes.credentials import Caller, OidcTokenExchange
|
||||
from litellm.proxy.control_plane.dataplanes.identity import USER_INFO_PATH
|
||||
from litellm.proxy.control_plane.dataplanes.registry import Dataplane
|
||||
from litellm.proxy.control_plane.fleet.service import fetch_teams
|
||||
from tests.unit.proxy.control_plane.support import RoutedTransport, response, rig
|
||||
|
||||
EU_MODELS: Final = "https://eu.test/model/info"
|
||||
US_MODELS: Final = "https://us.test/model/info"
|
||||
EU_TEAMS: Final = "https://eu.test/team/list"
|
||||
US_TEAMS: Final = "https://us.test/team/list"
|
||||
EU_USER_INFO: Final = f"https://eu.test{USER_INFO_PATH}"
|
||||
|
||||
|
||||
def _model(name: str, provider: str, model_id: str, upstream: str | None = None) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": upstream or f"{provider}/{name}", "custom_llm_provider": provider},
|
||||
"model_info": {"id": model_id, "mode": "chat", "db_model": True},
|
||||
}
|
||||
|
||||
|
||||
async def _exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str:
|
||||
assert config.audience == "dataplane"
|
||||
assert subject_token == "idp-token"
|
||||
return "dataplane-token"
|
||||
|
||||
|
||||
def test_fleet_models_unions_deployments_and_groups_them_by_public_name() -> None:
|
||||
client, _, minter = rig(
|
||||
{
|
||||
EU_MODELS: [
|
||||
response(200, {"data": [_model("gpt-4o", "azure", "m1"), _model("claude", "anthropic", "m2")]})
|
||||
],
|
||||
US_MODELS: [response(200, {"data": [_model("gpt-4o", "openai", "m3")]})],
|
||||
}
|
||||
)
|
||||
body: Final = client.get("/api/v1/fleet/models", headers=minter.headers()).json()
|
||||
assert [(m["dataplane_id"], m["model_name"], m["provider"], m["model_id"]) for m in body["items"]] == [
|
||||
("eu", "gpt-4o", "azure", "m1"),
|
||||
("eu", "claude", "anthropic", "m2"),
|
||||
("us", "gpt-4o", "openai", "m3"),
|
||||
]
|
||||
assert body["groups"] == [
|
||||
{"model_name": "claude", "dataplanes": ["eu"]},
|
||||
{"model_name": "gpt-4o", "dataplanes": ["eu", "us"]},
|
||||
]
|
||||
assert body["partial"] is False
|
||||
|
||||
|
||||
def test_fleet_models_accepts_a_single_deployment_object_from_a_dataplane() -> None:
|
||||
client, _, minter = rig(
|
||||
{
|
||||
EU_MODELS: [response(200, {"data": _model("gpt-4o", "azure", "m1")})],
|
||||
US_MODELS: [response(200, {"data": []})],
|
||||
}
|
||||
)
|
||||
|
||||
body: Final = client.get("/api/v1/fleet/models", headers=minter.headers()).json()
|
||||
|
||||
assert body["partial"] is False
|
||||
assert body["unreachable"] == []
|
||||
assert body["items"] == [
|
||||
{
|
||||
"dataplane_id": "eu",
|
||||
"model_name": "gpt-4o",
|
||||
"litellm_model": "azure/gpt-4o",
|
||||
"provider": "azure",
|
||||
"model_id": "m1",
|
||||
"mode": "chat",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_fleet_models_tolerates_an_unparseable_dataplane_and_reports_it() -> None:
|
||||
client, _, minter = rig(
|
||||
{
|
||||
EU_MODELS: [response(200, {"data": [_model("gpt-4o", "azure", "m1")]})],
|
||||
US_MODELS: [response(200, {"data": "x"})],
|
||||
}
|
||||
)
|
||||
body: Final = client.get("/api/v1/fleet/models", headers=minter.headers()).json()
|
||||
assert [m["model_id"] for m in body["items"]] == ["m1"] and body["partial"] is True
|
||||
assert body["unreachable"][0]["id"] == "us" and "validation error" in body["unreachable"][0]["reason"]
|
||||
|
||||
|
||||
def test_fleet_teams_counts_members_and_keeps_budget_fields() -> None:
|
||||
team: Final = {
|
||||
"team_id": "t1",
|
||||
"team_alias": "ml",
|
||||
"organization_id": None,
|
||||
"spend": 12.5,
|
||||
"max_budget": 100.0,
|
||||
"members_with_roles": [{"user_id": "a", "role": "admin"}, {"user_id": "b", "role": "user"}],
|
||||
"models": ["gpt-4o"],
|
||||
"blocked": False,
|
||||
"metadata": {"ignored": True},
|
||||
}
|
||||
client, _, minter = rig({EU_TEAMS: [response(200, [team])], US_TEAMS: [response(200, [])]})
|
||||
body: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()).json()
|
||||
assert body["items"] == [
|
||||
{
|
||||
"dataplane_id": "eu",
|
||||
"team_id": "t1",
|
||||
"team_alias": "ml",
|
||||
"organization_id": None,
|
||||
"spend": 12.5,
|
||||
"max_budget": 100.0,
|
||||
"members": 2,
|
||||
"models": ["gpt-4o"],
|
||||
"blocked": False,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_fleet_teams_falls_back_to_the_callers_own_teams_when_the_dataplane_refuses_to_list_all() -> None:
|
||||
own: Final = {"team_id": "t-own", "members_with_roles": [{"user_id": "alice"}]}
|
||||
client, transport, minter = rig(
|
||||
{
|
||||
EU_TEAMS: [
|
||||
response(401, {"detail": {"error": "Only admin users can query all teams"}}),
|
||||
response(200, [own]),
|
||||
],
|
||||
EU_USER_INFO: [response(200, {"user_id": "alice"})],
|
||||
US_TEAMS: [response(200, [])],
|
||||
}
|
||||
)
|
||||
body: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()).json()
|
||||
assert [(t["dataplane_id"], t["team_id"], t["members"]) for t in body["items"]] == [("eu", "t-own", 1)]
|
||||
assert body["partial"] is False
|
||||
eu_calls: Final = [(url, params) for url, params, _ in transport.calls if url in {EU_TEAMS, EU_USER_INFO}]
|
||||
assert eu_calls == [
|
||||
(EU_TEAMS, None),
|
||||
(EU_USER_INFO, None),
|
||||
(EU_TEAMS, (("user_id", "alice"),)),
|
||||
]
|
||||
|
||||
|
||||
def test_fleet_teams_treats_a_user_info_not_found_as_an_empty_reached_dataplane() -> None:
|
||||
client, _, minter = rig(
|
||||
{
|
||||
EU_TEAMS: [response(403)],
|
||||
EU_USER_INFO: [response(404)],
|
||||
US_TEAMS: [response(200, [])],
|
||||
}
|
||||
)
|
||||
|
||||
body: Final = client.get("/api/v1/fleet/teams", headers=minter.headers()).json()
|
||||
|
||||
assert body["items"] == []
|
||||
assert body["partial"] is False
|
||||
assert body["unreachable"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fleet_teams_uses_dataplane_identity_for_oidc_exchange_callers() -> None:
|
||||
team: Final = {"team_id": "t-own", "members_with_roles": [{"user_id": "dataplane-user"}]}
|
||||
routes: Final = {
|
||||
EU_TEAMS: iter(
|
||||
[
|
||||
response(401, {"detail": {"error": "Only admin users can query all teams"}}),
|
||||
response(200, [team]),
|
||||
]
|
||||
),
|
||||
EU_USER_INFO: iter([response(200, {"user_id": "dataplane-user"})]),
|
||||
}
|
||||
transport: Final = RoutedTransport(routes)
|
||||
client: Final = DataplaneClient(transport=transport, timeout_seconds=1, retry=RetryPolicy(attempts=1))
|
||||
credential: Final = OidcTokenExchange(
|
||||
OAuthTokenExchangeConfig(
|
||||
token_endpoint="https://identity.test/token",
|
||||
client_id="control-plane",
|
||||
client_secret="secret",
|
||||
profile="rfc8693",
|
||||
audience="dataplane",
|
||||
scope=None,
|
||||
),
|
||||
_exchange,
|
||||
)
|
||||
dataplane: Final = Dataplane(
|
||||
id="eu",
|
||||
name="EU",
|
||||
url="https://eu.test",
|
||||
credential=credential,
|
||||
credential_kinds=("oidc_exchange",),
|
||||
)
|
||||
principal: Final = build_principal(
|
||||
default_policy(),
|
||||
"idp-subject",
|
||||
"idp@example.com",
|
||||
frozenset({VIEWER_ROLE}),
|
||||
None,
|
||||
frozenset({"eu"}),
|
||||
"oidc",
|
||||
)
|
||||
caller: Final = Caller(principal=principal, bearer="idp-token")
|
||||
|
||||
teams: Final = await fetch_teams(client, dataplane, caller)
|
||||
|
||||
assert isinstance(teams, tuple)
|
||||
assert tuple(team.team_id for team in teams) == ("t-own",)
|
||||
assert [(url, params) for url, params, _ in transport.calls] == [
|
||||
(EU_TEAMS, None),
|
||||
(EU_USER_INFO, None),
|
||||
(EU_TEAMS, (("user_id", "dataplane-user"),)),
|
||||
]
|
||||
assert all(
|
||||
headers is not None and headers["Authorization"] == "Bearer dataplane-token"
|
||||
for _, _, headers in transport.calls
|
||||
)
|
||||
|
||||
|
||||
def test_fleet_routes_need_fleet_read_and_503_when_nobody_answers() -> None:
|
||||
client, _, minter = rig({EU_TEAMS: [httpx.ConnectError("down")] * 3, US_TEAMS: [response(502, "bad")] * 3})
|
||||
denied: Final = client.get("/api/v1/fleet/teams", headers=minter.headers(role=USAGE_VIEWER_ROLE))
|
||||
assert denied.status_code == 403 and "fleet:read" in denied.json()["detail"]
|
||||
assert client.get("/api/v1/fleet/teams", headers=minter.headers()).status_code == 503
|
||||
|
||||
|
||||
def test_a_team_with_null_spend_lists_as_zero_instead_of_failing_the_fleet() -> None:
|
||||
fresh: Final = {"team_id": "new", "spend": None}
|
||||
client, _, minter = rig({EU_TEAMS: [response(200, [fresh])], US_TEAMS: [response(200, [])]})
|
||||
reply: Final = client.get("/api/v1/fleet/teams", headers=minter.headers())
|
||||
assert reply.status_code == 200, reply.text
|
||||
body: Final = reply.json()
|
||||
assert [(item["team_id"], item["spend"]) for item in body["items"]] == [("new", 0.0)]
|
||||
assert body["partial"] is False
|
||||
93
tests/unit/proxy/control_plane/support.py
Normal file
93
tests/unit/proxy/control_plane/support.py
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
"""Shared rig for control plane route tests: a routed fake transport and a licensed app over two dataplanes."""
|
||||
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig
|
||||
from litellm.proxy.control_plane.app import create_app
|
||||
from litellm.proxy.control_plane.auth.jwt import SessionTokenCodec
|
||||
from litellm.proxy.control_plane.auth.rbac import VIEWER_ROLE, RolePolicy, build_principal
|
||||
from litellm.proxy.control_plane.container import build_container
|
||||
from litellm.proxy.control_plane.settings import ControlPlaneSettings
|
||||
|
||||
MASTER_KEY: Final = "01234567890123456789012345678901"
|
||||
DATAPLANES: Final = {
|
||||
"eu": {"url": "https://eu.test", "service_token": {"token": "sk-eu"}},
|
||||
"us": {"url": "https://us.test", "name": "US East", "service_token": {"token": "sk-us"}},
|
||||
}
|
||||
|
||||
|
||||
class Licensed:
|
||||
def is_premium(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@dataclass
|
||||
class RoutedTransport: # mutable-ok: records calls for assertions
|
||||
"""Answers each URL from its own queue of outcomes, so fan-out order does not matter to a test."""
|
||||
|
||||
routes: Mapping[str, Iterator[httpx.Response | Exception]]
|
||||
calls: list[tuple[str, tuple[tuple[str, str], ...] | None, dict[str, str] | None]] = field(default_factory=list)
|
||||
|
||||
async def request(
|
||||
self,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
params: Sequence[tuple[str, str]] | None = None,
|
||||
json: dict[str, object] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
timeout: float,
|
||||
) -> httpx.Response:
|
||||
assert method == "GET", "fleet reads never write"
|
||||
self.calls.append((url, None if params is None else tuple(params), headers))
|
||||
outcome: Final = next(self.routes[url], None)
|
||||
if outcome is None:
|
||||
raise AssertionError(f"no scripted outcome left for {url}")
|
||||
if isinstance(outcome, Exception):
|
||||
raise outcome
|
||||
return outcome
|
||||
|
||||
|
||||
async def no_exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str:
|
||||
raise AssertionError("no oidc dataplane configured")
|
||||
|
||||
|
||||
def response(status: int, body: object = None) -> httpx.Response:
|
||||
return httpx.Response(status, json=body, request=httpx.Request("GET", "https://x"))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Minter:
|
||||
codec: SessionTokenCodec
|
||||
policy: RolePolicy
|
||||
|
||||
def headers(self, *dataplanes: str, role: str = VIEWER_ROLE) -> dict[str, str]:
|
||||
principal: Final = build_principal(
|
||||
self.policy, "alice", "alice@x.test", frozenset({role}), None, frozenset(dataplanes), "session"
|
||||
)
|
||||
return {"Authorization": f"Bearer {self.codec.mint(principal)}"}
|
||||
|
||||
|
||||
def rig(
|
||||
routes: Mapping[str, list[httpx.Response | Exception]], usage_max_window_days: int = 366
|
||||
) -> tuple[TestClient, RoutedTransport, Minter]:
|
||||
settings: Final = ControlPlaneSettings.model_validate(
|
||||
{
|
||||
"LITELLM_MASTER_KEY": MASTER_KEY,
|
||||
"dataplanes": DATAPLANES,
|
||||
"usage_max_window_days": usage_max_window_days,
|
||||
}
|
||||
)
|
||||
queues: defaultdict[str, Iterator[httpx.Response | Exception]] = defaultdict(lambda: iter(()))
|
||||
queues.update({url: iter(outcomes) for url, outcomes in routes.items()})
|
||||
transport: Final = RoutedTransport(queues)
|
||||
container: Final = build_container(settings, transport=transport, exchanger=no_exchange)
|
||||
codec: Final = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy)
|
||||
return TestClient(create_app(container, Licensed())), transport, Minter(codec, container.policy)
|
||||
|
|
@ -60,3 +60,22 @@ def test_dataplanes_environment_and_file_cannot_both_be_configured(monkeypatch:
|
|||
|
||||
with pytest.raises(ValueError, match="CONTROL_PLANE_DATAPLANES and CONTROL_PLANE_DATAPLANES_FILE"):
|
||||
load_settings()
|
||||
|
||||
|
||||
def test_fanout_concurrency_must_be_positive() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, "fanout_concurrency": 0})
|
||||
|
||||
|
||||
def test_master_key_minimum_length_uses_utf8_bytes() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "é" * 15 + "x"})
|
||||
settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "é" * 16})
|
||||
assert len(settings.master_key.get_secret_value().encode("utf-8")) == 32
|
||||
|
||||
|
||||
def test_usage_window_maximum_is_configurable() -> None:
|
||||
settings: Final = ControlPlaneSettings.model_validate(
|
||||
{"LITELLM_MASTER_KEY": MASTER_KEY, "usage_max_window_days": 12}
|
||||
)
|
||||
assert settings.usage_max_window_days == 12
|
||||
|
|
|
|||
0
tests/unit/proxy/control_plane/usage/__init__.py
Normal file
0
tests/unit/proxy/control_plane/usage/__init__.py
Normal file
185
tests/unit/proxy/control_plane/usage/test_router.py
Normal file
185
tests/unit/proxy/control_plane/usage/test_router.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
from datetime import date
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_MAX
|
||||
from litellm.proxy.control_plane.auth.rbac import USAGE_VIEWER_ROLE
|
||||
from tests.unit.proxy.control_plane.support import MASTER_KEY, response, rig
|
||||
|
||||
EU: Final = "https://eu.test/user/daily/activity/aggregated"
|
||||
US: Final = "https://us.test/user/daily/activity/aggregated"
|
||||
WINDOW: Final = {"start_date": "2026-10-01", "end_date": "2026-10-02"}
|
||||
|
||||
|
||||
def _page(
|
||||
rows: list[dict[str, object]], api_key_limit: int | None = None, total_api_keys: int | None = None
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"results": rows,
|
||||
"metadata": {
|
||||
"has_more": False,
|
||||
"page": 1,
|
||||
"total_pages": 1,
|
||||
"api_key_limit": api_key_limit,
|
||||
"total_api_keys": total_api_keys,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _row(day: str, spend: float, models: dict[str, float] | None = None) -> dict[str, object]:
|
||||
return {
|
||||
"date": day,
|
||||
"metrics": {"spend": spend, "api_requests": 1},
|
||||
"breakdown": {"models": {name: {"metrics": {"spend": cost}} for name, cost in (models or {}).items()}},
|
||||
}
|
||||
|
||||
|
||||
def test_summary_adds_both_dataplanes_and_forwards_the_window_and_caller() -> None:
|
||||
client, transport, minter = rig(
|
||||
{
|
||||
EU: [response(200, _page([_row("2026-10-01", 1.5), _row("2026-10-02", 0.5)]))],
|
||||
US: [response(200, _page([_row("2026-10-01", 2.0)]))],
|
||||
}
|
||||
)
|
||||
body: Final = client.get(
|
||||
"/api/v1/usage/summary", params=WINDOW, headers=minter.headers(role=USAGE_VIEWER_ROLE)
|
||||
).json()
|
||||
assert body["window"] == WINDOW
|
||||
assert (body["totals"]["spend"], body["totals"]["api_requests"]) == (4.0, 3)
|
||||
assert [(dp["id"], dp["metrics"]["spend"]) for dp in body["by_dataplane"]] == [("eu", 2.0), ("us", 2.0)]
|
||||
assert body["partial"] is False and body["unreachable"] == []
|
||||
url, params, headers = transport.calls[0]
|
||||
assert params == (
|
||||
("start_date", "2026-10-01"),
|
||||
("end_date", "2026-10-02"),
|
||||
("api_key_limit", str(USAGE_TOP_API_KEYS_MAX)),
|
||||
)
|
||||
assert headers is not None and headers["X-LiteLLM-Impersonate-User"] == "alice"
|
||||
|
||||
|
||||
def test_summary_does_not_report_a_dataplane_with_many_aggregated_rows_as_missed() -> None:
|
||||
rows: Final = [_row("2026-10-01", 1.0) for _ in range(1001)]
|
||||
client, transport, minter = rig({EU: [response(200, _page(rows))], US: [response(200, _page([]))]})
|
||||
body: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()).json()
|
||||
assert body["totals"]["spend"] == 1001.0
|
||||
assert body["partial"] is False and body["unreachable"] == []
|
||||
assert sum(1 for url, _, _ in transport.calls if url == EU) == 1
|
||||
|
||||
|
||||
def test_one_dataplane_down_gives_partial_numbers_and_names_it() -> None:
|
||||
client, _, minter = rig(
|
||||
{EU: [response(200, _page([_row("2026-10-01", 1.0)]))], US: [httpx.ConnectError("refused")] * 3}
|
||||
)
|
||||
body: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()).json()
|
||||
assert body["totals"]["spend"] == 1.0 and body["partial"] is True
|
||||
assert [u["id"] for u in body["unreachable"]] == ["us"]
|
||||
assert "ConnectError" in body["unreachable"][0]["reason"]
|
||||
|
||||
|
||||
def test_every_dataplane_down_is_a_503_error() -> None:
|
||||
client, _, minter = rig({EU: [response(503, "maintenance")] * 3, US: [httpx.ReadTimeout("slow")] * 3})
|
||||
down: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers())
|
||||
assert down.status_code == 503
|
||||
assert down.json()["type"] == "urn:litellm:error:service-unavailable"
|
||||
assert "eu:" in down.json()["detail"] and "us:" in down.json()["detail"]
|
||||
|
||||
|
||||
def test_binding_and_dataplane_filter_narrow_the_fan_out() -> None:
|
||||
client, transport, minter = rig(
|
||||
{
|
||||
EU: [response(200, _page([_row("2026-10-01", 1.0)]))] * 2,
|
||||
US: [response(200, _page([_row("2026-10-01", 9.0)]))],
|
||||
}
|
||||
)
|
||||
bound: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers("eu")).json()
|
||||
assert [dp["id"] for dp in bound["by_dataplane"]] == ["eu"]
|
||||
filtered: Final = client.get(
|
||||
"/api/v1/usage/summary", params={**WINDOW, "dataplane": "eu"}, headers={"Authorization": f"Bearer {MASTER_KEY}"}
|
||||
).json()
|
||||
assert [dp["id"] for dp in filtered["by_dataplane"]] == ["eu"]
|
||||
assert all(url == EU for url, _, _ in transport.calls)
|
||||
assert (
|
||||
client.get(
|
||||
"/api/v1/usage/summary", params={**WINDOW, "dataplane": "us"}, headers=minter.headers("eu")
|
||||
).status_code
|
||||
== 403
|
||||
)
|
||||
|
||||
|
||||
def test_daily_groups_by_model_across_dataplanes_and_top_ranks_them() -> None:
|
||||
client, _, minter = rig(
|
||||
{
|
||||
EU: [response(200, _page([_row("2026-10-01", 3.0, {"gpt-4o": 2.0, "claude": 1.0})]))] * 2,
|
||||
US: [response(200, _page([_row("2026-10-01", 5.0, {"gpt-4o": 5.0})]))] * 2,
|
||||
}
|
||||
)
|
||||
daily: Final = client.get(
|
||||
"/api/v1/usage/daily", params={**WINDOW, "group_by": "model"}, headers=minter.headers()
|
||||
).json()
|
||||
assert [(row["date"], row["group"], row["metrics"]["spend"]) for row in daily["rows"]] == [
|
||||
("2026-10-01", "claude", 1.0),
|
||||
("2026-10-01", "gpt-4o", 7.0),
|
||||
]
|
||||
top: Final = client.get("/api/v1/usage/top", params={**WINDOW, "limit": 1}, headers=minter.headers()).json()
|
||||
assert top["group_by"] == "model" and top["limit"] == 1
|
||||
assert [(item["group"], item["metrics"]["spend"]) for item in top["items"]] == [("gpt-4o", 7.0)]
|
||||
|
||||
|
||||
def test_usage_routes_need_the_usage_read_scope_and_default_the_window_to_a_month() -> None:
|
||||
client, transport, minter = rig({EU: [response(200, _page([]))], US: [response(200, _page([]))]})
|
||||
assert client.get("/api/v1/usage/daily", params=WINDOW).status_code == 401
|
||||
body: Final = client.get("/api/v1/usage/summary", headers=minter.headers()).json()
|
||||
start, end = date.fromisoformat(body["window"]["start_date"]), date.fromisoformat(body["window"]["end_date"])
|
||||
assert (end - start).days == 29
|
||||
|
||||
|
||||
def test_default_usage_window_respects_a_smaller_configured_maximum() -> None:
|
||||
client, _, minter = rig(
|
||||
{EU: [response(200, _page([]))], US: [response(200, _page([]))]}, usage_max_window_days=7
|
||||
)
|
||||
result: Final = client.get("/api/v1/usage/summary", headers=minter.headers())
|
||||
assert result.status_code == 200
|
||||
body: Final = result.json()
|
||||
start, end = date.fromisoformat(body["window"]["start_date"]), date.fromisoformat(body["window"]["end_date"])
|
||||
assert (end - start).days + 1 == 7
|
||||
|
||||
|
||||
def test_reversed_window_is_a_400_before_any_dataplane_is_asked() -> None:
|
||||
client, transport, minter = rig({EU: [], US: []})
|
||||
reversed_window: Final = client.get(
|
||||
"/api/v1/usage/summary",
|
||||
params={"start_date": "2026-10-03", "end_date": "2026-10-01"},
|
||||
headers=minter.headers(),
|
||||
)
|
||||
assert reversed_window.status_code == 400 and "after end_date" in reversed_window.json()["detail"]
|
||||
assert transport.calls == []
|
||||
|
||||
|
||||
def test_window_over_the_configured_maximum_is_rejected_before_fanout() -> None:
|
||||
client, transport, minter = rig({EU: [], US: []})
|
||||
too_wide: Final = client.get(
|
||||
"/api/v1/usage/summary",
|
||||
params={"start_date": "2025-09-30", "end_date": "2026-10-01"},
|
||||
headers=minter.headers(),
|
||||
)
|
||||
assert too_wide.status_code == 400
|
||||
assert transport.calls == []
|
||||
|
||||
|
||||
def test_truncated_api_key_aggregation_marks_the_dataplane_partial() -> None:
|
||||
client, transport, minter = rig(
|
||||
{
|
||||
EU: [
|
||||
response(
|
||||
200,
|
||||
_page([_row("2026-10-01", 1.0)], api_key_limit=USAGE_TOP_API_KEYS_MAX, total_api_keys=1001),
|
||||
)
|
||||
],
|
||||
US: [response(200, _page([]))],
|
||||
}
|
||||
)
|
||||
body: Final = client.get("/api/v1/usage/summary", params=WINDOW, headers=minter.headers()).json()
|
||||
assert body["partial"] is True
|
||||
assert body["truncated_dataplanes"] == ["eu"]
|
||||
assert any(("api_key_limit", str(USAGE_TOP_API_KEYS_MAX)) in params for _, params, _ in transport.calls)
|
||||
113
tests/unit/proxy/control_plane/usage/test_service.py
Normal file
113
tests/unit/proxy/control_plane/usage/test_service.py
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
from datetime import date
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.control_plane.usage.service import (
|
||||
NO_TEAM,
|
||||
DataplaneDays,
|
||||
add_metrics,
|
||||
daily_rows,
|
||||
group_day,
|
||||
ranked,
|
||||
totals,
|
||||
)
|
||||
from litellm.types.proxy.control_plane.usage import UsageGroupBy
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
BreakdownMetrics,
|
||||
DailySpendData,
|
||||
KeyMetadata,
|
||||
KeyMetricWithMetadata,
|
||||
MetricWithMetadata,
|
||||
SpendMetrics,
|
||||
)
|
||||
|
||||
D1: Final = date(2026, 10, 1)
|
||||
D2: Final = date(2026, 10, 2)
|
||||
|
||||
|
||||
def _m(spend: float, requests: int = 1) -> SpendMetrics:
|
||||
return SpendMetrics(spend=spend, api_requests=requests, successful_requests=requests, total_tokens=requests * 10)
|
||||
|
||||
|
||||
def _day(
|
||||
day: date,
|
||||
spend: float,
|
||||
models: dict[str, float] | None = None,
|
||||
keys: dict[str, tuple[float, str | None, str | None]] | None = None,
|
||||
) -> DailySpendData:
|
||||
return DailySpendData(
|
||||
date=day,
|
||||
metrics=_m(spend),
|
||||
breakdown=BreakdownMetrics(
|
||||
models={name: MetricWithMetadata(metrics=_m(cost)) for name, cost in (models or {}).items()},
|
||||
api_keys={
|
||||
key: KeyMetricWithMetadata(metrics=_m(cost), metadata=KeyMetadata(key_alias=alias, team_id=team))
|
||||
for key, (cost, alias, team) in (keys or {}).items()
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_add_metrics_sums_every_numeric_field() -> None:
|
||||
left: Final = SpendMetrics.model_validate({name: 1 for name in SpendMetrics.model_fields})
|
||||
right: Final = SpendMetrics.model_validate({name: 2 for name in SpendMetrics.model_fields})
|
||||
total: Final = add_metrics(left, right)
|
||||
assert total.model_dump() == {name: 3 for name in SpendMetrics.model_fields}
|
||||
|
||||
|
||||
def test_totals_add_up_across_dataplanes_and_days() -> None:
|
||||
fleet: Final = (
|
||||
DataplaneDays("eu", (_day(D1, 1.0), _day(D2, 2.0))),
|
||||
DataplaneDays("us", (_day(D1, 4.0),)),
|
||||
)
|
||||
assert totals(fleet) == _m(7.0, requests=3)
|
||||
|
||||
|
||||
def test_daily_rows_merge_the_same_model_on_the_same_day_from_two_dataplanes() -> None:
|
||||
fleet: Final = (
|
||||
DataplaneDays(
|
||||
"eu", (_day(D1, 3.0, models={"gpt-4o": 2.0, "claude": 1.0}), _day(D2, 1.0, models={"claude": 1.0}))
|
||||
),
|
||||
DataplaneDays("us", (_day(D1, 5.0, models={"gpt-4o": 5.0}),)),
|
||||
)
|
||||
rows: Final = daily_rows(fleet, "model")
|
||||
assert [(day, group, metrics.spend) for (day, group), metrics in rows] == [
|
||||
(D1, "claude", 1.0),
|
||||
(D1, "gpt-4o", 7.0),
|
||||
(D2, "claude", 1.0),
|
||||
]
|
||||
assert daily_rows(fleet, "dataplane") == (
|
||||
((D1, "eu"), _m(3.0)),
|
||||
((D1, "us"), _m(5.0)),
|
||||
((D2, "eu"), _m(1.0)),
|
||||
)
|
||||
|
||||
|
||||
def test_team_grouping_folds_keys_into_their_team_and_labels_keys_by_alias() -> None:
|
||||
day: Final = _day(
|
||||
D1,
|
||||
6.0,
|
||||
keys={"hash-a": (1.0, "alice-key", "t1"), "hash-b": (2.0, None, "t1"), "hash-c": (3.0, "svc", None)},
|
||||
)
|
||||
assert {team: metrics.spend for team, metrics in group_day(day, "team", "eu").items()} == {"t1": 3.0, NO_TEAM: 3.0}
|
||||
assert set(group_day(day, "api_key", "eu")) == {"alice-key", "hash-b", "svc"}
|
||||
|
||||
|
||||
def test_ranked_orders_by_fleet_spend_and_honours_the_limit() -> None:
|
||||
fleet: Final = (
|
||||
DataplaneDays("eu", (_day(D1, 0, models={"a": 1.0, "b": 5.0}), _day(D2, 0, models={"a": 1.0}))),
|
||||
DataplaneDays("us", (_day(D1, 0, models={"c": 3.0, "a": 2.0}),)),
|
||||
)
|
||||
assert [(group, metrics.spend) for group, metrics in ranked(fleet, "model", limit=2)] == [("b", 5.0), ("a", 4.0)]
|
||||
assert len(ranked(fleet, "model", limit=10)) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dimension", ["model_group", "provider", "endpoint"])
|
||||
def test_empty_breakdowns_group_to_nothing(dimension: UsageGroupBy) -> None:
|
||||
assert dict(group_day(_day(D1, 1.0), dimension, "eu")) == {}
|
||||
|
||||
|
||||
def test_two_keys_sharing_an_alias_add_up_instead_of_overwriting() -> None:
|
||||
day: Final = _day(D1, 8.0, keys={"hash-a": (3.0, "shared", None), "hash-b": (5.0, "shared", None)})
|
||||
assert {label: metrics.spend for label, metrics in group_day(day, "api_key", "eu").items()} == {"shared": 8.0}
|
||||
Loading…
Add table
Reference in a new issue