From c2eb549ee685b442b74bc7c1cf9c4c1a7698a5ec Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:04:12 -0500 Subject: [PATCH] feat(usage): search team keys beyond the top-N in the Team usage view (#42857) --- litellm/proxy/_types.py | 5 + .../management_endpoints/team_endpoints.py | 109 +++++++++ .../management_endpoints/team_endpoints.py | 24 ++ .../endpointaudit/coverage_allowlist.txt | 1 + .../test_team_daily_activity_key_search.py | 128 +++++++++++ .../management/test_team_daily_activity.py | 15 +- .../proxy/auth/test_route_checks.py | 49 ++++ .../test_team_endpoints.py | 212 ++++++++++++++++++ .../components/EntityUsage/EntityUsage.tsx | 14 +- .../src/components/networking.tsx | 25 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 60 ++++- 11 files changed, 635 insertions(+), 7 deletions(-) create mode 100644 tests/integration/spend/test_team_daily_activity_key_search.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 51a3eb03067..73e3e0e6ee0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6cec3e714ec..493c83c730a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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 "" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 4524c47ec38..aac2703e918 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -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 diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index d0f9c31ecaf..b0f28a9c740 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -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 diff --git a/tests/integration/spend/test_team_daily_activity_key_search.py b/tests/integration/spend/test_team_daily_activity_key_search.py new file mode 100644 index 00000000000..2b0395cf4ba --- /dev/null +++ b/tests/integration/spend/test_team_daily_activity_key_search.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py index d84cc4c94af..9bbc8fdde29 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -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": diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 571e066e947..7bb79a115dd 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -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", [ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b066b3b80e6..38241926f8e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27460b21108..16dc41c3ba8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -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 = ({ 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 = ({ keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} apiKeyTruncation={apiKeyTruncation} + searchKeys={entityType === "team" ? searchTeamKeys : undefined} /> ), }, diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index e6923a26d4c..edaf56a8a16 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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 ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79d008c1b48..a398b2db93c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;