mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(usage): search team keys beyond the top-N in the Team usage view (#42857)
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
5a8ec13786
commit
c2eb549ee6
11 changed files with 635 additions and 7 deletions
|
|
@ -307,6 +307,7 @@ class KeyManagementRoutes(str, enum.Enum):
|
|||
# team usage routes
|
||||
TEAM_DAILY_ACTIVITY = "/team/daily/activity"
|
||||
TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated"
|
||||
TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH = "/team/daily/activity/aggregated/search"
|
||||
|
||||
# team spend-log viewing
|
||||
SPEND_LOGS = "/spend/logs"
|
||||
|
|
@ -673,6 +674,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
|
||||
KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value,
|
||||
KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED.value,
|
||||
KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH.value,
|
||||
KeyManagementRoutes.SPEND_LOGS.value,
|
||||
KeyManagementRoutes.SPEND_LOGS_V2.value,
|
||||
KeyManagementRoutes.KEY_RESET_SPEND.value,
|
||||
|
|
@ -717,6 +719,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/team/permissions_bulk_update",
|
||||
"/team/daily/activity",
|
||||
"/team/daily/activity/aggregated",
|
||||
"/team/daily/activity/aggregated/search",
|
||||
"/team/spend/by_user",
|
||||
# gateway request counts (SGR); deployment-wide, admin-only
|
||||
"/gateway/daily/activity",
|
||||
|
|
@ -887,6 +890,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/team/permissions_update",
|
||||
"/team/daily/activity",
|
||||
"/team/daily/activity/aggregated",
|
||||
"/team/daily/activity/aggregated/search",
|
||||
"/team/spend/by_user",
|
||||
"/team/{team_id}/members/me",
|
||||
# POST/GET the team's logging callbacks, and DELETE one of them. Every
|
||||
|
|
@ -986,6 +990,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/user/daily/activity",
|
||||
"/team/daily/activity",
|
||||
"/team/daily/activity/aggregated",
|
||||
"/team/daily/activity/aggregated/search",
|
||||
"/tag/daily/activity",
|
||||
"/tag/list",
|
||||
"/audit",
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_LIMIT
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -196,6 +197,7 @@ from litellm.repositories.verification_token_repository import (
|
|||
from litellm.router import Router
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
DailySpendMetadata,
|
||||
SpendAnalyticsPaginatedResponse,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.team_endpoints import (
|
||||
|
|
@ -204,7 +206,9 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
|
|||
BulkUpdateTeamMemberPermissionsRequest,
|
||||
BulkUpdateTeamMemberPermissionsResponse,
|
||||
GetTeamMemberPermissionsResponse,
|
||||
TeamIdSearchFilter,
|
||||
TeamIdSearchMatch,
|
||||
TeamKeyActivitySearchWhere,
|
||||
TeamListItem,
|
||||
TeamListResponse,
|
||||
TeamMemberAddResult,
|
||||
|
|
@ -6805,6 +6809,111 @@ async def get_team_daily_activity_aggregated(
|
|||
)
|
||||
|
||||
|
||||
def _team_key_search_where(*, search: str, scope: _TeamDailyActivityScope) -> TeamKeyActivitySearchWhere:
|
||||
"""Caller scoping lives inside the same Prisma where as the search term so `take`
|
||||
never trims visible matches in favour of keys the caller is not allowed to see."""
|
||||
search_or: Final = (
|
||||
{"token": search}, # mutable-ok: prisma where clause leaf
|
||||
{"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf
|
||||
{"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf
|
||||
)
|
||||
own_keys: Final = tuple(scope.api_key_filter) if isinstance(scope.api_key_filter, list) else None
|
||||
team_filter: Final[TeamIdSearchFilter | None] = (
|
||||
{ # mutable-ok: prisma where clause leaf
|
||||
"in": tuple(scope.team_ids),
|
||||
"notIn": tuple(scope.exclude_team_ids),
|
||||
}
|
||||
if scope.team_ids is not None and scope.exclude_team_ids is not None
|
||||
else {"in": tuple(scope.team_ids)} # mutable-ok: prisma where clause leaf
|
||||
if scope.team_ids is not None
|
||||
else {"notIn": tuple(scope.exclude_team_ids)} # mutable-ok: prisma where clause leaf
|
||||
if scope.exclude_team_ids is not None
|
||||
else None
|
||||
)
|
||||
if team_filter is None and own_keys is None:
|
||||
return {"OR": search_or} # mutable-ok: prisma where clause root
|
||||
if team_filter is None and own_keys is not None:
|
||||
return {"token": {"in": own_keys}, "OR": search_or} # mutable-ok: prisma where clause root
|
||||
if team_filter is not None and own_keys is None:
|
||||
return {"team_id": team_filter, "OR": search_or} # mutable-ok: prisma where clause root
|
||||
assert team_filter is not None and own_keys is not None
|
||||
return { # mutable-ok: prisma where clause root
|
||||
"team_id": team_filter,
|
||||
"token": {"in": own_keys}, # mutable-ok: prisma where clause leaf
|
||||
"OR": search_or,
|
||||
}
|
||||
|
||||
|
||||
@router.get(
|
||||
"/team/daily/activity/aggregated/search",
|
||||
response_model=SpendAnalyticsPaginatedResponse,
|
||||
tags=["team management"], # mutable-ok: FastAPI route tags shape
|
||||
)
|
||||
async def search_team_daily_activity_keys(
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
search: str = fastapi.Query(
|
||||
...,
|
||||
min_length=1,
|
||||
description="Exact token hash, or a case-insensitive substring of the key alias or owning user id",
|
||||
),
|
||||
team_ids: str | None = None,
|
||||
start_date: str | None = None,
|
||||
end_date: str | None = None,
|
||||
exclude_team_ids: str | None = None,
|
||||
timezone: int | None = None,
|
||||
) -> SpendAnalyticsPaginatedResponse:
|
||||
"""Aggregated daily team activity for the keys matching `search`, across every key the caller may
|
||||
see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend."""
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
range_error: Final = _aggregated_date_range_error(start_date, end_date)
|
||||
if range_error is not None:
|
||||
raise _daily_activity_error(status_code=400, message=range_error)
|
||||
|
||||
scope: Final = await _resolve_team_daily_activity_scope(
|
||||
team_ids=team_ids,
|
||||
exclude_team_ids=exclude_team_ids,
|
||||
api_key=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
matched_keys: Final = await _tokens_db(prisma_client).find_many(
|
||||
where=_team_key_search_where(search=search, scope=scope),
|
||||
take=USAGE_TOP_API_KEYS_LIMIT,
|
||||
order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict
|
||||
)
|
||||
tokens: Final = [key.token for key in matched_keys] # mutable-ok: get_daily_activity_aggregated takes list[str]
|
||||
if not tokens:
|
||||
return SpendAnalyticsPaginatedResponse(
|
||||
results=[], # mutable-ok: response model field shape
|
||||
metadata=DailySpendMetadata(api_key_limit=USAGE_TOP_API_KEYS_LIMIT, total_api_keys=0),
|
||||
)
|
||||
|
||||
return await get_daily_activity_aggregated(
|
||||
prisma_client=prisma_client,
|
||||
table_name="litellm_dailyteamspend",
|
||||
entity_id_field="team_id",
|
||||
entity_id=scope.team_ids,
|
||||
entity_metadata_field=scope.team_alias_metadata,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
model=None,
|
||||
api_key=tokens,
|
||||
exclude_entity_ids=scope.exclude_team_ids,
|
||||
timezone_offset_minutes=timezone,
|
||||
include_entity_breakdown=True,
|
||||
)
|
||||
|
||||
|
||||
def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str:
|
||||
team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count))
|
||||
user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else ""
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.proxy._types import (
|
||||
KeyManagementRoutes,
|
||||
|
|
@ -11,10 +13,32 @@ from litellm.proxy._types import (
|
|||
MemberDeleteRequest,
|
||||
)
|
||||
from litellm.proxy.common_utils.timezone_utils import budget_duration_error
|
||||
from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse
|
||||
|
||||
TeamIdSearchMatch = Literal["exact", "prefix"]
|
||||
|
||||
|
||||
TeamIdSearchFilter = TypedDict(
|
||||
"TeamIdSearchFilter",
|
||||
{ # mutable-ok: functional TypedDict field map
|
||||
"in": NotRequired[ReadOnly[Sequence[str]]],
|
||||
"notIn": NotRequired[ReadOnly[Sequence[str]]],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TeamKeyActivitySearchWhere(TypedDict):
|
||||
"""Prisma filter behind `/team/daily/activity/aggregated/search`: exact token hash, or key alias
|
||||
or user id containing the term, case-insensitive, narrowed to the teams and keys the caller may see."""
|
||||
|
||||
team_id: NotRequired[ReadOnly[TeamIdSearchFilter]]
|
||||
token: NotRequired[ReadOnly[Mapping[Literal["in"], Sequence[str]]]]
|
||||
OR: ReadOnly[
|
||||
tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...]
|
||||
]
|
||||
|
||||
|
||||
MAX_BULK_TEAM_MEMBER_DELETES: Final = 500
|
||||
|
||||
MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES: Final = 500
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics
|
|||
GET /tag/wau
|
||||
GET /team/daily/activity
|
||||
GET /team/daily/activity/aggregated
|
||||
GET /team/daily/activity/aggregated/search
|
||||
GET /team/spend/by_user
|
||||
GET /team/spend/report
|
||||
GET /user/daily/activity
|
||||
|
|
|
|||
128
tests/integration/spend/test_team_daily_activity_key_search.py
Normal file
128
tests/integration/spend/test_team_daily_activity_key_search.py
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from pydantic import JsonValue
|
||||
|
||||
_SEARCH_PATH: Final = "/team/daily/activity/aggregated/search"
|
||||
|
||||
|
||||
def _range_around_today() -> dict[str, str]:
|
||||
today: Final = datetime.now(timezone.utc)
|
||||
return {
|
||||
"start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d"),
|
||||
"end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d"),
|
||||
"timezone": "0",
|
||||
}
|
||||
|
||||
|
||||
def _team_key_breakdown(body: dict[str, JsonValue], team: str) -> dict[str, JsonValue]:
|
||||
results: Final = body["results"]
|
||||
assert isinstance(results, list) and len(results) == 1, body
|
||||
entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"])
|
||||
return object_value(object_value(entities[team])["api_key_breakdown"])
|
||||
|
||||
|
||||
def test_team_key_search_returns_only_the_matching_key_spend_by_alias_and_by_hash(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
team: Final = scenario.team(models=[model])
|
||||
needle_alias: Final = f"needle-{uuid.uuid4().hex}"
|
||||
needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias)
|
||||
other: Final = scenario.key(team_id=team, models=[model], key_alias=f"other-{uuid.uuid4().hex}")
|
||||
needle_digest: Final = sha256(needle.encode()).hexdigest()
|
||||
other_digest: Final = sha256(other.encode()).hexdigest()
|
||||
for key in (needle, other):
|
||||
reply: Final = gateway.chat(model, key=key, text=f"key search {uuid.uuid4().hex}")
|
||||
assert object_value(reply["usage"])["total_tokens"] == 40, reply
|
||||
daily: Final = eventually(
|
||||
lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)),
|
||||
lambda values: sorted(row["api_key"] for row in values) == sorted((needle_digest, other_digest)),
|
||||
seconds=70,
|
||||
)
|
||||
assert all(float(row["spend"]) == pytest.approx(0.06) for row in daily), daily
|
||||
for search in (needle_alias.upper(), needle_digest):
|
||||
response: Final = gateway.request(
|
||||
"GET", _SEARCH_PATH, params={"team_ids": team, "search": search, **_range_around_today()}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = object_value(response.json())
|
||||
assert object_value(body["metadata"])["total_spend"] == pytest.approx(0.06), response.text
|
||||
per_key: Final = _team_key_breakdown(body, team)
|
||||
assert set(per_key) == {needle_digest}, response.text
|
||||
assert object_value(object_value(per_key[needle_digest])["metrics"])["spend"] == pytest.approx(0.06)
|
||||
|
||||
|
||||
def test_team_key_search_is_scoped_to_the_teams_the_caller_belongs_to(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
team: Final = scenario.team(models=[model])
|
||||
needle_alias: Final = f"needle-{uuid.uuid4().hex}"
|
||||
needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias)
|
||||
needle_digest: Final = sha256(needle.encode()).hexdigest()
|
||||
reply: Final = gateway.chat(model, key=needle, text=f"key search {uuid.uuid4().hex}")
|
||||
assert object_value(reply["usage"])["total_tokens"] == 40, reply
|
||||
eventually(
|
||||
lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)),
|
||||
lambda values: [row["api_key"] for row in values] == [needle_digest],
|
||||
seconds=70,
|
||||
)
|
||||
outsider: Final = scenario.user(user_role="internal_user")
|
||||
outsider_team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": outsider, "role": "user"}])
|
||||
outsider_key: Final = scenario.key(user_id=outsider, team_id=outsider_team, models=[model])
|
||||
params: Final = {"search": needle_alias, **_range_around_today()}
|
||||
admin_view: Final = gateway.request("GET", _SEARCH_PATH, params={"team_ids": team, **params})
|
||||
assert admin_view.status_code == 200, admin_view.text
|
||||
assert set(_team_key_breakdown(object_value(admin_view.json()), team)) == {needle_digest}, admin_view.text
|
||||
own_teams_view: Final = gateway.request("GET", _SEARCH_PATH, params=params, key=outsider_key)
|
||||
assert own_teams_view.status_code == 200, own_teams_view.text
|
||||
own_teams_body: Final = object_value(own_teams_view.json())
|
||||
assert own_teams_body["results"] == [], own_teams_view.text
|
||||
assert object_value(own_teams_body["metadata"])["total_api_keys"] == 0, own_teams_view.text
|
||||
foreign_team_view: Final = gateway.request(
|
||||
"GET", _SEARCH_PATH, params={"team_ids": team, **params}, key=outsider_key
|
||||
)
|
||||
assert foreign_team_view.status_code == 404, foreign_team_view.text
|
||||
|
||||
|
||||
def test_team_key_search_excludes_teams_inside_the_where(gateway: Gateway) -> None:
|
||||
"""The dashboard always sends exclude_team_ids; a matching key in an excluded
|
||||
team with higher spend must not consume a take slot nor appear in the result."""
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
team_keep: Final = scenario.team(models=[model])
|
||||
team_drop: Final = scenario.team(models=[model])
|
||||
shared_alias: Final = f"needle-{uuid.uuid4().hex}"
|
||||
keep: Final = scenario.key(team_id=team_keep, models=[model], key_alias=f"{shared_alias}-keep")
|
||||
drop: Final = scenario.key(team_id=team_drop, models=[model], key_alias=f"{shared_alias}-drop")
|
||||
keep_digest: Final = sha256(keep.encode()).hexdigest()
|
||||
drop_digest: Final = sha256(drop.encode()).hexdigest()
|
||||
for _ in range(2):
|
||||
reply: Final = gateway.chat(model, key=drop, text=f"key search {uuid.uuid4().hex}")
|
||||
assert object_value(reply["usage"])["total_tokens"] == 40, reply
|
||||
reply = gateway.chat(model, key=keep, text=f"key search {uuid.uuid4().hex}")
|
||||
assert object_value(reply["usage"])["total_tokens"] == 40, reply
|
||||
eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id IN (%s, %s)',
|
||||
(team_keep, team_drop),
|
||||
),
|
||||
lambda values: sorted(row["api_key"] for row in values) == sorted((keep_digest, drop_digest)),
|
||||
seconds=70,
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"GET",
|
||||
_SEARCH_PATH,
|
||||
params={"search": shared_alias, "exclude_team_ids": team_drop, **_range_around_today()},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = object_value(response.json())
|
||||
results: Final = body["results"]
|
||||
assert isinstance(results, list) and len(results) == 1, body
|
||||
entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"])
|
||||
assert set(entities) == {team_keep}, response.text
|
||||
assert set(_team_key_breakdown(body, team_keep)) == {keep_digest}, response.text
|
||||
|
|
@ -5,8 +5,9 @@ from .actors import Actor
|
|||
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
||||
|
||||
|
||||
# GET /team/daily/activity and its /aggregated variant (same shared scope
|
||||
# resolver, so the matrix must hold for both). A proxy admin (admin view) sees
|
||||
# GET /team/daily/activity, its /aggregated variant, and the key-search
|
||||
# variant (same shared scope resolver, so the matrix must hold for all
|
||||
# three). A proxy admin (admin view) sees
|
||||
# activity for any team. A non-admin is scoped to user_info.teams: a bare query
|
||||
# defaults to its own teams (200), and an explicit team_ids filter naming a
|
||||
# team it does not belong to is 404 (the VERIA-43 fix). Org admins have no
|
||||
|
|
@ -43,8 +44,12 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31"
|
|||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint",
|
||||
("/team/daily/activity", "/team/daily/activity/aggregated"),
|
||||
ids=("paginated", "aggregated"),
|
||||
(
|
||||
"/team/daily/activity",
|
||||
"/team/daily/activity/aggregated",
|
||||
"/team/daily/activity/aggregated/search",
|
||||
),
|
||||
ids=("paginated", "aggregated", "search"),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"actor,team,expected_status",
|
||||
|
|
@ -54,7 +59,7 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31"
|
|||
async def test_team_daily_activity_matrix(
|
||||
actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world
|
||||
):
|
||||
query = _DATES
|
||||
query = _DATES + ("&search=x" if endpoint.endswith("/search") else "")
|
||||
if team == "alpha":
|
||||
query += f"&team_ids={world.team_alpha_id}"
|
||||
elif team == "beta":
|
||||
|
|
|
|||
|
|
@ -3600,6 +3600,55 @@ def test_user_daily_activity_aggregated_not_covered_by_prefix_match():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
[
|
||||
"/team/daily/activity",
|
||||
"/team/daily/activity/aggregated",
|
||||
"/team/daily/activity/aggregated/search",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"user_role",
|
||||
[
|
||||
LitellmUserRoles.INTERNAL_USER.value,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
|
||||
],
|
||||
)
|
||||
def test_team_daily_activity_routes_reachable_by_non_admin(route, user_role):
|
||||
"""The Team Usage dashboard calls all three team daily-activity routes, and
|
||||
each handler self-scopes to the caller's teams and own keys
|
||||
(_resolve_team_daily_activity_scope). self_managed_routes is the only list
|
||||
granting them to a non-admin, and check_route_access is exact-match, so each
|
||||
sub-path needs its own entry: dropping one 401s the dashboard before the
|
||||
handler ever runs.
|
||||
"""
|
||||
user_obj = LiteLLM_UserTable(
|
||||
user_id="test_user",
|
||||
user_email="test@example.com",
|
||||
user_role=user_role,
|
||||
)
|
||||
valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role)
|
||||
request = MagicMock(spec=Request)
|
||||
request.query_params = {}
|
||||
|
||||
def outcome() -> str:
|
||||
try:
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=user_role,
|
||||
route=route,
|
||||
request=request,
|
||||
valid_token=valid_token,
|
||||
request_data={},
|
||||
)
|
||||
except Exception as exc:
|
||||
return f"denied: {exc}"
|
||||
return "allowed"
|
||||
|
||||
assert outcome() == "allowed"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"user_role",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -14646,6 +14646,218 @@ async def test_get_team_daily_activity_aggregated_rejects_bad_ranges(
|
|||
mock_aggregated.assert_not_called()
|
||||
|
||||
|
||||
def _key_search_team_setup(mock_db_client, user_id: str, team_id: str):
|
||||
mock_user_info = LiteLLM_UserTable(
|
||||
user_id=user_id,
|
||||
teams=[team_id],
|
||||
max_budget=1000.0,
|
||||
spend=0.0,
|
||||
user_email="test@example.com",
|
||||
user_role="internal_user",
|
||||
)
|
||||
mock_team = MagicMock(spec=LiteLLM_TeamTable)
|
||||
mock_team.team_id = team_id
|
||||
mock_team.team_alias = "Test Team"
|
||||
mock_team.members_with_roles = [Member(user_id=user_id, role="user")]
|
||||
mock_team.model_dump.return_value = {
|
||||
"team_id": team_id,
|
||||
"team_alias": "Test Team",
|
||||
"members_with_roles": [{"user_id": user_id, "role": "user"}],
|
||||
}
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
|
||||
return mock_user_info
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_team_daily_activity_keys_scopes_where_before_take(mock_db_client):
|
||||
"""A member's search must put the team and own-key scoping inside the same
|
||||
Prisma where as the term, because `take` trims rows before Python sees them:
|
||||
scoped outside the where, the top-N slice could be spent entirely on keys
|
||||
the caller is not allowed to see."""
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_LIMIT
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
search_team_daily_activity_keys,
|
||||
)
|
||||
|
||||
user_id = "test_user_123"
|
||||
team_id = "test_team_456"
|
||||
user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
mock_user_info = _key_search_team_setup(mock_db_client, user_id, team_id)
|
||||
|
||||
user_key_1 = MagicMock()
|
||||
user_key_1.token = "user_key_1"
|
||||
matched = MagicMock()
|
||||
matched.token = "user_key_1"
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=[[user_key_1], [matched]])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_user_object:
|
||||
mock_get_user_object.return_value = mock_user_info
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_aggregated:
|
||||
mock_aggregated.return_value = MagicMock()
|
||||
|
||||
await search_team_daily_activity_keys(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
search="Needle",
|
||||
team_ids=team_id,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-31",
|
||||
exclude_team_ids=None,
|
||||
timezone=480,
|
||||
)
|
||||
|
||||
token_calls = mock_db_client.db.litellm_verificationtoken.find_many.call_args_list
|
||||
assert len(token_calls) == 2
|
||||
search_kwargs = token_calls[1][1]
|
||||
assert search_kwargs["where"] == {
|
||||
"team_id": {"in": (team_id,)},
|
||||
"token": {"in": ("user_key_1",)},
|
||||
"OR": (
|
||||
{"token": "Needle"},
|
||||
{"key_alias": {"contains": "Needle", "mode": "insensitive"}},
|
||||
{"user_id": {"contains": "Needle", "mode": "insensitive"}},
|
||||
),
|
||||
}
|
||||
assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT
|
||||
assert search_kwargs["order"] == {"spend": "desc"}
|
||||
|
||||
call_kwargs = mock_aggregated.call_args[1]
|
||||
assert call_kwargs["api_key"] == ["user_key_1"]
|
||||
assert call_kwargs["entity_id"] == [team_id]
|
||||
assert call_kwargs["table_name"] == "litellm_dailyteamspend"
|
||||
assert call_kwargs["include_entity_breakdown"] is True
|
||||
assert call_kwargs["timezone_offset_minutes"] == 480
|
||||
assert call_kwargs["model"] is None
|
||||
assert call_kwargs["entity_metadata_field"] == {team_id: {"team_alias": "Test Team"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_team_daily_activity_keys_admin_unscoped_where(mock_db_client):
|
||||
"""An admin's search has no caller scoping, so the where is the bare OR over
|
||||
token, key alias and user id; every matched hash is passed through to the
|
||||
aggregation."""
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_LIMIT
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
search_team_daily_activity_keys,
|
||||
)
|
||||
|
||||
match_1 = MagicMock()
|
||||
match_1.token = "h1"
|
||||
match_2 = MagicMock()
|
||||
match_2.token = "h2"
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[match_1, match_2])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_aggregated:
|
||||
mock_aggregated.return_value = MagicMock()
|
||||
|
||||
await search_team_daily_activity_keys(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
search="Needle",
|
||||
team_ids=None,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-31",
|
||||
exclude_team_ids=None,
|
||||
timezone=None,
|
||||
)
|
||||
|
||||
search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1]
|
||||
assert search_kwargs["where"] == {
|
||||
"OR": (
|
||||
{"token": "Needle"},
|
||||
{"key_alias": {"contains": "Needle", "mode": "insensitive"}},
|
||||
{"user_id": {"contains": "Needle", "mode": "insensitive"}},
|
||||
)
|
||||
}
|
||||
assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT
|
||||
assert mock_aggregated.call_args[1]["api_key"] == ["h1", "h2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_team_daily_activity_keys_no_match_returns_empty_without_aggregating(
|
||||
mock_db_client,
|
||||
):
|
||||
"""A term matching no key still owes the caller the standard metadata shape
|
||||
(api_key_limit, total_api_keys), and the aggregated query must not run."""
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_LIMIT
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
search_team_daily_activity_keys,
|
||||
)
|
||||
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_aggregated:
|
||||
result = await search_team_daily_activity_keys(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
search="Needle",
|
||||
team_ids=None,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-31",
|
||||
exclude_team_ids=None,
|
||||
timezone=None,
|
||||
)
|
||||
|
||||
assert result.results == []
|
||||
assert result.metadata.total_api_keys == 0
|
||||
assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT
|
||||
mock_aggregated.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_team_daily_activity_keys_excludes_teams_in_where(mock_db_client):
|
||||
"""The dashboard always sends exclude_team_ids=litellm-dashboard; if that
|
||||
filter stayed out of the where, matching keys in excluded teams could fill
|
||||
the take=N slice and push visible matches out."""
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
search_team_daily_activity_keys,
|
||||
)
|
||||
|
||||
matched = MagicMock()
|
||||
matched.token = "h1"
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[matched])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_aggregated:
|
||||
mock_aggregated.return_value = MagicMock()
|
||||
|
||||
await search_team_daily_activity_keys(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
search="Needle",
|
||||
team_ids=None,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-31",
|
||||
exclude_team_ids="litellm-dashboard",
|
||||
timezone=None,
|
||||
)
|
||||
|
||||
search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1]
|
||||
assert search_kwargs["where"] == {
|
||||
"team_id": {"notIn": ("litellm-dashboard",)},
|
||||
"OR": (
|
||||
{"token": "Needle"},
|
||||
{"key_alias": {"contains": "Needle", "mode": "insensitive"}},
|
||||
{"user_id": {"contains": "Needle", "mode": "insensitive"}},
|
||||
),
|
||||
}
|
||||
assert mock_aggregated.call_args[1]["exclude_entity_ids"] == ["litellm-dashboard"]
|
||||
|
||||
|
||||
def _wire_new_team_prisma(mock_db_client):
|
||||
mock_db_client.jsonify_team_object = lambda db_data: db_data
|
||||
mock_db_client.get_data = AsyncMock(return_value=None)
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import type { ColumnDef } from "@tanstack/react-table";
|
|||
import PaginationStatusAlerts from "@/components/shared/PaginationStatusAlerts";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
|
||||
import React, { type ReactNode, useMemo, useState } from "react";
|
||||
import React, { type ReactNode, useCallback, useMemo, useState } from "react";
|
||||
import TeamMultiSelect from "@/components/common_components/team_multi_select";
|
||||
import UserDropdown from "@/components/common_components/UserDropdown";
|
||||
import { ActivityMetrics, processActivityData } from "@/components/activity_metrics";
|
||||
|
|
@ -34,6 +34,7 @@ import {
|
|||
tagDailyActivityCall,
|
||||
teamDailyActivityAggregatedCall,
|
||||
teamDailyActivityCall,
|
||||
teamDailyActivityKeySearchCall,
|
||||
userDailyActivityCall,
|
||||
} from "@/components/networking";
|
||||
import { Logo } from "@/components/molecules/logo/Logo";
|
||||
|
|
@ -182,6 +183,16 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
|
|||
const modelBreakdownKey = modelViewType === "groups" ? "model_groups" : "models";
|
||||
const modelMetrics = processActivityData(spendData, modelBreakdownKey, teams || []);
|
||||
const keyMetrics = processActivityData(spendData, "api_keys", teams || []);
|
||||
const searchTeamKeys = useCallback(
|
||||
(query: string) => {
|
||||
if (!accessToken || !startTime || !endTime) return Promise.resolve({});
|
||||
const teamIds = Array.isArray(entityFilterArg) ? entityFilterArg : null;
|
||||
return teamDailyActivityKeySearchCall(accessToken, startTime, endTime, query, teamIds).then((data) =>
|
||||
processActivityData(data, "api_keys", teams || []),
|
||||
);
|
||||
},
|
||||
[accessToken, startTime, endTime, entityFilterArg, teams],
|
||||
);
|
||||
const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {};
|
||||
|
||||
const getAllTags = () => {
|
||||
|
|
@ -667,6 +678,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
|
|||
keyMetrics={keyMetrics}
|
||||
hidePromptCachingMetrics={entityType === "agent"}
|
||||
apiKeyTruncation={apiKeyTruncation}
|
||||
searchKeys={entityType === "team" ? searchTeamKeys : undefined}
|
||||
/>
|
||||
),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1467,6 +1467,31 @@ export const teamDailyActivityAggregatedCall = async (
|
|||
}
|
||||
};
|
||||
|
||||
export const teamDailyActivityKeySearchCall = async (
|
||||
accessToken: string,
|
||||
startTime: Date,
|
||||
endTime: Date,
|
||||
...options: [search: string, teamIds?: string[] | null]
|
||||
) => {
|
||||
const [search, teamIds = null] = options;
|
||||
try {
|
||||
return await apiClient.get(`/team/daily/activity/aggregated/search`, {
|
||||
accessToken,
|
||||
query: {
|
||||
start_date: formatDate(startTime),
|
||||
end_date: formatDate(endTime),
|
||||
timezone: new Date().getTimezoneOffset().toString(),
|
||||
search,
|
||||
team_ids: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined,
|
||||
exclude_team_ids: "litellm-dashboard",
|
||||
},
|
||||
});
|
||||
} catch (error) {
|
||||
console.error("Failed to search team daily activity keys:", error);
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
export type TeamUserSpendResponse = components["schemas"]["TeamUserSpendResponse"];
|
||||
|
||||
export const teamSpendByUserCall = async (
|
||||
|
|
|
|||
60
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
60
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -15645,6 +15645,27 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/team/daily/activity/aggregated/search": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/**
|
||||
* Search Team Daily Activity Keys
|
||||
* @description Aggregated daily team activity for the keys matching `search`, across every key the caller may
|
||||
* see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend.
|
||||
*/
|
||||
get: operations["search_team_daily_activity_keys_team_daily_activity_aggregated_search_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/team/delete": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -31378,7 +31399,7 @@ export interface components {
|
|||
* @description Enum for key management routes
|
||||
* @enum {string}
|
||||
*/
|
||||
KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/spend/logs" | "/spend/logs/v2";
|
||||
KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2";
|
||||
/**
|
||||
* KeyManagementSystem
|
||||
* @enum {string}
|
||||
|
|
@ -66624,6 +66645,43 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
search_team_daily_activity_keys_team_daily_activity_aggregated_search_get: {
|
||||
parameters: {
|
||||
query: {
|
||||
/** @description Exact token hash, or a case-insensitive substring of the key alias or owning user id */
|
||||
search: string;
|
||||
team_ids?: string | null;
|
||||
start_date?: string | null;
|
||||
end_date?: string | null;
|
||||
exclude_team_ids?: string | null;
|
||||
timezone?: number | null;
|
||||
};
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
delete_team_team_delete_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue