diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 2b328b455dc..57a43b43ee6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4284,7 +4284,7 @@ async def can_key_call_resolved_model( llm_model_list: list | None, valid_token: UserAPIKeyAuth, llm_router: litellm.Router | None, -) -> None: +) -> tuple[str, ...]: from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, @@ -4355,8 +4355,9 @@ async def can_key_call_resolved_model( proxy_logging_obj=proxy_logging_obj, ) + project_object: LiteLLM_ProjectTableCachedObj | None = None if valid_token.project_id is not None: - project_object: Final = await get_project_object( + project_object = await get_project_object( project_id=valid_token.project_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -4369,6 +4370,24 @@ async def can_key_call_resolved_model( llm_router=llm_router, ) + matched_model_access_groups: Final = await collect_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if matched_model_access_groups: + await _model_access_group_max_budget_check( + matched_model_access_groups=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return matched_model_access_groups + def can_org_access_model( model: str, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6a9c3a095e6..a511b715b75 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -741,10 +741,18 @@ async def _update_database_and_spend_counters( ) raise - if defer_budget_counter_update: - return - try: + if defer_budget_counter_update: + if model_access_groups: + from litellm.proxy.proxy_server import increment_fusion_model_access_group_spend_counters + + await increment_fusion_model_access_group_spend_counters( + model_access_groups=model_access_groups, + response_cost=response_cost, + budget_reservation=budget_reservation, + ) + return + await increment_spend_counters( token=user_api_key, team_id=team_id, @@ -756,6 +764,10 @@ async def _update_database_and_spend_counters( tags=request_tags, request_started_at=start_time, model_access_groups=model_access_groups, + # Global scopes reconcile the whole logical Fusion request. Access + # groups are deployment-specific, so the final provider call must + # be charged only its own cost; hidden calls were charged above. + model_access_group_response_cost=response_cost, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b16037ca55e..5e7ff306b6c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2670,6 +2670,7 @@ async def increment_spend_counters( tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + model_access_group_response_cost: float | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2692,6 +2693,9 @@ async def increment_spend_counters( return cost: Final[float] = response_cost + model_access_group_cost: Final = ( + cost if model_access_group_response_cost is None else model_access_group_response_cost + ) async def _key_scope(key_token: str) -> None: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow @@ -2825,10 +2829,10 @@ async def increment_spend_counters( else None, _increment_model_access_group_spend_counters( model_access_groups=model_access_groups, - response_cost=cost, + response_cost=model_access_group_cost, reserved_counter_keys=reserved_counter_keys, ) - if model_access_groups + if model_access_groups and model_access_group_cost != 0 else None, _increment_org_spend_counter( org_id=org_id, @@ -2853,6 +2857,27 @@ async def increment_spend_counters( budget_reservation["finalized"] = True +async def increment_fusion_model_access_group_spend_counters( + model_access_groups: Sequence[str], + response_cost: float, + budget_reservation: dict | None, # mutable-ok: shared reservation ledger is read here to avoid duplicate charges +) -> None: + """Charge one deferred Fusion provider call to only its serving access groups. + + Fusion defers the shared key/team/user reservation until its final outer call, + but model access groups are deployment-specific. Updating those counters per + provider call preserves attribution while skipping any group already reserved + for the virtual Fusion model itself. + """ + from litellm.proxy.spend_tracking.budget_reservation import get_reserved_counter_keys + + await _increment_model_access_group_spend_counters( + model_access_groups=model_access_groups, + response_cost=response_cost, + reserved_counter_keys=get_reserved_counter_keys(budget_reservation=budget_reservation), + ) + + async def _reconcile_budget_reservation_for_counter_update( budget_reservation: dict | None, response_cost: float | None, diff --git a/litellm/router.py b/litellm/router.py index 5a3fae75a1d..3ae9712b3f5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,16 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Callable, + Generator, + Iterator, + Mapping, + MutableMapping, + Sequence, +) from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -68,6 +77,7 @@ from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.ptu_pricing import ( PTU_COST_ATTRIBUTION_ENV_VAR, @@ -9470,9 +9480,11 @@ class Router: async def _authorize_fusion_dependencies( self, fusion_router: FusionRouter, - request_kwargs: Mapping[str, object], + request_kwargs: MutableMapping[ # mutable-ok: request-local carrier is enriched before hidden calls dispatch + str, object + ], ) -> None: - """Apply the originating proxy caller's model access to every hidden call.""" + """Apply the originating proxy caller's model access and group budgets to every hidden call.""" metadata_values: Final = tuple(request_kwargs.get(key) for key in ("litellm_metadata", "metadata")) raw_user_api_key_auth: Final = next( ( @@ -9512,13 +9524,51 @@ class Router: ) ) ) - for dependency_model in dependency_models: - await can_key_call_resolved_model( - model=dependency_model, - llm_model_list=self.model_list, - valid_token=user_api_key_auth, - llm_router=self, + matched_dependency_groups: Final = await asyncio.gather( + *( + can_key_call_resolved_model( + model=dependency_model, + llm_model_list=self.model_list, + valid_token=user_api_key_auth, + llm_router=self, + ) + for dependency_model in dependency_models ) + ) + dependency_access_groups: Final = frozenset( + group for matched_groups in matched_dependency_groups for group in matched_groups + ) + if not dependency_access_groups: + return + + # Keep the group gating the virtual Fusion model and add every group + # gating a hidden dependency. The normal spend writer narrows this + # authorization upper bound to groups serving each provider call. + all_groups: Final = tuple( + dict.fromkeys( + ( + *(user_api_key_auth.matched_model_access_groups or ()), + *sorted(dependency_access_groups), + ) + ) + ) + user_api_key_auth.matched_model_access_groups = list( # mutable-ok: auth schema requires a list carrier + all_groups + ) + for metadata_key in ("litellm_metadata", "metadata"): + metadata = request_kwargs.get(metadata_key) + if not isinstance(metadata, Mapping): + continue + mutable_metadata = ( + metadata + if isinstance(metadata, dict) + else dict(metadata) # mutable-ok: SDK metadata boundary requires a native mapping + ) + mutable_metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list( # mutable-ok: metadata JSON requires a list + all_groups + ) + mutable_metadata["user_api_key_auth"] = user_api_key_auth + request_kwargs[metadata_key] = mutable_metadata # rebind-ok: enrich request-local metadata for dispatch async def _fusion_asearch( # kwargs-ok: bridge preserves the Router.asearch keyword surface self, diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index fb82d8708fd..d1e65d781ce 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -25,6 +25,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ( _model_access_group_max_budget_check, + can_key_call_resolved_model, collect_matched_model_access_groups, common_checks, stamp_matched_model_access_groups, @@ -381,6 +382,34 @@ async def test_group_exactly_at_its_max_budget_blocks_the_request(): assert exc_info.value.current_cost == 10.0 +@pytest.mark.asyncio +async def test_resolved_model_authorization_enforces_its_access_group_budget( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.proxy import proxy_server + + cache = await _cache() + prisma = _RecordingPrismaClient(_MagBudgetRow("tier-a", spend=10.0, max_budget=10.0)) + router = Router(model_list=MODEL_LIST) + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"]) + read, seen = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}) + + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "get_current_spend", read) + + with pytest.raises(litellm.BudgetExceededError): + await can_key_call_resolved_model( + model="gpt-4o", + llm_model_list=MODEL_LIST, + valid_token=valid_token, + llm_router=router, + ) + + assert seen == [MODEL_ACCESS_GROUP_COUNTER_KEY] + + @pytest.mark.asyncio async def test_group_just_under_its_max_budget_passes(): """Asserting the counter was read is what keeps this honest: a group that got skipped entirely, diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index c40fa9322fe..971edf12282 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -607,6 +607,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda tags=["tag-a"], request_started_at=start_time, model_access_groups=("premium",), + model_access_group_response_cost=0.2, ) @@ -730,6 +731,7 @@ async def test_fusion_hidden_costs_accumulate_then_continuation_reconciles_once( increment.assert_awaited_once() assert increment.await_args.kwargs["response_cost"] == pytest.approx(0.7) + assert increment.await_args.kwargs["model_access_group_response_cost"] == pytest.approx(0.4) assert increment.await_args.kwargs["budget_reservation"] is reservation diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index d63391ffb10..b553726f922 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1192,7 +1192,7 @@ def test_fusion_reservation_does_not_return_a_partial_additive_estimate() -> Non def child_estimate(*, model: str, **_: object) -> float | None: return None if model == "panel" else 1.0 - with patch( + with patch( # test-quality-ok: isolates one unavailable dependency estimate to verify all-or-nothing Fusion admission "litellm.proxy.spend_tracking.budget_reservation._estimate_request_max_cost_for_model", side_effect=child_estimate, ): @@ -3445,6 +3445,53 @@ async def test_model_access_group_counter_accumulates_across_calls_without_a_res assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None +@pytest.mark.asyncio +async def test_model_access_group_can_use_its_provider_call_cost_during_fusion_reconciliation(spend_counter_state): + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "panel", spend=1.0, max_budget=25.0) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.7, + model_access_groups=["panel"], + model_access_group_response_cost=0.2, + ) + + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("panel") + ) == pytest.approx(1.2) + + +@pytest.mark.asyncio +async def test_deferred_fusion_access_group_cost_skips_a_group_reserved_by_the_virtual_model(spend_counter_state): + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "fusion", spend=1.0, max_budget=25.0) + await _cache_model_access_group_budget(key_cache, "panel", spend=2.0, max_budget=25.0) + await counter_cache.async_set_cache(key=model_access_group_spend_counter_key("fusion"), value=1.5) + reservation = { + "entries": [{"counter_key": model_access_group_spend_counter_key("fusion")}], + } + + from litellm.proxy.proxy_server import increment_fusion_model_access_group_spend_counters + + await increment_fusion_model_access_group_spend_counters( + model_access_groups=["fusion", "panel"], + response_cost=0.2, + budget_reservation=reservation, + ) + + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("fusion") + ) == pytest.approx(1.5) + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("panel") + ) == pytest.approx(2.2) + + @pytest.mark.asyncio async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state): """The reservation already wrote this counter, so the post-call pass has to skip it.""" diff --git a/tests/test_litellm/test_fusion_router.py b/tests/test_litellm/test_fusion_router.py index 129bfc6ff8d..0bee42d5b86 100644 --- a/tests/test_litellm/test_fusion_router.py +++ b/tests/test_litellm/test_fusion_router.py @@ -17,6 +17,7 @@ from litellm.fusion_router import ( fusion_router_dependencies, validate_fusion_router_write, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.router import Router from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponseStream @@ -683,18 +684,32 @@ async def test_proxy_fusion_authorizes_every_hidden_model(monkeypatch: pytest.Mo ) model_list[-1]["litellm_params"]["fusion_router_config"]["analyst_model"] = "analyst" router = Router(model_list=model_list) - authorize = AsyncMock(return_value=None) + authorize = AsyncMock(side_effect=[("outer-budget",), ("panel-budget",), ("analyst-budget",)]) monkeypatch.setattr(auth_checks, "can_key_call_resolved_model", authorize) + auth = UserAPIKeyAuth(models=["*"], matched_model_access_groups=["fusion-budget"]) + metadata: dict[str, object] = {"user_api_key_auth": auth} await router._authorize_fusion_dependencies( # pyright: ignore[reportPrivateUsage] fusion_router=router.fusion_routers["fusion/test"], request_kwargs={ - "metadata": {"user_api_key_auth": UserAPIKeyAuth(models=["*"])}, + "metadata": metadata, "proxy_server_request": {"body": {"model": "fusion/test"}}, }, ) assert [call.kwargs["model"] for call in authorize.await_args_list] == ["outer", "panel-a", "analyst"] + assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == [ + "fusion-budget", + "analyst-budget", + "outer-budget", + "panel-budget", + ] + assert auth.matched_model_access_groups == [ + "fusion-budget", + "analyst-budget", + "outer-budget", + "panel-budget", + ] @pytest.mark.asyncio