from __future__ import annotations import asyncio import json from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional, Sequence, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth, ) from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @dataclass class _BudgetCounter: counter_key: str max_budget: float fallback_spend: float entity_type: str entity_id: str source_cache_key: Optional[str] = None spend_log_entity_id: Optional[str] = None window_start: Optional[datetime] = None class _CounterReservationUnavailable(Exception): def __init__( self, touched_counter: bool = False, counter_invalidated: bool = False, ) -> None: self.touched_counter = touched_counter self.counter_invalidated = counter_invalidated super().__init__("Counter reservation unavailable") def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: if not budget_reservation: return set() entries = budget_reservation.get("entries") or [] return { entry["counter_key"] for entry in entries if isinstance(entry, dict) and entry.get("counter_key") is not None } async def reserve_budget_for_request( request_body: dict, route: str, llm_router: Optional[Router], valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, ) -> Optional[dict]: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None if route in {"/models", "/v1/models", "/utils/token_counter"}: return None if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None counters = await _get_budget_counters( request_body=request_body, valid_token=valid_token, team_object=team_object, user_object=user_object, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, ) if not counters: return None current_spend_by_counter_key: Dict[str, float] = {} reservation_cost = estimate_request_max_cost( request_body=request_body, route=route, llm_router=llm_router, ) # estimate_request_max_cost still returns None when the model is unknown # to the cost map (no token-priced cost fields, e.g. image/audio routes). # In that case we fall back to read-time enforcement only. if reservation_cost is None or reservation_cost <= 0: return None applied_entries: List[Dict[str, Any]] = [] try: for counter in counters: entry = _counter_to_reservation_entry( counter=counter, reserved_cost=reservation_cost, ) applied_entries.append(entry) try: reserved_value = await _reserve_counter( counter=counter, reservation_cost=reservation_cost, ) except _CounterReservationUnavailable as exc: if exc.touched_counter and not exc.counter_invalidated: await _release_applied_entries_best_effort( entries=[entry], default_reserved_cost=reservation_cost, ) applied_entries.remove(entry) continue if reserved_value is not None: current_spend = reserved_value else: cached_spend = current_spend_by_counter_key.get(counter.counter_key) if cached_spend is None: cached_spend = await _get_current_counter_value(counter=counter) current_spend = cached_spend + reservation_cost if current_spend > counter.max_budget: remaining_before_reservation = counter.max_budget - ( current_spend - reservation_cost ) if remaining_before_reservation > 1e-12: await _resize_applied_reservation( entries=applied_entries, current_reserved_cost=reservation_cost, new_reserved_cost=remaining_before_reservation, ) reservation_cost = remaining_before_reservation continue raise litellm.BudgetExceededError( current_cost=current_spend, max_budget=counter.max_budget, message=( "Budget has been exceeded! " f"{counter.entity_type}={counter.entity_id} " f"Current cost: {current_spend}, " f"Max budget: {counter.max_budget}" ), ) except Exception: await _release_applied_entries_best_effort( entries=applied_entries, default_reserved_cost=reservation_cost, ) raise if not applied_entries: return None input_cost = estimate_request_input_cost( request_body=request_body, route=route, llm_router=llm_router ) return { "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), } async def reconcile_budget_reservation( budget_reservation: Optional[dict], actual_cost: Optional[float], finalize: bool = True, ) -> None: if not budget_reservation or budget_reservation.get("finalized") is True: return reserved_cost = float(budget_reservation.get("reserved_cost") or 0.0) actual = float(actual_cost or 0.0) await _set_reserved_entries_actual_cost( entries=budget_reservation.get("entries") or [], actual_cost=actual, default_reserved_cost=reserved_cost, ) if finalize: budget_reservation["finalized"] = True async def release_budget_reservation(budget_reservation: Optional[dict]) -> None: await reconcile_budget_reservation( budget_reservation=budget_reservation, actual_cost=0.0, ) async def release_budget_reservation_on_cancel( budget_reservation: dict | None, ) -> None: """Reconcile a still-open reservation when the request is cancelled mid-flight. A client disconnect or timeout cancels the request task, which surfaces as CancelledError / GeneratorExit rather than a normal exception, so neither the success cost callback nor the failure hook runs and the pre-call reservation is never reconciled. Left alone it pins the spend counter above real spend and 429s subsequent requests until the counter's TTL expires. Reconcile to the request's input-token cost rather than refunding to zero: by the time a request is cancelled in-flight the provider call was already dispatched, so the input tokens were billed even if no chunk reached the client. Refunding to zero would let a caller abort pre-token to dodge that charge; the worst-case output portion of the reservation is still released. asyncio.shield keeps the reconcile running to completion even though the surrounding task is being cancelled. The `finalized` guard makes this a no-op when success/failure handling already reconciled, so calling it on every cancellation path is safe. """ if not budget_reservation or budget_reservation.get("finalized") is True: return incurred_cost = float(budget_reservation.get("input_cost") or 0.0) try: await asyncio.shield( reconcile_budget_reservation( budget_reservation=budget_reservation, actual_cost=incurred_cost ) ) except (asyncio.CancelledError, Exception): pass async def invalidate_budget_reservation_counters( budget_reservation: Optional[dict], ) -> None: if budget_reservation is None: return from litellm.proxy.proxy_server import _invalidate_spend_counter for counter_key in get_reserved_counter_keys(budget_reservation=budget_reservation): await _invalidate_spend_counter(counter_key=counter_key) async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, proxy_logging_obj: ProxyLogging, end_user_id: Optional[str] = None, end_user_object: Optional[Any] = None, ) -> List[_BudgetCounter]: counters: List[_BudgetCounter] = [] if valid_token.token is not None: if valid_token.max_budget is not None and valid_token.max_budget > 0: counters.append( _BudgetCounter( counter_key=f"spend:key:{valid_token.token}", source_cache_key=valid_token.token, max_budget=float(valid_token.max_budget), fallback_spend=float(valid_token.spend or 0.0), entity_type="Key", entity_id=valid_token.token, ) ) counters.extend( _get_budget_limit_counters( entity_prefix=f"spend:key:{valid_token.token}", entity_type="Key", entity_id=valid_token.token, budget_limits=valid_token.budget_limits, fallback_spend=float(valid_token.spend or 0.0), ) ) if team_object is not None and team_object.team_id is not None: team_id = team_object.team_id if team_object.max_budget is not None and team_object.max_budget > 0: counters.append( _BudgetCounter( counter_key=f"spend:team:{team_id}", source_cache_key=f"team_id:{team_id}", max_budget=float(team_object.max_budget), fallback_spend=float(team_object.spend or 0.0), entity_type="Team", entity_id=team_id, ) ) counters.extend( _get_budget_limit_counters( entity_prefix=f"spend:team:{team_id}", entity_type="Team", entity_id=team_id, budget_limits=team_object.budget_limits, fallback_spend=float(team_object.spend or 0.0), ) ) if ( (team_object is None or team_object.team_id is None) and user_object is not None and user_object.user_id is not None and user_object.max_budget is not None and user_object.max_budget > 0 ): counters.append( _BudgetCounter( counter_key=f"spend:user:{user_object.user_id}", source_cache_key=user_object.user_id, max_budget=float(user_object.max_budget), fallback_spend=float(user_object.spend or 0.0), entity_type="User", entity_id=user_object.user_id, ) ) end_user_counter = await _get_end_user_budget_counter( valid_token=valid_token, end_user_id=end_user_id, end_user_object=end_user_object, ) if end_user_counter is not None: counters.append(end_user_counter) counters.extend( await _get_tag_budget_counters( request_body=request_body, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) ) team_member_counter = await _get_team_member_budget_counter( valid_token=valid_token, team_object=team_object, user_object=user_object, user_api_key_cache=user_api_key_cache, ) if team_member_counter is not None: counters.append(team_member_counter) org_counter = await _get_org_budget_counter( valid_token=valid_token, team_object=team_object, user_api_key_cache=user_api_key_cache, ) if org_counter is not None: counters.append(org_counter) return counters async def _get_end_user_budget_counter( valid_token: UserAPIKeyAuth, end_user_id: Optional[str], end_user_object: Optional[Any], ) -> Optional[_BudgetCounter]: end_user_id = end_user_id or valid_token.end_user_id if end_user_id is None: return None source_cache_key = f"end_user_id:{end_user_id}" max_budget = _to_float(valid_token.end_user_max_budget) fallback_spend = 0.0 if end_user_object is not None: fallback_spend = _to_float(_get_value(end_user_object, "spend")) or 0.0 if max_budget is None: budget_table = _get_value(end_user_object, "litellm_budget_table") max_budget = _to_float(_get_value(budget_table, "max_budget")) if max_budget is None or max_budget <= 0: return None return _BudgetCounter( counter_key=f"spend:end_user:{end_user_id}", source_cache_key=source_cache_key, max_budget=max_budget, fallback_spend=fallback_spend, entity_type="EndUser", entity_id=end_user_id, ) async def _get_tag_budget_counters( request_body: dict, prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, proxy_logging_obj: ProxyLogging, ) -> List[_BudgetCounter]: from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body from litellm.proxy.auth.auth_checks import get_tag_objects_batch tag_names = _dedupe_tags(get_tags_from_request_body(request_body=request_body)) if not tag_names: return [] tag_objects = await get_tag_objects_batch( tag_names=tag_names, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) counters: List[_BudgetCounter] = [] for tag_name in tag_names: tag_object = tag_objects.get(tag_name) if tag_object is None: continue budget_table = _get_value(tag_object, "litellm_budget_table") max_budget = _to_float(_get_value(budget_table, "max_budget")) if max_budget is None or max_budget <= 0: continue counters.append( _BudgetCounter( counter_key=f"spend:tag:{tag_name}", source_cache_key=f"tag:{tag_name}", max_budget=max_budget, fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0, entity_type="Tag", entity_id=tag_name, ) ) return counters def _dedupe_tags(tags: List[str]) -> List[str]: seen = set() deduped_tags = [] for tag in tags: if tag in seen: continue seen.add(tag) deduped_tags.append(tag) return deduped_tags async def _get_team_member_budget_counter( valid_token: UserAPIKeyAuth, team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], user_api_key_cache: DualCache, ) -> Optional[_BudgetCounter]: if ( team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None ): return None membership_cache_key = ( f"team_membership:{valid_token.user_id}:{team_object.team_id}" ) cached_team_membership = await user_api_key_cache.async_get_cache( key=membership_cache_key ) team_membership: Optional[LiteLLM_TeamMembership] = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): team_membership = cached_team_membership elif isinstance(cached_team_membership, dict): team_membership = LiteLLM_TeamMembership(**cached_team_membership) team_member_budget: Optional[float] = None if team_membership is not None and team_membership.litellm_budget_table is not None: team_member_budget = team_membership.litellm_budget_table.max_budget else: default_budget_id = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): default_budget = await user_api_key_cache.async_get_cache( key=f"team_member_default_budget:{default_budget_id}", ) team_member_budget = _to_float(_get_value(default_budget, "max_budget")) if team_member_budget is None or team_member_budget <= 0: return None team_member_spend = ( cast(LiteLLM_TeamMembership, team_membership).spend if team_membership is not None else 0.0 ) return _BudgetCounter( counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", source_cache_key=membership_cache_key, max_budget=float(team_member_budget), fallback_spend=float(team_member_spend or 0.0), entity_type="TeamMember", entity_id=f"{valid_token.user_id}:{team_object.team_id}", ) async def _get_org_budget_counter( valid_token: UserAPIKeyAuth, team_object: Optional[LiteLLM_TeamTable], user_api_key_cache: DualCache, ) -> Optional[_BudgetCounter]: org_id: Optional[str] = None if valid_token.org_id is not None: org_id = valid_token.org_id elif team_object is not None and team_object.organization_id is not None: org_id = team_object.organization_id if org_id is None: return None org_table = await user_api_key_cache.async_get_cache( key=f"org_id:{org_id}:with_budget", ) if org_table is None: return None org_budget_table = _get_value(org_table, "litellm_budget_table") if org_budget_table is None: return None org_max_budget = _to_float(_get_value(org_budget_table, "max_budget")) if org_max_budget is None or org_max_budget <= 0: return None org_spend = _to_float(_get_value(org_table, "spend")) or 0.0 return _BudgetCounter( counter_key=f"spend:org:{org_id}", source_cache_key=f"org_id:{org_id}:with_budget", max_budget=org_max_budget, fallback_spend=org_spend, entity_type="Organization", entity_id=org_id, ) def _get_budget_limit_counters( entity_prefix: str, entity_type: str, entity_id: str, budget_limits: Optional[Sequence[Any]], fallback_spend: float, ) -> List[_BudgetCounter]: counters: List[_BudgetCounter] = [] if not budget_limits: return counters for window in budget_limits: window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") max_budget = window_dict.get("max_budget") if not budget_duration or max_budget is None or max_budget <= 0: continue window_start = get_budget_window_start(window_dict) if window_start is None: verbose_proxy_logger.warning( "Skipping budget window with invalid duration for %s=%s: %s", entity_type, entity_id, budget_duration, ) continue counters.append( _BudgetCounter( counter_key=f"{entity_prefix}:window:{budget_duration}", max_budget=float(max_budget), fallback_spend=0.0, entity_type=entity_type, entity_id=f"{entity_id}:{budget_duration}", spend_log_entity_id=entity_id, window_start=window_start, ) ) return counters def _coerce_window(window: Any) -> dict: if isinstance(window, dict): return window if isinstance(window, str): try: parsed = json.loads(window) return parsed if isinstance(parsed, dict) else {} except Exception: return {} if hasattr(window, "model_dump"): return window.model_dump() return {} async def _reserve_counter( counter: _BudgetCounter, reservation_cost: float, ) -> Optional[float]: from litellm.proxy.proxy_server import ( _ensure_spend_counter_initialized, _ensure_window_spend_counter_initialized, _invalidate_spend_counter, _increment_spend_counter_cache, ) attempted_increment = False try: if counter.source_cache_key is not None: await _ensure_spend_counter_initialized( counter_key=counter.counter_key, source_cache_key=counter.source_cache_key, ) elif ( counter.spend_log_entity_id is not None and counter.window_start is not None ): initialized = await _ensure_window_spend_counter_initialized( counter_key=counter.counter_key, entity_type=counter.entity_type, entity_id=counter.spend_log_entity_id, window_start=counter.window_start, ) if initialized is False: verbose_proxy_logger.warning( "Skipping budget reservation for %s because window spend could not be loaded", counter.counter_key, ) raise _CounterReservationUnavailable attempted_increment = True reserved_value = await _increment_spend_counter_cache( counter_key=counter.counter_key, increment=reservation_cost, ) return float(reserved_value) if reserved_value is not None else None except _CounterReservationUnavailable: raise except Exception: verbose_proxy_logger.warning( "Skipping budget reservation for %s because spend counter reservation failed", counter.counter_key, exc_info=True, ) counter_invalidated = False try: await _invalidate_spend_counter(counter_key=counter.counter_key) counter_invalidated = True except Exception: verbose_proxy_logger.warning( "Failed to invalidate spend counter after budget reservation failure for %s", counter.counter_key, exc_info=True, ) raise _CounterReservationUnavailable( touched_counter=attempted_increment, counter_invalidated=counter_invalidated, ) async def _get_current_counter_value(counter: _BudgetCounter) -> float: from litellm.proxy.proxy_server import get_current_spend return await get_current_spend( counter_key=counter.counter_key, fallback_spend=counter.fallback_spend, ) async def _set_reserved_entries_actual_cost( entries: List[dict], actual_cost: float, default_reserved_cost: float, reseed_on_inconsistent: bool = True, ) -> None: for entry in entries: await _set_reserved_entry_actual_cost( entry=entry, actual_cost=actual_cost, default_reserved_cost=default_reserved_cost, reseed_on_inconsistent=reseed_on_inconsistent, ) async def _set_reserved_entry_actual_cost( entry: dict, actual_cost: float, default_reserved_cost: float, reseed_on_inconsistent: bool = True, ) -> None: from litellm.proxy.proxy_server import ( _increment_spend_counter_cache, reseed_spend_counter_from_db, ) counter_key = entry.get("counter_key") if counter_key is None: return reserved_cost = _get_entry_reserved_cost( entry=entry, default_reserved_cost=default_reserved_cost, ) target_adjustment = actual_cost - reserved_cost applied_adjustment = float(entry.get("applied_adjustment") or 0.0) adjustment = target_adjustment - applied_adjustment if adjustment == 0: return if await _counter_can_apply_adjustment( counter_key=counter_key, adjustment=adjustment, ): await _increment_spend_counter_cache( counter_key=counter_key, increment=adjustment, ) elif reseed_on_inconsistent: # Post-call reconcile / release: the counter was flushed or reseeded # between reservation and reconcile (Redis restart / cross-pod reset), # so the optimistic delta no longer applies. Recover by reseeding from # the DB's lagging authoritative floor rather than deleting the counter # and failing open — deleting it is what left budgets unenforced after a # Redis reload. await reseed_spend_counter_from_db(counter_key=counter_key) else: # Pre-call admission resize: the in-flight reservation cost is not yet # persisted, so the DB floor would discard it. Keep the original # fail-closed behavior (raise -> reserve_budget_for_request releases and # denies) rather than admitting against an inconsistent counter. raise RuntimeError( f"Cannot resize budget reservation against inconsistent counter {counter_key}" ) entry["applied_adjustment"] = target_adjustment async def _counter_can_apply_adjustment( counter_key: str, adjustment: float, ) -> bool: from litellm.proxy.proxy_server import spend_counter_cache current_value = await spend_counter_cache.async_get_cache(key=counter_key) if current_value is None: return False try: current_float = float(current_value) except (TypeError, ValueError): return False return not (adjustment < 0 and current_float + adjustment < -1e-12) async def _release_applied_entries_best_effort( entries: List[dict], default_reserved_cost: float, ) -> None: for entry in entries: try: await _set_reserved_entry_actual_cost( entry=entry, actual_cost=0.0, default_reserved_cost=default_reserved_cost, ) except Exception: counter_key = entry.get("counter_key") verbose_proxy_logger.exception( "Failed to release partial budget reservation during exception cleanup" ) if counter_key is None: continue try: from litellm.proxy.proxy_server import _invalidate_spend_counter await _invalidate_spend_counter(counter_key=counter_key) except Exception: verbose_proxy_logger.exception( "Failed to invalidate partial budget reservation counter during exception cleanup" ) async def _resize_applied_reservation( entries: List[dict], current_reserved_cost: float, new_reserved_cost: float, ) -> None: await _set_reserved_entries_actual_cost( entries=entries, actual_cost=new_reserved_cost, default_reserved_cost=current_reserved_cost, reseed_on_inconsistent=False, ) for entry in entries: entry["reserved_cost"] = new_reserved_cost entry["applied_adjustment"] = 0.0 def _counter_to_reservation_entry( counter: _BudgetCounter, reserved_cost: float, ) -> Dict[str, Any]: return { "counter_key": counter.counter_key, "entity_type": counter.entity_type, "entity_id": counter.entity_id, "reserved_cost": reserved_cost, "applied_adjustment": 0.0, } def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float: try: return float(entry.get("reserved_cost", default_reserved_cost) or 0.0) except (TypeError, ValueError): return default_reserved_cost def get_budget_window_start(window: Any) -> Optional[datetime]: window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") if budget_duration is None: return None try: duration_seconds = duration_in_seconds(str(budget_duration)) except Exception: return None reset_at = _coerce_datetime(window_dict.get("reset_at")) if reset_at is None: return datetime.now(timezone.utc) - timedelta(seconds=duration_seconds) if reset_at.tzinfo is None: reset_at = reset_at.replace(tzinfo=timezone.utc) return reset_at - timedelta(seconds=duration_seconds) def _coerce_datetime(value: Any) -> Optional[datetime]: if value is None: return None if isinstance(value, datetime): return value if isinstance(value, str): try: return datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError: return None return None def estimate_request_max_cost( request_body: dict, route: str, llm_router: Optional[Router], ) -> Optional[float]: model = get_model_from_request(request_body, route, llm_router=llm_router) if model is None: return None models = [model] if isinstance(model, str) else model estimates = [ _estimate_request_max_cost_for_model( request_body=request_body, route=route, model=model_name, llm_router=llm_router, ) for model_name in models ] estimates = [estimate for estimate in estimates if estimate is not None] if not estimates: return None return max(cast(List[float], estimates)) def estimate_request_input_cost( request_body: dict, route: str, llm_router: Router | None, ) -> float | None: """Cost of the request's input tokens alone. Once the provider request is dispatched the input tokens are billed even if the client disconnects before the first chunk, so this is the cost floor a cancelled in-flight request has already incurred. A cancelled reservation is reconciled to this instead of being refunded to zero. """ model = get_model_from_request(request_body, route, llm_router=llm_router) if model is None: return None models = [model] if isinstance(model, str) else model estimates = [ _estimate_request_input_cost_for_model( request_body=request_body, route=route, model=model_name, llm_router=llm_router, ) for model_name in models ] estimates = [estimate for estimate in estimates if estimate is not None] if not estimates: return None return max(cast("list[float]", estimates)) def _estimate_request_input_cost_for_model( request_body: dict, route: str, model: str, llm_router: Router | None, ) -> float | None: model_info = _get_model_cost_info(model=model, llm_router=llm_router) if model_info is None: return None input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) if input_cost_per_token is None: return None input_tokens = _estimate_input_tokens( request_body=request_body, route=route, model=model, model_info=model_info, ) if input_tokens is None: return None return input_tokens * input_cost_per_token def _estimate_request_max_cost_for_model( request_body: dict, route: str, model: str, llm_router: Optional[Router], ) -> Optional[float]: model_info = _get_model_cost_info(model=model, llm_router=llm_router) if model_info is None: return None image_cost = _estimate_image_generation_cost( request_body=request_body, model_info=model_info, ) if image_cost is not None: return image_cost input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) output_cost_per_token = _to_float(model_info.get("output_cost_per_token")) input_tokens = _estimate_input_tokens( request_body=request_body, route=route, model=model, model_info=model_info, ) output_tokens = _estimate_output_tokens( request_body=request_body, route=route, model_info=model_info, ) if input_tokens is None or output_tokens is None: return None cost = 0.0 if input_cost_per_token is not None: cost += input_tokens * input_cost_per_token elif input_tokens > 0: return None output_multiplier = _get_output_multiplier(request_body=request_body) if output_cost_per_token is not None: cost += output_tokens * output_multiplier * output_cost_per_token elif output_tokens > 0: return None return cost def _estimate_image_generation_cost( request_body: dict, model_info: Dict[str, Any], ) -> Optional[float]: """ Reserve `n × per-image cost` for image-generation requests so concurrent requests against a depleted budget cannot all slip past the admission gate onto the provider. Token-based pricing (e.g. gpt-image-1) is handled by the chat-route token path; per-pixel and size/quality-tiered pricing (DALL-E 2 size variants, premium tiers) are not handled here and fall through to read-time enforcement. The "output" vs "input" cost-per-image naming is inconsistent across providers — OpenAI's dall-e-3 entry uses ``input_cost_per_image`` while aiml/dall-e-3 uses ``output_cost_per_image`` — so both are summed. """ # Gate strictly on `mode`. Several chat and embedding models carry # ``input_cost_per_image`` / ``output_cost_per_image`` to price multimodal # *vision input* (e.g. ``gemini-3.1-pro-preview``, ``azure/gpt-realtime-*``, # ``amazon.titan-embed-image-v1``). Falling back to "treat as image-gen if # an image cost field is present" would short-circuit the token-priced # path for those models and reserve a fraction of a cent instead of the # true per-token cost. All real image-generation entries in # ``model_prices_and_context_window.json`` carry ``mode: image_generation`` # or ``mode: image_edit``, so the field-presence fallback is unnecessary. if model_info.get("mode") not in ("image_generation", "image_edit"): return None output_cost_per_image = _to_float(model_info.get("output_cost_per_image")) input_cost_per_image = _to_float(model_info.get("input_cost_per_image")) cost_per_image = (output_cost_per_image or 0.0) + (input_cost_per_image or 0.0) if cost_per_image <= 0: return None n = _to_int(request_body.get("n")) or 1 return cost_per_image * max(n, 1) def _get_model_cost_info( model: str, llm_router: Optional[Router], ) -> Optional[Dict[str, Any]]: if llm_router is not None: try: model_group_info = llm_router.get_model_group_info(model_group=model) if model_group_info is not None: return model_group_info.model_dump() except Exception: verbose_proxy_logger.debug( "Unable to load router model group info for budget reservation", exc_info=True, ) try: return dict(litellm.get_model_info(model=model)) except Exception: return None def _estimate_input_tokens( request_body: dict, route: str, model: str, model_info: Dict[str, Any], ) -> Optional[int]: try: if "messages" in request_body: return litellm.token_counter( model=model, messages=request_body.get("messages") or [], tools=request_body.get("tools"), tool_choice=request_body.get("tool_choice"), ) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: return _count_text_tokens(model=model, text=request_body.get("input")) if "query" in request_body or "documents" in request_body: query_tokens = _count_text_tokens( model=model, text=request_body.get("query") ) document_tokens = _count_text_tokens( model=model, text=request_body.get("documents"), ) return query_tokens + document_tokens except Exception: verbose_proxy_logger.debug( "Unable to count input tokens for budget reservation", exc_info=True ) max_input_tokens = _to_int(model_info.get("max_input_tokens")) if max_input_tokens is not None: return max_input_tokens return None DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK = 16384 def _estimate_output_tokens( request_body: dict, route: str, model_info: Dict[str, Any], ) -> Optional[int]: if _is_input_only_route(route=route): return 0 requested: Optional[int] = None for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): requested = _to_int(request_body.get(key)) if requested is not None: break # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the # atomic admission counter actually bounds concurrent in-flight cost # (mirrors parallel_request_limiter_v3's DEFAULT_MAX_TOKENS_ESTIMATE). # (2) An adversarial caller cannot send max_tokens=999999999 to inflate # the reservation up to remaining team headroom and pin the counter # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. model_ceiling = ( _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK ) if requested is None: requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK return min(requested, model_ceiling) def _count_text_tokens(model: str, text: Any) -> int: if text is None: return 0 token_count = 0 stack = [text] while stack: item = stack.pop() if item is None: continue if isinstance(item, list): stack.extend(item) continue if isinstance(item, dict): token_count += litellm.token_counter(model=model, text=json.dumps(item)) continue token_count += litellm.token_counter(model=model, text=str(item)) return token_count def _get_output_multiplier(request_body: dict) -> int: output_multiplier = 1 for key in ("n", "best_of"): value = _to_int(request_body.get(key)) if value is not None: output_multiplier = max(output_multiplier, value) return output_multiplier def _is_input_only_route(route: str) -> bool: return any( route_part in route for route_part in ( "embeddings", "rerank", "moderations", ) ) def _to_float(value: Any) -> Optional[float]: if value is None: return None try: return float(value) except (TypeError, ValueError): return None def _to_int(value: Any) -> Optional[int]: if value is None: return None try: return int(value) except (TypeError, ValueError): return None def _get_value(obj: Any, key: str) -> Any: if isinstance(obj, dict): return obj.get(key) return getattr(obj, key, None)