diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 15ca658ee38..e5068b020d1 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -34,7 +34,7 @@ from litellm.proxy.management_endpoints.daily_activity_scopes import ( EntityScopeResolver, ResolvedScope, ) -from litellm.proxy.management_endpoints.ptu_consumption import with_ptu_consumption +from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, with_ptu_consumption from litellm.proxy.management_endpoints.team_endpoints import aggregated_date_range_error from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import PrismaClient, get_prisma_client_or_throw @@ -272,7 +272,9 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver, include_entity_breakdown=resolver.include_entity_breakdown, api_key_limit=api_key_limit, ) - return with_ptu_consumption(activity, llm_router) if resolver.entity == "team" else activity + if resolver.entity != "team": + return activity + return with_ptu_consumption(activity, llm_router, single_team_id(resolved.scope.entity_ids)) except HTTPException: raise except Exception as exc: diff --git a/litellm/proxy/management_endpoints/ptu_consumption.py b/litellm/proxy/management_endpoints/ptu_consumption.py index 7d13c293d70..7b3752d184a 100644 --- a/litellm/proxy/management_endpoints/ptu_consumption.py +++ b/litellm/proxy/management_endpoints/ptu_consumption.py @@ -5,7 +5,7 @@ cached, and completion tokens a team sent to a PTU model group convert back to t a PTU-hour it consumed, reported next to the raw token counts. """ -from collections.abc import Callable +from collections.abc import Callable, Sequence from types import MappingProxyType from typing import Final @@ -13,7 +13,11 @@ from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabl from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours from litellm.router import Router from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.ptu_shares import model_group_ptu_capacity, routed_deployments +from litellm.router_utils.ptu_shares import ( + model_group_ptu_capacity, + routed_deployments, + team_servable_deployments, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, MetricWithMetadata, @@ -99,31 +103,38 @@ def attach_ptu_hours( ) -def capacity_by_requested_name(llm_router: Router) -> Callable[[str], PTUCapacity | None]: +def capacity_by_requested_name(llm_router: Router, team_id: str | None) -> Callable[[str], PTUCapacity | None]: """The sizing row behind the name a usage row is keyed by, resolved the way the ceiling resolves a request: an alias to its group, then a group, a routing group, a deployment id, - or a provider model to the deployments it is served from.""" + or a provider model to the deployments it is served from. For one team those are narrowed + to the deployments it can be served from, its shared one first, so a team with no share on + a name reads no PTU-hours for it; a page spanning teams keeps every deployment behind the name.""" listed_rows: Final = llm_router.get_model_list() or () aliases: Final = llm_router.model_group_alias def capacity_for(requested_model: str) -> PTUCapacity | None: model_group: Final = resolve_model_group_alias(aliases, requested_model) or requested_model - return model_group_ptu_capacity( - routed_deployments( - listed_rows, - llm_router.model_list, # pyright: ignore[reportUnknownArgumentType] # Router.model_list is a bare list - model_group, - ) + routed: Final = routed_deployments( + listed_rows, + llm_router.model_list, # pyright: ignore[reportUnknownArgumentType] # Router.model_list is a bare list + model_group, ) + return model_group_ptu_capacity(routed if team_id is None else team_servable_deployments(routed, team_id)) return capacity_for +def single_team_id(team_ids: Sequence[str] | None) -> str | None: + """The one team a usage page is scoped to, else None for a page spanning several or all teams.""" + return team_ids[0] if team_ids is not None and len(team_ids) == 1 else None + + def with_ptu_consumption( - activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None + activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None, team_id: str | None ) -> SpendAnalyticsPaginatedResponse: - """``activity`` with PTU-hours attached from the router's sized deployments, untouched while - PTU cost attribution is off or no router is loaded.""" + """``activity`` with PTU-hours attached from the deployments ``team_id`` is served from (every + sized deployment when the page spans teams), untouched while PTU cost attribution is off or no + router is loaded.""" if llm_router is None or not is_ptu_cost_attribution_enabled(): return activity - return attach_ptu_hours(activity, capacity_by_requested_name(llm_router)) + return attach_ptu_hours(activity, capacity_by_requested_name(llm_router, team_id)) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3d7609401df..e8486925d54 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -143,7 +143,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, ) -from litellm.proxy.management_endpoints.ptu_consumption import with_ptu_consumption +from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, with_ptu_consumption from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, @@ -6689,7 +6689,7 @@ async def get_team_daily_activity( page=page, page_size=page_size, ) - return with_ptu_consumption(activity, llm_router) + return with_ptu_consumption(activity, llm_router, single_team_id(scope.team_ids)) _MAX_AGGREGATED_RANGE_DAYS: Final = 400 diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index ed59a3a40d1..b763b470ff2 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -150,15 +150,19 @@ def _model_group_of( ) -> str: """The group of the deployment behind ``requested_model`` this team can be served from, one holding its share first.""" - routed: Final = routed_deployments(listed_rows, deployments, requested_model) - servable: Final = filter_ptu_shared_deployments(routed, team_id).deployments - shared_first: Final = sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None) + servable: Final = team_servable_deployments(routed_deployments(listed_rows, deployments, requested_model), team_id) return next( - (group for deployment in shared_first if (group := _deployment_model_group(deployment)) is not None), + (group for deployment in servable if (group := _deployment_model_group(deployment)) is not None), requested_model, ) +def team_servable_deployments(deployments: Sequence[_DeploymentT], team_id: str) -> tuple[_DeploymentT, ...]: + """The deployments ``team_id`` can be served from, the ones it holds a share on first.""" + servable: Final = filter_ptu_shared_deployments(deployments, team_id).deployments + return tuple(sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None)) + + def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]]) -> PTUCapacity | None: """The sizing row of the group's first reserved deployment, single-team or shared, so a team's tokens on that group convert to PTU-hours.""" diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 69caaba6005..5cb37b16221 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -3,7 +3,7 @@ import io from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from dataclasses import dataclass, fields from itertools import chain -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock @@ -578,7 +578,7 @@ def test_aggregated_routes_return_scoped_results( assert len(body["results"]) == 2, response.text -def _ptu_sized_router(model_group: str) -> Router: +def _ptu_sized_router(model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20})) -> Router: return Router( model_list=[ { @@ -590,7 +590,7 @@ def _ptu_sized_router(model_group: str) -> Router: "ptu_count": 50, "cost_per_ptu_per_hour": 1.0, "ptu_effective_from": "2024-01-01T00:00:00Z", - "ptu_shares": {"team-a": 30, "team-b": 20}, + "ptu_shares": dict(shares), }, } ] @@ -637,6 +637,26 @@ def test_team_aggregated_route_reports_ptu_hours_for_the_sized_model_group_only_ assert body["metadata"]["total_tokens"] == 90, response.text +def test_team_aggregated_route_reports_no_ptu_hours_to_a_team_holding_no_share( + daily_activity_client: tuple[TestClient, _FakeRepository], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """rare-group is split between team-b and team-c, so team-a's page alone converts none of its + rare-group tokens even though the group is sized.""" + client, _ = daily_activity_client + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + _inject_llm_router(client, _ptu_sized_router("rare-group", shares=MappingProxyType({"team-b": 30, "team-c": 20}))) + + response: Final = client.get("/team/daily/activity/aggregated", params=_entity_params("team_ids", "team-a")) + + assert response.status_code == 200, response.text + body: Final = response.json() + rare_day: Final = {day["date"]: day for day in body["results"]}["2025-01-01"] + assert body["metadata"]["total_ptu_hours"] == 0, response.text + assert rare_day["breakdown"]["model_groups"]["rare-group"]["metrics"]["ptu_hours"] == 0, response.text + assert body["metadata"]["total_tokens"] == 90, response.text + + def test_ptu_hours_stay_zero_outside_the_team_routes( daily_activity_client: tuple[TestClient, _FakeRepository], monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py index 0b4b048aa1b..d05135d94f3 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py @@ -159,9 +159,79 @@ def test_rows_keyed_by_an_alias_a_deployment_id_or_a_provider_model_are_sized_li names: Final = ("gpt-4.1-ptu", "ptu-alias", "shared-ptu", "azure/gpt-4.1") one_hour_each: Final = {name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in names} - attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _shared_ptu_router()) + attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _shared_ptu_router(), "team-a") groups: Final = attached.results[0].breakdown.model_groups assert [groups[name].metrics.ptu_hours for name in names] == [pytest.approx(1.0)] * len(names) assert attached.results[0].metrics.ptu_hours == pytest.approx(float(len(names))) assert attached.metadata.total_ptu_hours == pytest.approx(float(len(names))) + + +def _mixed_ptu_router() -> Router: + """One group split between team-a on a gpt-4.1 PTU deployment and team-b on a gpt-5.5 one, + beside an open pay-as-you-go deployment of gpt-4.1 in its own group.""" + return Router( + model_list=[ + { + "model_name": "ptu", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "ptu-41", + "base_model": "azure/gpt-4.1", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 50}, + }, + }, + { + "model_name": "ptu", + "litellm_params": {"model": "azure/gpt-5.5", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "ptu-55", + "base_model": "azure/gpt-5.5", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-b": 50}, + }, + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-payg", "api_base": "https://payg.example"}, + "model_info": {"id": "payg"}, + }, + ] + ) + + +_ONE_GPT_41_PTU_HOUR_ON_GPT_55: Final = _ONE_PTU_HOUR_OF_INPUT / AZURE_PTU_CAPACITY["gpt-5.5"].normalized_tokens_per_ptu_hour + + +@pytest.mark.parametrize( + ("team_id", "expected_on_the_group", "expected_on_the_provider_model"), + [ + ("team-a", 1.0, 1.0), + ("team-b", _ONE_GPT_41_PTU_HOUR_ON_GPT_55, 0.0), + ("team-c", 0.0, 0.0), + (None, 1.0, 1.0), + ], +) +def test_a_teams_rows_are_sized_by_the_deployment_it_is_served_from( + monkeypatch, team_id: str | None, expected_on_the_group: float, expected_on_the_provider_model: float +): + """A team's page converts its tokens through the deployment the ceiling served it from: team-b's + share is on the gpt-5.5 deployment so its group row counts at that rate, its provider-model row + reached only the open deployment so it counts nothing, team-c holds no share so it counts nothing, + and a page spanning teams keeps the group's first sized deployment.""" + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + one_hour_each: Final = { + name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in ("ptu", "azure/gpt-4.1") + } + + attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _mixed_ptu_router(), team_id) + + groups: Final = attached.results[0].breakdown.model_groups + assert groups["ptu"].metrics.ptu_hours == pytest.approx(expected_on_the_group) + assert groups["azure/gpt-4.1"].metrics.ptu_hours == pytest.approx(expected_on_the_provider_model) + assert attached.metadata.total_ptu_hours == pytest.approx(expected_on_the_group + expected_on_the_provider_model) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 52ed68f2aa9..6c9ed92808a 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -16964,6 +16964,39 @@ async def test_team_daily_activity_sizes_a_team_scoped_deployment_by_its_public_ assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 1.0 +@pytest.mark.asyncio +async def test_team_daily_activity_reports_no_ptu_hours_to_a_team_holding_no_share( + mock_db_client, mock_admin_auth, monkeypatch +): + """The deployment is split between team-a and team-b, so team-c's page alone converts none of + its tokens on the group even though the group is sized.""" + from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity + + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + page = _ptu_activity_page() + + with ( + patch("litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", AsyncMock(return_value=page)), + patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router()), + ): + result = await get_team_daily_activity( + team_ids="team-c", + start_date="2026-09-23", + end_date="2026-09-24", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=mock_admin_auth, + ) + + assert result.metadata.total_ptu_hours == 0.0 + assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 0.0 + assert result.results[0].metrics.total_tokens == _ONE_PTU_HOUR_OF_INPUT + + @pytest.mark.parametrize( ("start_date", "end_date"), (