From 35d3a1784b90cb2dd6e20afcb74ebf1f4bedbac4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 14:27:23 -0700 Subject: [PATCH 01/31] fix(proxy): preserve access windows when editing models (#49316) --- .../model_management_endpoints.py | 6 +-- .../test_model_management_endpoints.py | 51 ++++++++++++++++++- 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index dbd77e45d57..2c7453ebb1c 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1515,14 +1515,10 @@ def update_db_model( # convert to prisma compatible format - for key, value in merged_model_info.items(): - if isinstance(value, datetime.datetime): - merged_model_info[key] = value.isoformat() - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( model_name=merged_model_name, litellm_params=json.dumps(merged_litellm_params), - model_info=json.dumps(merged_model_info), + model_info=TypeAdapter(dict[str, object]).dump_json(merged_model_info).decode(), ) if updated_patch.blocked is not None: diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index abaa802de2c..ab6cfcfebfe 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4,8 +4,9 @@ import inspect import json import socket from collections.abc import AsyncIterator, Iterator, Mapping +from datetime import datetime, time, timezone from types import MappingProxyType, SimpleNamespace -from typing import Dict, Final, Optional +from typing import Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -46,7 +47,14 @@ from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIM from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.router_api import ApplyResult, ChangeSet -from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelAccessWindow, + ModelInfo, + updateDeployment, + updateLiteLLMParams, +) async def _passthrough_row(update_data): @@ -4048,6 +4056,45 @@ class TestUpdateDBModelBlocked: assert "blocked" not in result +@pytest.mark.parametrize("window_patch", ["omit", "echo", "replace", "clear"]) +def test_model_edit_preserves_or_explicitly_changes_access_windows( + window_patch: Literal["omit", "echo", "replace", "clear"], +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + saved_window: Final = ModelAccessWindow( + start=time(13), end=time(17), timezone="Asia/Almaty", team_ids=("team-window", "team-shared") + ) + replacement_window: Final = ModelAccessWindow( + start=time(22, 0, 0, 123456), end=time(6, 0, 0, 654321), + timezone="America/New_York", team_ids=("team-other",), + ) + created_at: Final = datetime(2026, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc) + requested_windows: Final = { + "omit": None, + "echo": (saved_window,), + "replace": (replacement_window,), + "clear": (), + }[window_patch] + expected_windows: Final = (saved_window,) if requested_windows is None else requested_windows + result: Final = update_db_model( + db_model=Deployment( + model_name="window-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="window-deployment", access_windows=(saved_window,), created_at=created_at), + ), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(max_parallel_requests=7), + model_info=None if requested_windows is None else ModelInfo(access_windows=requested_windows), + ), + ) + + assert json.loads(result["litellm_params"])["max_parallel_requests"] == 7 + saved_info: Final = ModelInfo.model_validate_json(result["model_info"]) + assert saved_info.access_windows == expected_windows + assert saved_info.created_at == created_at + + class TestUpdateDBModelKeepsLegacyDropParams: def test_partial_patch_keeps_encrypted_string_drop_params(self, monkeypatch): from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper From ec4f66bb94ff706bb7f2fbb55261228baed29f77 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 14:38:14 -0700 Subject: [PATCH 02/31] fix(router): budget classifier input by provider limits (#49306) --- litellm/decisions/call.py | 8 +- litellm/integrations/otel/routing.py | 5 + .../llms/azure_ai/decisions/transformation.py | 12 +- .../llms/base_llm/decisions/transformation.py | 21 +- .../llms/typesafe/decisions/transformation.py | 13 +- litellm/router_api.py | 3 + .../complexity_router/complexity_router.py | 119 ++++++---- .../complexity_router/jev_classifier.py | 157 ++++++++++++- litellm/types/utils.py | 10 + .../integrations/otel/test_otel_v2_logger.py | 5 + tests/unit/integrations/test_opentelemetry.py | 5 + .../test_base_decisions_transformation.py | 60 +++++ .../complexity_router/test_jev_classifier.py | 216 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 + 14 files changed, 589 insertions(+), 55 deletions(-) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 3b73f65c917..ae01aa8d0f7 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -218,7 +218,7 @@ def supported_decisions_providers() -> tuple[str, ...]: ) -def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig: +def get_decisions_provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig: provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None) provider_config: Final = ( None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider) @@ -408,7 +408,7 @@ def _provider_ir_request( return supported -def _outbound_headers( +def decisions_outbound_headers( extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object] ) -> Mapping[str, str]: forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( @@ -466,7 +466,7 @@ def prepare_call( api_base=call_api_base, api_key=api_key, ) - provider_config: Final = _provider_config(upstream_model, provider) + provider_config: Final = get_decisions_provider_config(upstream_model, provider) if not upstream_model: raise litellm.BadRequestError( message="A model name is required for the Decisions API", @@ -478,7 +478,7 @@ def prepare_call( except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: - headers: Final = _outbound_headers(extra_headers, provider, kwargs) + headers: Final = decisions_outbound_headers(extra_headers, provider, kwargs) request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs) request: Final = _fields_without_unsupported_media( fields, provider_config, canonical_model, provider, kwargs diff --git a/litellm/integrations/otel/routing.py b/litellm/integrations/otel/routing.py index 82dd6237edc..339775855be 100644 --- a/litellm/integrations/otel/routing.py +++ b/litellm/integrations/otel/routing.py @@ -21,6 +21,11 @@ _ROUTING_SCALAR_FIELDS: Final = frozenset( "classifier_cost", "classifier_failure_reason", "classifier_error_type", + "classifier_input_budget_tokens", + "classifier_input_tokens", + "classifier_input_count_method", + "classifier_input_truncated", + "classifier_input_budget_status", "classifier_confidence", "classifier_primary_rule", "classifier_capability_boundary", diff --git a/litellm/llms/azure_ai/decisions/transformation.py b/litellm/llms/azure_ai/decisions/transformation.py index fb5330ed715..22c4b59f5e8 100644 --- a/litellm/llms/azure_ai/decisions/transformation.py +++ b/litellm/llms/azure_ai/decisions/transformation.py @@ -12,7 +12,7 @@ from litellm.llms.azure_ai.common_utils import ( AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_auth_headers, ) -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits _FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$") _AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"}) @@ -26,6 +26,16 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig): api_base_env = ("AZURE_AI_API_BASE",) api_key_required = False + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + limits: Final = super().get_input_limits(model, custom_llm_provider, headers) + return ( + limits + if limits.max_request_tokens is not None + else super().get_input_limits("Microsoft-Decision-1", custom_llm_provider, headers) + ) + def missing_api_base_message(self, custom_llm_provider: str) -> str: return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://.services.ai.azure.com" diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index e8e4b77e757..efd3e3fe9b5 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -1,7 +1,7 @@ import json from abc import ABC from collections.abc import Callable, Mapping, Sequence -from dataclasses import replace +from dataclasses import dataclass, replace from itertools import chain from types import MappingProxyType from typing import Final @@ -511,6 +511,12 @@ def _is_decisions_model(model: str, custom_llm_provider: str) -> bool: return is_decisions_model_mode(model_mode) +@dataclass(frozen=True, slots=True) +class DecisionsInputLimits: + max_request_tokens: int | None = None + max_question_tokens: int | None = None + + class BaseDecisionsConfig(ABC): path: str = "/v1/systemone" api_key_env: tuple[str, ...] = () @@ -533,6 +539,19 @@ class BaseDecisionsConfig(ABC): def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: return None + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + try: + limit: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get( + "max_input_tokens" + ) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model + return DecisionsInputLimits() + return DecisionsInputLimits( + max_request_tokens=limit if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0 else None + ) + def missing_api_base_message(self, custom_llm_provider: str) -> str: return f"api_base is required for Decisions provider '{custom_llm_provider}'" diff --git a/litellm/llms/typesafe/decisions/transformation.py b/litellm/llms/typesafe/decisions/transformation.py index 952c874bc8e..871596b7f00 100644 --- a/litellm/llms/typesafe/decisions/transformation.py +++ b/litellm/llms/typesafe/decisions/transformation.py @@ -6,7 +6,7 @@ from types import MappingProxyType from typing import Final, Literal from litellm.constants import TYPESAFE_DANGER_ZONE_HEADER_PREFIX -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, image_media_type +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits, image_media_type from litellm.types.decisions import ( DecisionsAsset, DecisionsIRMessages, @@ -19,6 +19,7 @@ from litellm.types.decisions import ( ) _IMAGES_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}images" +_LONG_CONTEXT_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}long-context" _ASSET_FORMATS: Final[Mapping[str, Literal["png", "jpg"]]] = MappingProxyType( {"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg"} ) @@ -84,6 +85,16 @@ class TypeSafeDecisionsConfig(BaseDecisionsConfig): def get_default_api_base(self) -> str | None: return "https://api.typesafe.ai" + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + long_context_values: Final = tuple( + value.strip().lower() for name, value in headers.items() if name.lower() == _LONG_CONTEXT_HEADER + ) + if long_context_values and all(value == "true" for value in long_context_values): + return DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144) + return DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768) + def supports_images_in( self, decision_input: DecisionsIRState | DecisionsIRMessages, headers: Mapping[str, str] ) -> bool: diff --git a/litellm/router_api.py b/litellm/router_api.py index d62999cf261..19f12cec40d 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1216,6 +1216,9 @@ class RoutingContext: def registry(self) -> RouterRegistry: return self._router.registry() + def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]: + return self._router._params_sent_to_deployment(request, deployment) + def settings(self) -> RouterSettings: return self._router.settings() diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 30f1021251f..97ad5ffa05b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -126,10 +126,13 @@ from .jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevClassifierClient, + JevSystemOneRequest, JevVerdict, RouterJevClassifierClient, build_jev_request, + fit_jev_request, jev_classifier_cost, + jev_input_targets, ) from .llm_v2 import LLM_V2_PROMPT_VERSION, LLMV2Decision, LLMV2TaskContext, LLMV2Verdict, llm_v2_response_format from .stall_detector import detect_stalled_task @@ -1166,6 +1169,7 @@ class ClassificationOutcome(NamedTuple): heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None classifier_failure_reason: ClassifierFailureReason | None = None classifier_error_type: str | None = None + classifier_input_telemetry: StandardLoggingRoutingDecision | None = None def _with_classifier_failure( @@ -1207,23 +1211,24 @@ def _with_classifier_forecast( decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome ) -> StandardLoggingRoutingDecision: """Attach validated forecasts and their applied policy to the routing decision.""" + with_input: Final[StandardLoggingRoutingDecision] = {**decision, **(outcome.classifier_input_telemetry or {})} if outcome.heuristic_v2_forecast is not None: - return {**decision, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} + return {**with_input, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} if outcome.jev_verdict is not None: forecasted_decision: Final[StandardLoggingRoutingDecision] = { - **decision, + **with_input, "classifier_probabilities": outcome.jev_verdict.probabilities, "classifier_confidence": outcome.jev_verdict.confidence, } return forecasted_decision if outcome.llm_v2_forecast is not None: - return _with_llm_v2_forecast(decision, outcome.llm_v2_forecast) + return _with_llm_v2_forecast(with_input, outcome.llm_v2_forecast) forecast: Final = outcome.capability_forecast if forecast is None: - return decision + return with_input verdict: Final = forecast.verdict enriched: Final[StandardLoggingRoutingDecision] = { - **decision, + **with_input, "classifier_crux": verdict.crux, "classifier_primary_rule": verdict.primary_rule, "classifier_capability_boundary": verdict.capability_boundary, @@ -1974,19 +1979,21 @@ class ComplexityRouter(CustomLogger): `score` because the score is only interpretable against the boundaries that mapped it to a tier. """ - decision: Final = StandardLoggingRoutingDecision( - router_model_name=self.model_name, - router_type="complexity", - routed_model=routed_model, - cause=cause, - conversation_continuing=conversation_continuing, - **( - {"classifier_failure_reason": classifier_failure_reason} - if classifier_failure_reason is not None - else {} - ), - **({"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {}), + failure: Final[StandardLoggingRoutingDecision] = ( + {"classifier_failure_reason": classifier_failure_reason} if classifier_failure_reason is not None else {} ) + error: Final[StandardLoggingRoutingDecision] = ( + {"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {} + ) + decision: Final[StandardLoggingRoutingDecision] = { + "router_model_name": self.model_name, + "router_type": "complexity", + "routed_model": routed_model, + "cause": cause, + "conversation_continuing": conversation_continuing, + **failure, + **error, + } if (baseline := self.savings_baseline) is not None: decision["savings_baseline_model"] = baseline.model if baseline.deployment_id is not None: @@ -2371,16 +2378,6 @@ class ComplexityRouter(CustomLogger): system_prompt, failure_reason="unsupported_input", ) - breaker: Final = self._classifier_circuit_breaker - permit: Final = breaker.acquire_permit() if breaker is not None else None - if breaker is not None and permit is None: - return self._classifier_failure_outcome( - "jev classifier circuit is open", - prompt, - system_prompt, - failure_reason="circuit_open", - signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL, - ) criteria: Final[Mapping[str, str]] = ( MappingProxyType( { @@ -2395,13 +2392,40 @@ class ComplexityRouter(CustomLogger): ) ) timeout_s: Final = config.timeout_ms / 1000 - request: Final = build_jev_request( - prompt=self._classifier_context_payload(prompt, system_prompt, request_kwargs, messages), - system_prompt=None, - model=config.model, - instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, - criteria=criteria, + from litellm.litellm_core_utils.token_counter import offload_token_count + + def build_request(current_prompt: str, caller_constraints: str | None) -> JevSystemOneRequest: + return build_jev_request( + prompt=self._classifier_context_payload( + current_prompt, caller_constraints, request_kwargs, messages, current_ask_for_history=prompt + ), + system_prompt=None, + model=config.model, + instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, + criteria=criteria, + ) + + budgeted: Final = await offload_token_count(fit_jev_request)( + build_request, prompt, system_prompt, jev_input_targets(self.litellm_router_instance, config), _truncate ) + request: Final = budgeted.request + if request is None: + return self._classifier_failure_outcome( + "classifier input cannot fit the provider budget", + prompt, + system_prompt, + failure_reason="unsupported_input", + )._replace(classifier_input_telemetry=budgeted.telemetry) + breaker: Final = self._classifier_circuit_breaker + permit: Final = breaker.acquire_permit() if breaker is not None else None + if breaker is not None and permit is None: + return self._classifier_failure_outcome( + "jev classifier circuit is open", + prompt, + system_prompt, + failure_reason="circuit_open", + signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL, + )._replace(classifier_input_telemetry=budgeted.telemetry) try: response: Final = await asyncio.wait_for(client.evaluate(request, timeout_s, request_kwargs), timeout_s) answer: Final = response.answers.get("tier") @@ -2449,6 +2473,7 @@ class ComplexityRouter(CustomLogger): cause="jev_classifier", classifier_cost=verdict.cost, jev_verdict=verdict, + classifier_input_telemetry=budgeted.telemetry, ) except asyncio.CancelledError: if breaker is not None and permit is not None: @@ -2463,7 +2488,7 @@ class ComplexityRouter(CustomLogger): system_prompt, failure_reason=_classifier_failure_reason(e), error_type=type(e).__name__, - ) + )._replace(classifier_input_telemetry=budgeted.telemetry) def _classifier_failure_outcome( self, @@ -2654,6 +2679,7 @@ class ComplexityRouter(CustomLogger): messages: Sequence[Mapping[str, object]] | None, *, encrypted_task: bool = False, + current_ask_for_history: str | None = None, ) -> str: include_assistant: Final = self.config.classifier_context_include_assistant_turns marker_pairs: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) @@ -2661,7 +2687,7 @@ class ComplexityRouter(CustomLogger): prior_turns: Final = ( _extract_prior_turns( messages, - current_ask=prompt, + current_ask=prompt if current_ask_for_history is None else current_ask_for_history, window_size=self.config.classifier_context_window_size, budget_chars=self.config.classifier_context_budget_chars, per_turn_chars=self.config.classifier_context_per_turn_chars, @@ -4983,16 +5009,19 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, - routing_decision=self._build_routing_decision( - routed_model=fallback_model, - conversation_continuing=conversation_continuing, - cause=outcome.cause, - signals=outcome.signals, - classifier_failure_reason=outcome.classifier_failure_reason, - classifier_error_type=outcome.classifier_error_type, - escalation_keyword=escalation_keyword, - escalated=False, - tier_litellm_params=fallback_tier_params, + routing_decision=_with_classifier_forecast( + self._build_routing_decision( + routed_model=fallback_model, + conversation_continuing=conversation_continuing, + cause=outcome.cause, + signals=outcome.signals, + classifier_failure_reason=outcome.classifier_failure_reason, + classifier_error_type=outcome.classifier_error_type, + escalation_keyword=escalation_keyword, + escalated=False, + tier_litellm_params=fallback_tier_params, + ), + outcome, ), litellm_params=fallback_tier_params, ) diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 87bd031490a..25ea813229c 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias @@ -25,10 +26,13 @@ from litellm.llms.pass_through.typesafe_logging_handler import TypeSafePassthrou from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, StandardLoggingRoutingDecision if TYPE_CHECKING: from litellm.router import Router + from litellm.router_api import RoutingContext + from litellm.router_strategy.complexity_router.config import OpenSourceClassifierConfig + from litellm.types.router import Deployment JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] ClassifierProvider: TypeAlias = Literal["typesafe", "laya", "bespoke", "databricks"] @@ -52,6 +56,155 @@ class JevSystemOneRequest(LiteLLMBaseModel): questions: Mapping[str, JevChoiceQuestion] +@dataclass(frozen=True, slots=True) +class JevInputTarget: + model: str + max_tokens: int | None + + +@dataclass(frozen=True, slots=True) +class JevBudgetedRequest: + request: JevSystemOneRequest | None + telemetry: StandardLoggingRoutingDecision + + +def jev_input_targets(router: "Router", config: "OpenSourceClassifierConfig") -> tuple[JevInputTarget, ...]: + from litellm.router_api import RoutingContext + from litellm.types.router import Deployment + + if config.deployment_name is None: + provider: Final = "typesafe" if config.provider == "jev" else config.provider + return (_jev_input_target({"model": f"{provider}/{config.model}"}, None),) + context: Final = RoutingContext(router) + deployments: Final = context.registry().deployments(model_name=config.deployment_name) or () + return tuple( + _jev_deployment_input_target(context, Deployment.model_validate(deployment)) for deployment in deployments + ) + + +def _jev_deployment_input_target(context: "RoutingContext", deployment: "Deployment") -> JevInputTarget: + params: Final = context.deployment_params({}, deployment) + explicit_limit: Final = deployment.model_info.model_dump().get("max_input_tokens") + return _jev_input_target(params, explicit_limit) + + +def _jev_input_target(params: Mapping[str, object], explicit_limit: object) -> JevInputTarget: + from litellm.decisions.call import decisions_outbound_headers, get_decisions_provider_config + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + + qualified_model: Final = str(params.get("model") or "") + override: Final = params.get("custom_llm_provider") + try: + model, provider = resolve_model_and_provider_for_metadata( + qualified_model, override if isinstance(override, str) else None + ) + headers: Final = decisions_outbound_headers( + TypeAdapter[Mapping[str, str] | None](Mapping[str, str] | None).validate_python( + params.get("extra_headers") + ), + provider, + params, + ) + limits: Final = get_decisions_provider_config(model, provider).get_input_limits(model, provider, headers) + candidates: Final = tuple( + value + for value in (limits.max_request_tokens, limits.max_question_tokens, explicit_limit) + if isinstance(value, int) and not isinstance(value, bool) and value > 0 + ) + return JevInputTarget(f"{provider}/{model}", min(candidates) if candidates else None) + except Exception: # noqa: BLE001 # unavailable metadata must not turn an existing classifier into a request failure + return JevInputTarget(qualified_model, None) + + +def _jev_request_count(request: JevSystemOneRequest, models: Sequence[str], budget: int) -> tuple[int, str]: + from litellm._lazy_imports import get_token_counter_new + + payload: Final = request.model_dump_json(exclude={"model"}) + byte_count: Final = len(payload.encode("utf-8")) + if byte_count <= budget or litellm.disable_token_counter: + return byte_count, "utf8_bytes" + try: + counts: Final = tuple(get_token_counter_new()(model=model, text=payload) for model in models) + if counts and all(count > 0 for count in counts): + return max(counts), "estimated_tokens" + except Exception: # noqa: BLE001 # tokenization is best-effort; byte counting remains available + return byte_count, "utf8_bytes" + return byte_count, "utf8_bytes" + + +def fit_jev_request( + build_request: Callable[[str, str | None], JevSystemOneRequest], + prompt: str, + system_prompt: str | None, + targets: Sequence[JevInputTarget], + truncate: Callable[[str, int], str], +) -> JevBudgetedRequest: + known: Final = tuple(target for target in targets if target.max_tokens is not None) + if not known: + return JevBudgetedRequest( + build_request(prompt, system_prompt), + {"classifier_input_budget_status": "unknown_limits", "classifier_input_truncated": False}, + ) + budget: Final = max(1, int(min(target.max_tokens for target in known if target.max_tokens is not None) * 0.75)) + models: Final = tuple(dict.fromkeys(target.model for target in known)) + initial_cap: Final = max(120, budget * 8) + bounded_prompt: Final = truncate(prompt, initial_cap) + bounded_system: Final = truncate(system_prompt, initial_cap) if system_prompt is not None else None + return _fit_jev_components( + build_request, + bounded_prompt, + bounded_system, + models, + budget, + truncate, + len(known) != len(targets), + bounded_prompt != prompt or bounded_system != system_prompt, + ) + + +def _fit_jev_components( + build_request: Callable[[str, str | None], JevSystemOneRequest], + prompt: str, + system_prompt: str | None, + models: Sequence[str], + budget: int, + truncate: Callable[[str, int], str], + unknown_limits: bool, + truncated: bool, +) -> JevBudgetedRequest: + attempts: Final = len(prompt).bit_length() + len(system_prompt or "").bit_length() + 1 + components: tuple[str, str | None, bool, int] = ( # rebind-ok: bounded iterative fitting state + prompt, + system_prompt, + truncated, + attempts, + ) + while True: + current_prompt, current_system, was_truncated, remaining = components + request = build_request(current_prompt, current_system) + count, method = _jev_request_count(request, models, budget) + telemetry: StandardLoggingRoutingDecision = { + "classifier_input_budget_tokens": budget, + "classifier_input_tokens": count, + "classifier_input_count_method": method, + "classifier_input_truncated": was_truncated, + "classifier_input_budget_status": "unknown_limits" + if unknown_limits + else "truncated" + if was_truncated + else "within_budget", + } + if count <= budget: + return JevBudgetedRequest(request, telemetry) + if remaining == 0 or (len(current_system or "") <= 123 and len(current_prompt) <= 123): + return JevBudgetedRequest(None, {**telemetry, "classifier_input_budget_status": "untrimmable_input"}) + components = ( + (current_prompt, truncate(current_system, max(120, len(current_system) // 2)), True, remaining - 1) + if current_system is not None and len(current_system) > 123 + else (truncate(current_prompt, max(120, len(current_prompt) // 2)), current_system, True, remaining - 1) + ) + + class JevChoiceAnswer(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, allow_inf_nan=False) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index db865c7488a..1a8f5adf154 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3323,6 +3323,11 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_cost: float classifier_failure_reason: ReadOnly[ClassifierFailureReason] classifier_error_type: ReadOnly[str] + classifier_input_budget_tokens: ReadOnly[int] + classifier_input_tokens: ReadOnly[int] + classifier_input_count_method: ReadOnly[str] + classifier_input_truncated: ReadOnly[bool] + classifier_input_budget_status: ReadOnly[str] classifier_probabilities: ReadOnly[Mapping[str, float]] classifier_confidence: ReadOnly[float] heuristic_v2_forecast: ReadOnly[StandardLoggingHeuristicV2Forecast] @@ -3381,6 +3386,11 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_cost", "classifier_failure_reason", "classifier_error_type", + "classifier_input_budget_tokens", + "classifier_input_tokens", + "classifier_input_count_method", + "classifier_input_truncated", + "classifier_input_budget_status", "classifier_probabilities", "classifier_confidence", "heuristic_v2_forecast", diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 1f9d3af71d9..ac1f5557bf6 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -207,6 +207,11 @@ def test_routing_diagnostics_reach_llm_spans_without_promoting_prompt_text(fail: "classifier_cost": 0.001, "classifier_failure_reason": "timeout", "classifier_error_type": "TimeoutError", + "classifier_input_budget_tokens": 24576, + "classifier_input_tokens": 23000, + "classifier_input_count_method": "estimated_tokens", + "classifier_input_truncated": True, + "classifier_input_budget_status": "truncated", "escalated": False, } payload: Final = _payload( diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index f0fca94c370..5f749e951b1 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -344,6 +344,11 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada "cause": "heuristic", "classifier_failure_reason": "timeout", "classifier_error_type": "TimeoutError", + "classifier_input_budget_tokens": 24576, + "classifier_input_tokens": 23000, + "classifier_input_count_method": "estimated_tokens", + "classifier_input_truncated": True, + "classifier_input_budget_status": "truncated", "classifier_cost": 0.001, } with provider.get_tracer(__name__).start_as_current_span("inference") as span: diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 03459ba16c8..ec6f118837f 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -6,14 +6,19 @@ import httpx import pytest from pydantic import TypeAdapter, ValidationError +import litellm +from litellm.llms.azure_ai.decisions.transformation import AzureAIDecisionsConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( + BaseDecisionsConfig, + DecisionsInputLimits, ir_to_systemone_request, ir_to_systemone_response, parse_systemone_response, systemone_request_to_ir, ) from litellm.llms.openai.decisions.transformation import openai_request_to_ir +from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig from litellm.types.decisions import ( MAX_DECISION_QUESTIONS, DecisionsIRRefusal, @@ -49,6 +54,61 @@ _SYSTEMONE_REQUEST: Final[Mapping[str, object]] = { } +@pytest.mark.parametrize("declared", (4096, 0, -1, True, None)) +def test_decision_input_limits_use_only_positive_catalog_limits( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, declared: int | None +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bespoke/test-input-limit", + {"litellm_provider": "bespoke", "mode": "evaluation", "max_input_tokens": declared}, + ) + + limits: Final = BaseDecisionsConfig().get_input_limits("test-input-limit", "bespoke", {}) + + assert limits == DecisionsInputLimits(max_request_tokens=4096 if declared == 4096 else None) + + +def test_unmapped_decision_model_limits_remain_unknown(local_model_cost_map: None) -> None: + assert BaseDecisionsConfig().get_input_limits("unmapped-input-limit", "databricks", {}) == DecisionsInputLimits() + + +@pytest.mark.parametrize( + ("headers", "long_context"), + ( + ({}, False), + ({"TypeSafe-Danger-Zone-Long-Context": "false"}, False), + ({"typesafe-danger-zone-long-context": " TRUE "}, True), + ({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "true"}, True), + ({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "false"}, False), + ({"typesafe-danger-zone-long-context": "true,false"}, False), + ), +) +def test_typesafe_limits_require_unambiguous_long_context_opt_in( + headers: Mapping[str, str], long_context: bool +) -> None: + """Provider contract from https://docs.typesafe.ai/models, verified 2026-10-30""" + limits: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", headers) + + assert limits == ( + DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144) + if long_context + else DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768) + ) + + +@pytest.mark.parametrize("model", ("Microsoft-Decision-1", "unmapped-decision-deployment")) +def test_azure_decision_deployments_use_catalog_limits_without_typesafe_modes( + local_model_cost_map: None, model: str +) -> None: + catalog_limit: Final = litellm.get_model_info("Microsoft-Decision-1", "azure_ai")["max_input_tokens"] + limits: Final = AzureAIDecisionsConfig().get_input_limits( + model, "azure_ai", {"TypeSafe-Danger-Zone-Long-Context": "true"} + ) + + assert limits == DecisionsInputLimits(max_request_tokens=catalog_limit) + + def _openai_ir(raw: Mapping[str, object]) -> DecisionsIRRequest: return openai_request_to_ir(_OPENAI_BODY.validate_python(raw)) diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 77d465526a4..a71f572254e 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -9,6 +9,7 @@ from unittest.mock import create_autospec import httpx import pytest import respx +import tiktoken import litellm from litellm._logging import verbose_router_logger @@ -16,12 +17,14 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter, _ClassifierCircuitBreaker from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, JevClassifierConfig from litellm.router_strategy.complexity_router.jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevChoiceAnswer, + JevSystemOneRequest, JevSystemOneResponse, JevUsage, build_jev_request, @@ -1055,3 +1058,214 @@ def test_switching_to_deployed_classifier_does_not_copy_legacy_credentials() -> ) assert write.effective is not None assert write.effective["opensource_classifier_config"] == {"deployment_name": "registered"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("large_field", ["user", "system"]) +@pytest.mark.parametrize("disable_counter", [False, True]) +async def test_jev_bounds_complete_input_without_changing_forwarded_messages( + large_field: str, disable_counter: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_token_counter", disable_counter) + limit: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", {}).max_question_tokens + assert limit is not None + rubric: Final = "Operator routing rubric " + "document the task category " * 500 + criterion: Final = ("Simple tasks " + "the classification criteria " * 15).rstrip() + large_text: Final = "BEGIN_SOURCE " + "budget 预算 function(x) return x; " * 18000 + " END_SOURCE" + ask: Final = large_text if large_field == "user" else "CURRENT_ASK Explain the example" + constraints: Final = large_text if large_field == "system" else "Use concise language" + messages: Final = [ + {"role": "system", "content": constraints}, + {"role": "user", "content": "PRIOR_TURN Explain the budget"}, + {"role": "user", "content": ask}, + ] + original: Final = deepcopy(messages) + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient() + classifier: Final = ComplexityRouter( + "bounded-classifier", litellm.Router(model_list=[]), + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"instructions": rubric}, + "tier_definitions": [{"name": "SIMPLE", "description": criterion}, {"name": "REASONING"}], + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + "fallback_tier": "SIMPLE", + "session_affinity": False, + "deployment_affinity": False, + }, + jev_client=HttpJevClassifierClient("test", "https://budget.test", handler), + derive_savings_baseline=False, + ) + try: + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("https://budget.test/v1/systemone").respond( + json={"answers": {"tier": _answer().model_dump()}} + ) + result: Final = await classifier.async_pre_routing_hook( + model="bounded-classifier", request_kwargs={"messages": messages}, messages=messages + ) + sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content) + complete_input: Final = sent.model_dump_json(exclude={"model"}) + measured: Final = len(tiktoken.get_encoding("o200k_base").encode(complete_input)) + assert measured < limit + assert sent.questions["tier"].instructions == rubric + assert sent.questions["tier"].criteria["SIMPLE"] == criterion + assert sent.state.count("BEGIN_SOURCE") == sent.state.count("END_SOURCE") == 1 + assert sent.state.count("PRIOR_TURN") == 1 + assert result is not None and result.routing_decision is not None + decision: Final = result.routing_decision + assert decision["cause"] == "jev_classifier" + assert decision["classifier_input_budget_status"] == "truncated" + assert decision["classifier_input_truncated"] is True + assert measured <= decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] < limit + assert messages == original + assert result.messages is None or result.messages == original + if disable_counter: + assert len(complete_input.encode()) < limit + assert decision["classifier_input_count_method"] == "utf8_bytes" + finally: + await GLOBAL_LOGGING_WORKER.flush() + await handler.client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mode", "default_header", "explicit_cap", "other_member", "expected_status"), + [ + ("false", None, None, None, "truncated"), + ("true", None, None, None, "within_budget"), + ("true", "false", None, None, "truncated"), + ("true", None, 8192, None, "truncated"), + ("true", None, 8192, "default", "truncated"), + ("true", None, None, "default", "truncated"), + ("false", None, None, "unknown", "unknown_limits"), + ], +) +async def test_jev_budget_uses_effective_headers_and_every_deployment( + mode: str, default_header: str | None, explicit_cap: int | None, + other_member: str | None, expected_status: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"])) + header: Final = "TypeSafe-Danger-Zone-Long-Context" + deployment: Final = { + "model_name": "budget-model", + "litellm_params": { + "model": "typesafe/jev-latest", "api_key": "test", "api_base": "https://budget.test", + "extra_headers": {header: mode}, + }, + "model_info": {"max_input_tokens": explicit_cap} if explicit_cap is not None else {}, + } + sibling: Final = { + "model_name": "budget-model", + "litellm_params": { + "model": "databricks/unknown-budget-endpoint" if other_member == "unknown" else "typesafe/jev-latest", + "api_key": "test", "api_base": "https://budget.test", + }, + } + router: Final = litellm.Router( + model_list=[deployment, sibling] if other_member else [deployment], + default_litellm_params={"extra_headers": {header: default_header}} if default_header else {}, + num_retries=0, + ) + classifier: Final = ComplexityRouter( + "bounded-classifier", router, + {"classifier_type": "oss_classifier", "opensource_classifier_config": {"deployment_name": "budget-model"}, + "tiers": {"SIMPLE": "cheap"}}, + derive_savings_baseline=False, + ) + ask: Final = "Explain this code " + "word001 word002 word003 " * 12000 + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.route(method="POST").respond( + json={"answers": {"tier": _answer().model_dump()}} + ) + outcome: Final = await classifier.aclassify(ask) + assert outcome.cause == "jev_classifier" + telemetry: Final = outcome.classifier_input_telemetry + assert telemetry is not None and telemetry["classifier_input_budget_status"] == expected_status + assert telemetry["classifier_input_tokens"] <= telemetry["classifier_input_budget_tokens"] + sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content) + if other_member is None: + assert route.calls[0].request.headers[header] == (default_header or mode) + if explicit_cap is not None: + assert telemetry["classifier_input_budget_tokens"] < explicit_cap + assert (ask in sent.state) is (expected_status == "within_budget") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("oversized_field", "fallback"), + [("instructions", "heuristic"), ("instructions", "default_model"), ("criteria", "custom")], +) +async def test_jev_untrimmable_rubric_uses_configured_fallback_without_a_provider_call( + oversized_field: str, fallback: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"])) + router: Final = litellm.Router(model_list=[{ + "model_name": "budget-model", + "litellm_params": {"model": "typesafe/jev-latest", "api_key": "test", "api_base": "https://budget.test"}, + "model_info": {"max_input_tokens": 128 if oversized_field == "criteria" else 1024}, + }]) + rubric: Final = "word001 word002 word003 " * 2000 + classifier: Final = ComplexityRouter( + "bounded-classifier", router, + {"classifier_type": "oss_classifier", "opensource_classifier_config": { + "deployment_name": "budget-model", "instructions": rubric if oversized_field == "instructions" else "Choose", + }, "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + **({"tier_definitions": [{"name": name, "description": "word001 " * 50} + for name in ("SIMPLE", "REASONING")], "fallback_tier": "SIMPLE"} + if fallback == "custom" else {"classifier_fallback": fallback, "default_model": "default"})}, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.route(method="POST").respond(json={"answers": {"tier": _answer().model_dump()}}) + outcome: Final = await classifier.aclassify("What is two plus two?") + assert route.call_count == 0 + assert outcome.cause == { + "heuristic": "heuristic_scorer", "default_model": "default_model_fallback", "custom": "classifier_fallback", + }[fallback] + assert outcome.classifier_failure_reason == "unsupported_input" + assert outcome.classifier_input_telemetry is not None + assert outcome.classifier_input_telemetry["classifier_input_budget_status"] == "untrimmable_input" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback", ["heuristic", "default_model"]) +async def test_jev_open_breaker_preserves_input_budget_in_fallback_decision(fallback: str) -> None: + def clock() -> float: + return 100.0 + + transport: Final = create_autospec(httpx.AsyncBaseTransport, instance=True) + transport.handle_async_request.side_effect = TimeoutError("controlled classifier timeout") + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=transport) + classifier: Final = ComplexityRouter( + "breaker-budget", litellm.Router(model_list=[]), + {"classifier_type": "oss_classifier", "opensource_classifier_config": {}, + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, "classifier_fallback": fallback, + "default_model": "default", "session_affinity": False, "deployment_affinity": False}, + jev_client=HttpJevClassifierClient("test", "https://budget.test", handler), + derive_savings_baseline=False, + ) + classifier._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=clock) + try: + first: Final = await classifier.async_pre_routing_hook( + model="breaker-budget", request_kwargs={}, messages=[{"role": "user", "content": "Explain budgets"}] + ) + second: Final = await classifier.async_pre_routing_hook( + model="breaker-budget", request_kwargs={}, + messages=[{"role": "user", "content": "Explain this code " + "word001 word002 word003 " * 12000}], + ) + assert first is not None and first.routing_decision is not None + assert first.routing_decision["classifier_failure_reason"] == "timeout" + assert second is not None and second.routing_decision is not None + decision: Final = second.routing_decision + assert decision["classifier_failure_reason"] == "circuit_open" + assert decision["cause"] == ("heuristic_scorer" if fallback == "heuristic" else "default_model_fallback") + assert decision.get("classifier_input_budget_status") == "truncated" + assert decision["classifier_input_truncated"] is True + assert 0 < decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] + transport.handle_async_request.assert_awaited_once() + finally: + await handler.client.aclose() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 96c7c34621b..24736cdb34e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46205,6 +46205,16 @@ export interface components { * @enum {string} */ classifier_failure_reason?: "timeout" | "circuit_open" | "not_configured" | "unsupported_input" | "invalid_response" | "declined" | "classifier_error"; + /** Classifier Input Budget Status */ + classifier_input_budget_status?: string; + /** Classifier Input Budget Tokens */ + classifier_input_budget_tokens?: number; + /** Classifier Input Count Method */ + classifier_input_count_method?: string; + /** Classifier Input Tokens */ + classifier_input_tokens?: number; + /** Classifier Input Truncated */ + classifier_input_truncated?: boolean; /** Classifier Max Quality Gap */ classifier_max_quality_gap?: number; /** Classifier Model */ From e8f572b39bff31e8b59cc01e00d5adbb8e278405 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Fri, 30 Oct 2026 14:39:25 -0700 Subject: [PATCH 03/31] feat(sso): create the team_id_jwt_field team on UI SSO login when team_id_upsert is set (#49256) * feat(sso): create the team_id_jwt_field team on UI SSO login when team_id_upsert is set * fix(sso): keep the upserted team_id_jwt_field team out of result.team_ids so restricted_sso_group and the access-token group fallback are unchanged * fix(sso): only upsert the team_id_jwt_field team on UI SSO login when enable_jwt_auth is on The API JWT path only reaches team_id_upsert inside the enable_jwt_auth branch, so a deployment with a leftover litellm_jwtauth block and JWT auth off never created teams. Gate the UI SSO upsert the same way * fix(sso): read enable_jwt_auth through the proxy_server module so the basedpyright gate stays at its ceiling --- litellm/proxy/management_endpoints/types.py | 1 + litellm/proxy/management_endpoints/ui_sso.py | 67 +++- .../management/test_sso_sign_in.py | 110 ++++++- .../proxy/management_endpoints/test_ui_sso.py | 309 ++++++++++++++++-- 4 files changed, 448 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 221a6221fab..d4fd7a31d13 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -93,5 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: class CustomOpenID(OpenID): team_ids: list[str] + team_id_to_upsert: str | None = None user_role: LitellmUserRoles | None = None extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 59c2df91d77..53079d84fdd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -87,7 +87,7 @@ from litellm.proxy._types import ( TeamMemberAddRequest, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_team_object, get_user_object from litellm.proxy.auth.auth_utils import ( # noqa: F401 # legacy module exports _get_request_ip_address, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export get_request_ip_address, @@ -1266,6 +1266,7 @@ def generic_response_convertor( sso_jwt_handler: JWTHandler | None = None, role_mappings: Optional["RoleMappings"] = None, team_mappings: Optional["TeamMappings"] = None, + jwt_auth_enabled: bool = False, ) -> CustomOpenID: generic_user_id_attribute_name: Final = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username") generic_user_display_name_attribute_name: Final = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub") @@ -1286,14 +1287,15 @@ def generic_response_convertor( generic_user_email_attribute_name, ) - all_teams: Final = [] + token: Final = cast(dict[str, object], response) # cast-ok: the IdP userinfo JSON object fastapi-sso hands back + all_teams: Final = list[str]() if sso_jwt_handler is not None: - team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response)) + team_ids = sso_jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) if team_mappings is not None and team_mappings.team_ids_jwt_field is not None: team_ids_from_db_mapping: Final[list[str] | None] = get_nested_value( - data=cast(dict, response), + data=token, key_path=team_mappings.team_ids_jwt_field, default=[], ) @@ -1305,8 +1307,17 @@ def generic_response_convertor( team_ids_from_db_mapping, ) else: - team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response)) + team_ids = jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) + team_id_jwt_field: Final = jwt_handler.litellm_jwtauth.team_id_jwt_field + upsert_candidate: Final[object] = ( + get_nested_value(data=token, key_path=team_id_jwt_field) + if jwt_auth_enabled and jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None + else None + ) + team_id_to_upsert: Final = ( + upsert_candidate if isinstance(upsert_candidate, str) and upsert_candidate.strip() else None + ) # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider @@ -1375,6 +1386,7 @@ def generic_response_convertor( last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), team_ids=all_teams, + team_id_to_upsert=team_id_to_upsert, user_role=user_role, extra_fields=extra_fields, ) @@ -1614,6 +1626,8 @@ async def get_generic_sso_response( # make generic sso provider from fastapi_sso.sso.base import DiscoveryDocument + from litellm.proxy import proxy_server + received_response: dict | None = None sso_assertion: SSOIdentityAssertion | None = None @@ -1665,6 +1679,7 @@ async def get_generic_sso_response( sso_jwt_handler=sso_jwt_handler, role_mappings=role_mappings, team_mappings=team_mappings, + jwt_auth_enabled=proxy_server.general_settings.get("enable_jwt_auth", False) is True, ) SSOProvider: Final = create_generic_sso_provider( @@ -1884,6 +1899,29 @@ async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | No return frozenset(row.team_id for row in rows) +async def create_sso_team(team_id: str) -> None: + """Create the ``team_id_jwt_field`` team on UI SSO login the way the API JWT path does under ``team_id_upsert``.""" + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + return + await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + team_id_upsert=True, + ) + + +async def _create_sso_team_or_warn(team_id: str, create_team: Callable[[str], Awaitable[object]]) -> None: + """A concurrent login may create the team first, so a failed create is not fatal: the join still runs.""" + try: + await create_team(team_id) + except Exception as e: # noqa: BLE001 # fail-safe: a failed team create must not fail the login or skip the join + verbose_proxy_logger.warning("SSO login could not create team %s, trying to join it anyway: %s", team_id, e) + + async def add_missing_team_member( user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: Sequence[str], @@ -1892,21 +1930,30 @@ async def add_missing_team_member( [str, NewUserResponse | LiteLLM_UserTable], Awaitable[object] ] = create_team_member_add_task, existing_team_ids: Callable[[Sequence[str]], Awaitable[frozenset[str] | None]] = find_existing_team_ids, + team_id_to_upsert: str | None = None, + create_team: Callable[[str], Awaitable[object]] = create_sso_team, ) -> None: """ - Get missing teams (diff b/w user_info.team_ids and sso_teams) - Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them + - Join the ``team_id_to_upsert`` team too, creating it when it is missing (``team_id_upsert``), the way the API JWT path does; + it stays out of ``result.team_ids`` so restricted_sso_group and the access-token group fallback see only the claimed groups - Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order """ held_teams: Final = frozenset(user_info.teams or ()) - missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(sso_teams) if team_id not in held_teams) + claimed_teams: Final = sso_teams if team_id_to_upsert is None else (*sso_teams, team_id_to_upsert) + missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(claimed_teams) if team_id not in held_teams) if not missing_teams: return existing: Final = await existing_team_ids(missing_teams) teams_to_add: Final = ( - missing_teams if existing is None else tuple(team_id for team_id in missing_teams if team_id in existing) + missing_teams + if existing is None + else tuple(team_id for team_id in missing_teams if team_id in existing or team_id == team_id_to_upsert) ) for team_id in teams_to_add: + if team_id == team_id_to_upsert and (existing is None or team_id not in existing): + await _create_sso_team_or_warn(team_id, create_team) await add_team_member(team_id, user_info) @@ -3663,7 +3710,11 @@ class SSOAuthenticationHandler: verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition") return sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else [] - await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) + await add_missing_team_member( + user_info=user_info, + sso_teams=sso_teams, + team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None, + ) @staticmethod def verify_user_in_restricted_sso_group( diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index d0b6e2579c8..93cfa00ea13 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -16,7 +16,7 @@ import yaml from pydantic import JsonValue from tests.integration._support.client import Gateway, gateway_from_environment -from tests.integration._support.database import read_rows +from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState, owned_proxy from tests.integration._support.wire import Reply, Request, wire_server @@ -52,14 +52,29 @@ def _idp_reply(request: Request) -> Reply: "preferred_username": signed_in, "email": f"{signed_in}@example.com", "provider": IDP_PROVIDER, + "dept": _dept(signed_in), + "groups": [_group(signed_in)], } ).encode() ) return Reply(status=404, body=b'{"error": "not_found"}') +def _dept(subject: str) -> str: + return f"dept-{subject}" + + +def _group(subject: str) -> str: + return f"group-{subject}" + + @contextmanager -def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> Iterator[Gateway]: +def _sso_gateway( + directory: Path, + litellm_settings: Mapping[str, JsonValue], + general_settings: Mapping[str, JsonValue] | None = None, + db_sso_settings: Mapping[str, JsonValue] | None = None, +) -> Iterator[Gateway]: config: Final = directory / "sso_proxy_config.yaml" config.write_text( yaml.safe_dump( @@ -68,6 +83,7 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> "general_settings": { "master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL", + **(general_settings or {}), }, "litellm_settings": dict(litellm_settings), } @@ -83,7 +99,33 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> "OAUTHLIB_INSECURE_TRANSPORT": "1", } with owned_proxy(rig, directory, environment, config=config, remove_environment=("PROXY_BASE_URL",)) as gw: - yield gw + if db_sso_settings is None: + yield gw + return + with contextlib.ExitStack() as cleanups: + cleanups.callback( + write_rows, + 'DELETE FROM "LiteLLM_SSOConfig" WHERE id = %s', + ("sso_config",), + database_url=gw.database_url, + ) + stored: Final = _store_sso_settings(gw, {**_generic_sso_settings(environment), **db_sso_settings}) + assert stored.status_code == 200, stored.text + yield gw + + +def _generic_sso_settings(environment: Mapping[str, str]) -> Mapping[str, JsonValue]: + return { + "generic_client_id": environment["GENERIC_CLIENT_ID"], + "generic_client_secret": environment["GENERIC_CLIENT_SECRET"], + "generic_authorization_endpoint": environment["GENERIC_AUTHORIZATION_ENDPOINT"], + "generic_token_endpoint": environment["GENERIC_TOKEN_ENDPOINT"], + "generic_userinfo_endpoint": environment["GENERIC_USERINFO_ENDPOINT"], + } + + +def _store_sso_settings(gateway: Gateway, settings: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("PATCH", "/update/sso_settings", settings) def _sign_in(gateway: Gateway, subject: str) -> httpx.Response: @@ -165,3 +207,65 @@ def test_free_sso_admission_preserves_existing_users_and_serializes_the_last_sea refused: Final = _sign_in(gateway, "free-sso-extra") assert (above_capacity.status_code, refused.status_code) == (303, 403) assert gateway.rows('SELECT COUNT(*) AS count FROM "LiteLLM_UserTable"', ()) == [{"count": 6}] + + +@pytest.mark.parametrize( + ("enable_jwt_auth", "team_id_upsert", "team_exists_before", "expect_member"), + [(True, True, False, True), (False, True, False, False), (True, False, False, False), (True, False, True, True)], + ids=[ + "upsert-creates-missing-team", + "upsert-without-jwt-auth-skips-missing-team", + "no-upsert-skips-missing-team", + "no-upsert-joins-existing-team", + ], +) +def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert_is_on( + tmp_path: Path, enable_jwt_auth: bool, team_id_upsert: bool, team_exists_before: bool, expect_member: bool +) -> None: + subject: Final = f"sso-dept-user-{uuid.uuid4().hex[:12]}" + team_id: Final = _dept(subject) + jwtauth: Final = { + "enable_jwt_auth": enable_jwt_auth, + "litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": team_id_upsert}, + } + with _sso_gateway(tmp_path, {}, jwtauth) as gateway, contextlib.ExitStack() as cleanups: + cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id]}) + cleanups.callback(_delete_user, gateway, subject) + if team_exists_before: + created: Final = gateway.request("POST", "/team/new", {"team_id": team_id}) + assert created.status_code == 200, created.text + + callback: Final = _sign_in(gateway, subject) + + assert callback.status_code == 303, f"{callback.status_code} {callback.text}" + teams: Final = gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)) + user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + assert (len(teams), user) == (int(expect_member), [{"teams": [team_id] if expect_member else []}]) + + +def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_only_joins_groups( + tmp_path: Path, +) -> None: + subject: Final = f"sso-mapped-user-{uuid.uuid4().hex[:12]}" + team_id: Final = _dept(subject) + group_team_id: Final = _group(subject) + jwtauth: Final = { + "enable_jwt_auth": True, + "litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": True}, + } + team_mappings: Final = {"team_mappings": {"team_ids_jwt_field": "groups"}} + with ( + _sso_gateway(tmp_path, {}, jwtauth, db_sso_settings=team_mappings) as gateway, + contextlib.ExitStack() as cleanups, + ): + cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id, group_team_id]}) + cleanups.callback(_delete_user, gateway, subject) + + callback: Final = _sign_in(gateway, subject) + + assert callback.status_code == 303, f"{callback.status_code} {callback.text}" + teams: Final = gateway.rows( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id IN (%s, %s)', (team_id, group_team_id) + ) + user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + assert (teams, user) == ([{"team_id": team_id}], [{"teams": [team_id]}]) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 136cc86e67a..ba78fefd653 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -15,9 +15,15 @@ from fastapi import HTTPException, Request from pydantic import SecretStr import litellm +from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_UserTable, + NewUserResponse, + ProxyException, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -287,6 +293,13 @@ def test_get_microsoft_callback_response(stubbed_graph_api): assert sso_assertion.issuer == "https://idp.example.com" + +def _jwt_handler_double() -> MagicMock: + handler: Final = MagicMock(spec=JWTHandler) + handler.litellm_jwtauth = LiteLLM_JWTAuth() + return handler + + GRAPH_ERROR_BODIES: Final = ( {"error": {"code": "Authentication_MissingOrMalformed", "message": "Access Token missing or malformed."}}, {"error": {"code": "InvalidAuthenticationToken", "message": "IDX14100: JWT is not well formed."}}, @@ -1158,7 +1171,7 @@ def test_generic_response_convertor_normalizes_email(): } # Mock JWT handler - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -2226,7 +2239,7 @@ async def test_get_generic_sso_response_with_additional_headers(): # Arrange mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] generic_client_id = "test_client_id" @@ -2299,7 +2312,7 @@ async def test_get_generic_sso_response_with_empty_headers(): # Arrange mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] generic_client_id = "test_client_id" @@ -2359,8 +2372,8 @@ async def test_get_generic_sso_response_includes_token_claims_when_enabled(monke from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) - mock_sso_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() + mock_sso_jwt_handler = _jwt_handler_double() mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"] mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = [] @@ -2435,8 +2448,8 @@ async def test_get_generic_sso_response_does_not_include_token_claims_when_disab from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) - mock_sso_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() + mock_sso_jwt_handler = _jwt_handler_double() mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"] mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = [] access_token = pyjwt.encode({"upn": "token-user@example.com"}, "test-secret", algorithm="HS256") @@ -2578,7 +2591,7 @@ async def test_get_generic_sso_response_pkce_merges_token_claims_and_excludes_cr ): result, received_response, _, _ = await get_generic_sso_response( request=request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="test-client", redirect_url="http://test.com/callback", sso_jwt_handler=None, @@ -2600,7 +2613,7 @@ async def test_get_generic_sso_response_ignores_opaque_and_empty_token_claims(mo from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() userinfo = { "preferred_username": "userinfo-user", "email": "userinfo@example.com", @@ -4961,7 +4974,7 @@ class TestProcessSSOJWTAccessToken: @pytest.fixture def mock_jwt_handler(self): """Create a mock JWT handler for testing""" - mock_handler = MagicMock(spec=JWTHandler) + mock_handler = _jwt_handler_double() mock_handler.get_team_ids_from_jwt.return_value = ["team1", "team2", "team3"] return mock_handler @@ -5348,7 +5361,7 @@ class TestGenericResponseConvertorNestedAttributes: from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor # Mock JWT handler - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] # Payload with nested attributes structure @@ -5419,7 +5432,7 @@ class TestGenericResponseConvertorUserRole: from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] sso_response = { @@ -5449,7 +5462,7 @@ class TestGenericResponseConvertorUserRole: """ from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] sso_response = { @@ -6813,7 +6826,9 @@ class TestAddMissingTeamMember: # Verify add_missing_team_member was called with correct args mock_add_member.assert_called_once_with( - user_info=new_user_info, sso_teams=["entra-group-1", "entra-group-2"] + user_info=new_user_info, + sso_teams=["entra-group-1", "entra-group-2"], + team_id_to_upsert=None, ) @pytest.mark.asyncio @@ -7897,7 +7912,7 @@ def test_generic_response_convertor_with_extra_attributes(monkeypatch): "custom_field3": {"nested": "data"}, } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7931,7 +7946,7 @@ def test_generic_response_convertor_without_extra_attributes(monkeypatch): "custom_field2": "value2", } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7959,7 +7974,7 @@ def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypat "org_info": {"department": "Engineering", "manager": "Jane Smith"}, } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7986,7 +8001,7 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch): "email": "user@example.com", } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -8865,7 +8880,7 @@ class TestPKCEStateCookieBinding: ): await get_generic_sso_response( request=mock_request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="cid", redirect_url="https://proxy.example.com/sso/callback", sso_jwt_handler=None, @@ -8917,7 +8932,7 @@ class TestPKCEStateCookieBinding: ): await get_generic_sso_response( request=mock_request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="cid", redirect_url="https://proxy.example.com/sso/callback", sso_jwt_handler=None, @@ -8980,7 +8995,7 @@ class TestPKCEStateCookieBinding: }, ), ): - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] result, _, _, _ = await get_generic_sso_response( request=mock_request, @@ -9087,7 +9102,7 @@ async def test_debug_sso_callback_renders_full_jwt_claims(): patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), ): # Microsoft / Google envs may leak in from other tests — ensure only # the generic path runs. @@ -9173,7 +9188,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), ): for var in ("GENERIC_CLIENT_ID", "GOOGLE_CLIENT_ID"): os.environ.pop(var, None) @@ -9224,7 +9239,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), # test-quality-ok: debug endpoint reads proxy globals ] if force_inert: stack.append( @@ -9589,7 +9604,7 @@ async def test_pkce_arm_captures_sso_assertion(): }, ), ): - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] result, received_response, _, sso_assertion = await get_generic_sso_response( request=mock_request, @@ -9615,7 +9630,7 @@ async def test_verify_and_process_arm_captures_sso_assertion(): id_token = _ema_id_token() mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] mock_sso_instance = MagicMock() @@ -9680,7 +9695,7 @@ def _generic_sso_callback_request() -> Request: def _generic_sso_jwt_handler() -> MagicMock: - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] return jwt_handler @@ -11144,3 +11159,241 @@ async def test_cli_sso_login_never_serves_the_admin_login_form_on_a_ui_disabled_ body = response.body.decode() assert "Admin UI is Disabled" in body assert 'name="username"' not in body + + +def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler: + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=LiteLLM_JWTAuth(**jwtauth) + ) + return handler + + +@pytest.mark.parametrize( + ("jwtauth", "userinfo", "expected_team_ids", "expected_upsert"), + [ + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], "finance"), + ({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": " "}, [" "], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], None), + ( + {"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"}, + {}, + [], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_upsert": True}, + {"groups": ["g1", "g2"]}, + ["g1", "g2"], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_jwt_field": "dept", "team_id_upsert": True}, + {"groups": ["g1"], "dept": "finance"}, + ["g1", "finance"], + "finance", + ), + ], + ids=[ + "upsert-on", + "upsert-off", + "claim-missing", + "blank-claim-not-upserted", + "whitespace-claim-not-upserted", + "list-valued-claim-join-only", + "default-team-not-upserted", + "plural-claim-join-only", + "mixed", + ], +) +def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim( + jwtauth: dict[str, object], + userinfo: dict[str, object], + expected_team_ids: list[str], + expected_upsert: str | None, +) -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", **userinfo}, + jwt_handler=_upsert_jwt_handler(**jwtauth), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert) + + +def test_generic_sso_login_does_not_upsert_the_team_id_jwt_field_claim_when_jwt_auth_is_disabled() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "dept": "finance"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + jwt_auth_enabled=False, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (["finance"], None) + + +def test_generic_sso_login_marks_the_team_id_jwt_field_team_for_upsert_without_adding_it_to_team_ids() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "groups": ["g1"], "dept": "finance"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (["g1"], "finance") + + +def test_generic_sso_login_keeps_team_ids_empty_for_a_dept_only_user_so_the_access_token_group_fallback_runs() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "dept": "d1"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == ([], "d1") + + +def test_restricted_sso_group_still_rejects_a_user_whose_only_team_is_the_upserted_dept() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, generic_response_convertor + + response: Final = {"preferred_username": "alice", "email": "alice@example.com", "dept": "X", "groups": []} + result: Final = generic_response_convertor( + response=response, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + with pytest.raises(ProxyException) as denied: + SSOAuthenticationHandler.verify_user_in_restricted_sso_group( + general_settings={"ui_access_mode": {"type": "restricted_sso_group", "restricted_sso_group": "X"}}, + result=result, + received_response=response, + ) + + assert denied.value.code == "403" + + +def _recording_team_calls( + events: list[tuple[str, str]], +) -> tuple[Callable[[str, NewUserResponse | LiteLLM_UserTable], Awaitable[None]], Callable[[str], Awaitable[None]]]: + async def add(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None: + events.append(("join", team_id)) + + async def create(team_id: str) -> None: + events.append(("create", team_id)) + + return add, create + + +async def _only_group_teams_exist(team_ids: Sequence[str]) -> frozenset[str] | None: + return frozenset(team_id for team_id in team_ids if team_id.startswith("g")) + + +@pytest.mark.asyncio +async def test_sso_login_joins_the_upsert_team_from_team_id_to_upsert_and_creates_it_before_the_join() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["g1", "g2"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + team_id_to_upsert="finance", + create_team=create, + ) + + assert events == [("join", "g1"), ("join", "g2"), ("create", "finance"), ("join", "finance")] + + +@pytest.mark.asyncio +async def test_sso_login_joins_only_the_claimed_groups_when_team_id_upsert_is_off() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["g1", "g2"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + create_team=create, + ) + + assert events == [("join", "g1"), ("join", "g2")] + + +@pytest.mark.asyncio +async def test_sso_login_does_not_create_the_team_id_jwt_field_team_when_it_already_exists() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + async def every_team_exists(team_ids: Sequence[str]) -> frozenset[str] | None: + return frozenset(team_ids) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["finance"], + add_team_member=add, + existing_team_ids=every_team_exists, + team_id_to_upsert="finance", + create_team=create, + ) + + assert events == [("join", "finance")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "create_error", + [ + HTTPException(status_code=400, detail={"error": "Team id = finance already exists"}), + HTTPException(status_code=403, detail={"error": "team limit reached for this license"}), + ], + ids=["lost-the-create-race", "create-rejected"], +) +async def test_sso_login_still_joins_the_team_and_warns_when_creating_it_fails( + caplog: pytest.LogCaptureFixture, create_error: HTTPException +) -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, _ = _recording_team_calls(events) + + async def failed_create(team_id: str) -> None: + raise create_error + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["finance"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + team_id_to_upsert="finance", + create_team=failed_create, + ) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert events == [("join", "finance")] + assert len(warnings) == 1 and "finance" in warnings[0] and str(create_error.detail) in warnings[0] From 7e5a653c0129c5ba99e828d15e5e2ff1f5690e90 Mon Sep 17 00:00:00 2001 From: shrey Date: Fri, 30 Oct 2026 15:03:49 -0700 Subject: [PATCH 04/31] chore(cost-map): add Deepgram Aura-1 and Aura-2 text-to-speech rates (#49280) --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++++ model_prices_and_context_window.json | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1d18903bb07..843ab574da4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23078,6 +23078,24 @@ "output_cost_per_token_batches": 1e-06, "source": "https://developers.openai.com/api/docs/pricing" }, + "deepgram/aura": { + "input_cost_per_character": 1.5e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0150/1000 characters = $0.000015 per character", + "note": "Aura-1 text-to-speech voices (aura--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, + "deepgram/aura-2": { + "input_cost_per_character": 3e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.030/1000 characters = $0.00003 per character", + "note": "Aura-2 text-to-speech voices (aura-2--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, "deepgram/base": { "input_cost_per_second": 0.00020833, "litellm_provider": "deepgram", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1d18903bb07..843ab574da4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23078,6 +23078,24 @@ "output_cost_per_token_batches": 1e-06, "source": "https://developers.openai.com/api/docs/pricing" }, + "deepgram/aura": { + "input_cost_per_character": 1.5e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0150/1000 characters = $0.000015 per character", + "note": "Aura-1 text-to-speech voices (aura--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, + "deepgram/aura-2": { + "input_cost_per_character": 3e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.030/1000 characters = $0.00003 per character", + "note": "Aura-2 text-to-speech voices (aura-2--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, "deepgram/base": { "input_cost_per_second": 0.00020833, "litellm_provider": "deepgram", From faf9ee18e962c0dde19de46fc9190ab70f34a70f Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 15:05:32 -0700 Subject: [PATCH 05/31] fix(cost): preserve zero prices across batch, speech and search (#49294) * fix(cost): preserve zero cache-write prices in batch fallback * fix(cost): retain batch cache charges when ordinary input is free * fix(cost): preserve zero speech and Perplexity search rates * fix(cost): reject unset speech prices before zero defaults * test(cost): consolidate zero-price regression coverage --- CONTRIBUTING.md | 2 + litellm/cost_calculator.py | 13 +++- .../litellm_core_utils/llm_cost_calc/utils.py | 7 +- litellm/llms/perplexity/cost_calculator.py | 5 +- litellm/types/utils.py | 2 + litellm/utils.py | 9 ++- .../test_perplexity_cost_calculator.py | 51 ++++++++++++- tests/unit/test_cost_calculator.py | 73 ++++++++++++++----- ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 ++++ 9 files changed, 147 insertions(+), 30 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 531a78d8982..e956b27d5da 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -276,6 +276,8 @@ When the provider reports its own usage or spend numbers, compare the proxy's nu Provider work (a new provider, a new parameter, a new usage or cost field) is the most common kind of contribution, and it is where PRs get reworked most often. Follow these rules +**Preserve zero prices.** In [`model_prices_and_context_window.json`](model_prices_and_context_window.json), an explicit `0.0` means that the priced operation is free. A missing field or `null` means no rate is specified, so the relevant calculator may use a fallback. For example, when no batch input rate is configured, batch cache writes use `cache_creation_input_token_cost` if set, otherwise the regular input rate, then apply the batch discount. A zero ordinary input rate does not make separately priced cache reads or writes free. Test `rate is None` before substituting a fallback, since `row.get("cache_creation_input_token_cost") or fallback` would discard a valid zero. Speech keeps character pricing ahead of token pricing even when the character rate is zero. If neither rate is configured, speech cost calculation raises an error. Its model lookup preserves unset rates instead of treating compatibility defaults as explicit zero prices. Perplexity keeps a configured legacy search rate ahead of the current search rate, including zero. When registering custom prices, `None` leaves an existing rate unchanged and `0.0` overrides it + **Find the existing machinery before you add plumbing.** LiteLLM already has generic paths that move data between the endpoints, the provider transformations, and cost tracking. A second, hand-rolled path for the same data is a common reason a provider PR gets rewritten before merge. [`ARCHITECTURE.md`](ARCHITECTURE.md) shows where the layers live, and the decision records in [`docs/adrs/`](docs/adrs/AGENTS.md) explain why they are shaped that way. For example, [the ADR on provider usage fields](docs/adrs/adr_001_provider_usage_fields_ride_as_usage_extras.md) covers how provider-specific usage fields reach cost tracking. If you cannot find a mechanism for what you need, say so in the PR or ask in [#pr-review on Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA) before you build one **Bill a new provider usage field from the usage object.** Leave the field on the usage object as an extra and read it in that provider's `cost_calculator.py`. A provider without one gets a new `cost_calculator.py` and an arm for it in `cost_per_token` (`litellm/cost_calculator.py`), which picks the calculator by provider. If the provider's transform builds `Usage` from named fields, as the Anthropic, Gemini, and Bedrock Converse ones do, pass the new field there too. Streamed requests rebuild usage in `ChunkProcessor.calculate_usage` (`litellm/litellm_core_utils/streaming_chunk_builder_utils.py`), which does not copy extras yet the way the Responses-to-chat bridge does (minus a short exclusion list), so until it does, bill the field on streams by naming it there, as Anthropic's `inference_geo` is. A provider that reports its own charge for the request, as OpenRouter does in `usage.cost`, is billed at that charge from `_hidden_params` before any calculator runs, unless the deployment sets its own prices. When the field counts a built-in tool such as web search, add the provider to `get_cost_for_web_search_request` in `litellm/llms/__init__.py`, and add a gate to `litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py` only when no existing gate fires. Then show the field's cost in your proof on every endpoint it reaches, streaming included diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ee102e82ca4..f5b9014d87c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -623,7 +623,11 @@ def cost_per_token( ) if lyria_generation_cost is not None: return 0.0, lyria_generation_cost - speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) + speech_model_info: Final = cached_get_model_info_helper( + model=model_without_prefix, + custom_llm_provider=custom_llm_provider, + preserve_unset_costs=True, + ) cost_metric: Final = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 @@ -3084,7 +3088,7 @@ def batch_cost_calculator( batch_rates.cache_creation, batch_rates.cache_creation_rates_by_ttl, ) - elif input_cost_per_token: + elif input_cost_per_token is not None: details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = details["cache_hit_tokens"] cache_creation_tokens: Final = details["cache_creation_tokens"] @@ -3100,7 +3104,10 @@ def batch_cost_calculator( cache_read_cost_key: Final = get_service_tier_cost_key("cache_read_input_token_cost", None) total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2 - cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token + configured_cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") + cache_creation_cost: Final = ( + input_cost_per_token if configured_cache_creation_cost is None else configured_cache_creation_cost + ) total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2 text_rate: Final = _batch_or_half(batch_rates.output, output_cost_per_token) image_rate: Final = _batch_or_half( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 404dfab0a6d..50c26f4d4ca 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -28,6 +28,7 @@ from litellm.types.utils import ( DataResidency, ImageResponse, ModelInfo, + ModelInfoBase, PassthroughCallTypes, PromptTokensDetailsWrapper, ServiceTier, @@ -197,15 +198,15 @@ def get_billable_input_tokens(usage: Usage) -> int: def select_cost_metric_for_model( - model_info: ModelInfo, + model_info: ModelInfoBase, ) -> Literal["cost_per_character", "cost_per_token"]: """ Select 'cost_per_character' if model_info has 'input_cost_per_character' Select 'cost_per_token' if model_info has 'input_cost_per_token' """ - if model_info.get("input_cost_per_character"): + if model_info.get("input_cost_per_character") is not None: return "cost_per_character" - elif model_info.get("input_cost_per_token"): + elif model_info.get("input_cost_per_token") is not None: return "cost_per_token" else: raise ValueError( diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index c3d178771a5..c40a86b221b 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -100,8 +100,9 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non num_search_queries = getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 # Check both possible keys for search cost (legacy and current) - search_cost_value: Final = model_info.get("search_queries_cost_per_query") or model_info.get( - "search_context_cost_per_query" + legacy_search_cost: Final = model_info.get("search_queries_cost_per_query") + search_cost_value: Final = ( + model_info.get("search_context_cost_per_query") if legacy_search_cost is None else legacy_search_cost ) if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1a8f5adf154..0c3440e88f7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -464,6 +464,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): annotation_cost_per_page: float | None # for OCR models annotation_cost_per_page_batches: ReadOnly[float | None] search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool + search_queries_cost_per_query: ReadOnly[float | SearchContextCostPerQuery | None] web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None ) # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x) @@ -4168,6 +4169,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_video_per_second: float | None = None output_cost_per_audio_per_second: float | None = None search_context_cost_per_query: dict[str, Any] | None = None + search_queries_cost_per_query: float | SearchContextCostPerQuery | None = None google_maps_grounding_cost_per_query: float | None = None citation_cost_per_token: float | None = None cache_read_input_token_cost_above_272k_tokens: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 560d0c31e4a..c2633d7eb81 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6889,6 +6889,8 @@ def cached_get_model_info_helper( model: str, custom_llm_provider: str | None, api_base: str | None = None, + *, + preserve_unset_costs: bool = False, ) -> ModelInfoBase: """ _get_model_info_helper wrapped with lru_cache @@ -6899,6 +6901,7 @@ def cached_get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, + preserve_unset_costs=preserve_unset_costs, ) @@ -6962,6 +6965,7 @@ def get_model_info_helper( api_key: str | None = None, *, catalog_only: bool = False, + preserve_unset_costs: bool = False, ) -> ModelInfoBase: """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's @@ -7079,7 +7083,7 @@ def get_model_info_helper( verbose_logger.debug("model not found in litellm.model_cost, tried: %s", potential_model_names) raise ModelNotMappedError(_model_not_mapped_message(model, custom_llm_provider)) _input_cost_per_token: float | None = _model_info.get("input_cost_per_token") - if _input_cost_per_token is None: + if _input_cost_per_token is None and not preserve_unset_costs: # default value to 0, be noisy about this verbose_logger.debug( "model=%s, custom_llm_provider=%s has no input_cost_per_token in model_cost_map. Defaulting to 0.", @@ -7089,7 +7093,7 @@ def get_model_info_helper( _input_cost_per_token = 0 _output_cost_per_token: float | None = _model_info.get("output_cost_per_token") - if _output_cost_per_token is None: + if _output_cost_per_token is None and not preserve_unset_costs: # default value to 0, be noisy about this verbose_logger.debug( "model=%s, custom_llm_provider=%s has no output_cost_per_token in model_cost_map. Defaulting to 0.", @@ -7371,6 +7375,7 @@ def get_model_info_helper( supports_regex_lookaround=_model_info.get("supports_regex_lookaround", None), supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None), + search_queries_cost_per_query=_model_info.get("search_queries_cost_per_query"), web_search_billing_unit=_model_info.get("web_search_billing_unit", None), search_free_requests=_model_info.get("search_free_requests", None), google_maps_grounding_cost_per_query=_model_info.get("google_maps_grounding_cost_per_query", None), diff --git a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py index 380f85445e5..5e774b05207 100644 --- a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py @@ -9,6 +9,7 @@ import json import logging import math import os +from collections.abc import Mapping from datetime import datetime, timezone from typing import Final from unittest.mock import patch @@ -119,7 +120,8 @@ class TestPerplexityCostCalculator: assert completion_cost == 0.008 assert prompt_cost + completion_cost == 0.008 - def test_uses_perplexity_provided_cost_when_normalized_to_float(self): + @pytest.mark.parametrize("reported_cost", (0.0, 0.008)) + def test_uses_perplexity_provided_cost_when_normalized_to_float(self, reported_cost: float): """ Regression: for Responses API / Agent API models, `ResponseAPIUsage.parse_cost` (litellm/types/llms/openai.py) already flattens Perplexity's @@ -135,12 +137,55 @@ class TestPerplexityCostCalculator: perplexity/openai/gpt-5.2 before rates existed). """ usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) - usage.cost = 0.008 + usage.cost = reported_cost prompt_cost, completion_cost = perplexity_cost_per_token(model="sonar-pro", usage=usage) assert prompt_cost == 0.0 - assert completion_cost == 0.008 + assert completion_cost == reported_cost + + @pytest.mark.parametrize( + ("legacy_pricing", "expected_search_rate"), + ( + pytest.param({}, 0.01, id="missing"), + pytest.param({"search_queries_cost_per_query": None}, 0.01, id="null"), + pytest.param({"search_queries_cost_per_query": 0.0}, 0.0, id="free"), + pytest.param({"search_queries_cost_per_query": 0.004}, 0.004, id="paid"), + ), + ) + def test_search_cost_preserves_zero_and_defaults_only_when_unset( + self, + monkeypatch: pytest.MonkeyPatch, + legacy_pricing: Mapping[str, float | None], + expected_search_rate: float, + ) -> None: + model: Final = "perplexity/search-pricing-fixture" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "perplexity", + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": 8e-6, + "output_cost_per_reasoning_token": 3e-6, + "citation_cost_per_token": 2e-6, + "search_context_cost_per_query": {"search_context_size_low": 0.01}, + **legacy_pricing, + }, + ) + usage: Final = Usage( + prompt_tokens=100, + completion_tokens=50, + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=2), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=10), + citation_tokens=25, + ) + + prompt_cost, completion_cost = litellm.cost_per_token(model=model, usage_object=usage) + + assert prompt_cost == pytest.approx((100 + 25) * 2e-6) + assert completion_cost == pytest.approx(40 * 8e-6 + 10 * 3e-6 + 2 * expected_search_rate) @pytest.mark.parametrize( "reported_cost", diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index f10baccd1cb..c536522e736 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -37,7 +37,7 @@ from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject -from litellm.utils import _invalidate_model_cost_lowercase_map, supports_prompt_caching +from litellm.utils import _invalidate_model_cost_lowercase_map, cached_get_model_info_helper, supports_prompt_caching import os import threading from unittest.mock import MagicMock, patch @@ -4598,23 +4598,40 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate( assert completion_cost_value == pytest.approx(200 * 15e-6 / 2) -def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): +@pytest.mark.parametrize("input_rate", (3e-6, 0.0)) +@pytest.mark.parametrize( + ("cache_pricing", "expected_write_rate"), + [ + pytest.param({}, None, id="missing"), + pytest.param({"cache_creation_input_token_cost": None}, None, id="null"), + pytest.param({"cache_creation_input_token_cost": 0.0}, 0.0, id="free"), + pytest.param({"cache_creation_input_token_cost": 3.75e-6}, 3.75e-6, id="paid"), + ], +) +def test_batch_cost_calculator_preserves_cache_write_prices_and_defaults_only_when_unset( + cache_pricing: ModelInfo, + expected_write_rate: float | None, + input_rate: float, +) -> None: from litellm.cost_calculator import batch_cost_calculator - model_info: ModelInfo = { + model_info: Final[ModelInfo] = { "supported_openai_params": [], - "input_cost_per_token": 3e-6, + "input_cost_per_token": input_rate, "output_cost_per_token": 15e-6, "cache_read_input_token_cost": 3e-7, + **cache_pricing, } - prompt_cost, _ = batch_cost_calculator( + prompt_cost, completion_cost_value = batch_cost_calculator( usage=_batch_cache_usage(), - model="claude-sonnet-4-5-20250929", + model="batch-pricing-fixture", custom_llm_provider="anthropic", model_info=model_info, ) - assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) + write_rate: Final = input_rate if expected_write_rate is None else expected_write_rate + assert prompt_cost == pytest.approx((1000 * input_rate + 8000 * 3e-7 + 2000 * write_rate) / 2) + assert completion_cost_value == pytest.approx(200 * 15e-6 / 2) def test_completion_cost_bills_interactions_api_response(): @@ -7817,20 +7834,42 @@ def test_vertex_ai_mistral_predict_cost(usage): assert predictive_cost > 0 +@pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("call_type", ("speech", "aspeech")) +@pytest.mark.parametrize("model", ("speech-fixture", "openai/speech-fixture", "azure/speech-fixture")) @pytest.mark.parametrize( - "model", ["openai/tts-1", "azure/tts-1", "openai/gpt-4o-mini-tts"] + ("pricing", "expected_cost"), + ( + pytest.param({}, None, id="missing"), + pytest.param({"input_cost_per_character": None, "input_cost_per_token": None}, None, id="null"), + pytest.param({"input_cost_per_token": 0.0}, 0.0, id="free-tokens"), + pytest.param({"input_cost_per_character": 0.0, "input_cost_per_token": 2e-6}, 0.0, id="free-characters"), + pytest.param({"input_cost_per_token": 2e-6}, 100 * 2e-6, id="missing-character-rate"), + pytest.param({"input_cost_per_character": None, "input_cost_per_token": 2e-6}, 100 * 2e-6, id="null-character-rate"), + pytest.param({"input_cost_per_character": 3e-6, "input_cost_per_token": 2e-6}, 1000 * 3e-6, id="paid-characters"), + ), ) -def test_completion_cost_tts(model): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - cost = completion_cost( - model=model, - prompt="the quick brown fox jumped over the lazy dogs", - call_type="speech", +def test_registered_speech_prices_distinguish_unset_from_zero( + call_type: str, model: str, pricing: ModelInfo, expected_cost: float | None +) -> None: + provider: Final = "azure" if model.startswith("azure/") else "openai" + litellm.register_model( + {model: {"litellm_provider": provider, "mode": "audio_speech", **pricing}}, + persist_across_reloads=False, ) + compatibility_info: Final = cached_get_model_info_helper(model, provider) + assert compatibility_info["input_cost_per_token"] == (pricing.get("input_cost_per_token") or 0.0) + litellm.get_model_info(model) + response: Final = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=0, total_tokens=100)) - assert cost > 0 + if expected_cost is None: + with pytest.raises(ValueError, match="does not have 'input_cost_per_character' or 'input_cost_per_token'"): + completion_cost(completion_response=response, model=model, prompt="a" * 1000, call_type=call_type) + litellm.register_model({model: {"input_cost_per_token": 0.0}}, persist_across_reloads=False) + + assert completion_cost( + completion_response=response, model=model, prompt="a" * 1000, call_type=call_type + ) == pytest.approx(expected_cost or 0.0) def test_completion_cost_anthropic(): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 24736cdb34e..30d61f181cb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35894,6 +35894,8 @@ export interface components { search_context_cost_per_query?: { [key: string]: unknown; } | null; + /** Search Queries Cost Per Query */ + search_queries_cost_per_query?: number | components["schemas"]["SearchContextCostPerQuery"] | null; /** Stream Timeout */ stream_timeout?: number | string | null; /** Stream Timeout By Service Tier */ @@ -45445,6 +45447,17 @@ export interface components { } & { [key: string]: unknown; }; + /** SearchContextCostPerQuery */ + SearchContextCostPerQuery: { + /** Search Context Size High */ + search_context_size_high?: number; + /** Search Context Size Low */ + search_context_size_low?: number; + /** Search Context Size Medium */ + search_context_size_medium?: number; + } & { + [key: string]: unknown; + }; /** * SearchTool * @description Search tool configuration. @@ -51409,6 +51422,8 @@ export interface components { search_context_cost_per_query?: { [key: string]: unknown; } | null; + /** Search Queries Cost Per Query */ + search_queries_cost_per_query?: number | components["schemas"]["SearchContextCostPerQuery"] | null; /** Stream Timeout */ stream_timeout?: number | string | null; /** Stream Timeout By Service Tier */ From 298911d61bded3b81aa7d0face6e8538cb7114ed Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 15:25:47 -0700 Subject: [PATCH 06/31] fix(batches): cost an unmanaged gs:// batch through the deployment that created it (#49301) * fix(batches): cost an unmanaged gs:// batch through the deployment that created it The create now stores the router deployment id as model_id inside the managed object row's file_object JSON and CheckBatchCost reads it before falling back to the first deployment matching the base model in the gs:// path. Rows without the field keep the old lookup * fix(batches): tolerate a non-string or malformed file_object when reading the stored deployment * fix(batches): read the stored creating deployment from the writer, not a lagging read replica * fix(batches): resolve the writer database once so the type gate count does not grow * fix(batches): decode a dict-shaped file_object when reading the stored deployment, accumulate wire samples, annotate test locals --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 31 ++- .../proxy/hooks/managed_files.py | 50 ++++- ...ed_vertex_batch_deployment_routing_wire.py | 181 ++++++++++++++++++ .../proxy/test_managed_files_hook.py | 107 +++++++++++ .../common_utils/test_check_batch_cost.py | 47 +++++ 5 files changed, 408 insertions(+), 8 deletions(-) create mode 100644 tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index a1b4804ae29..a31e926256a 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast +from pydantic import TypeAdapter, ValidationError from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, @@ -50,6 +51,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool: CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) +_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( @@ -550,6 +552,10 @@ class CheckBatchCost: llm_provider: str, bare_model_name: str, ) -> Optional[Tuple[str, str]]: + stored_model_id: Final = self._get_stored_model_id(job) + if stored_model_id is not None and self._batch_deployment_exists(stored_model_id): + return stored_model_id, job.unified_object_id + deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider) if deployment_id is None: verbose_proxy_logger.info( @@ -662,23 +668,36 @@ class CheckBatchCost: ) @staticmethod - def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: - from litellm.types.utils import LiteLLMBatch - + def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]: file_object: Final = job.file_object try: - parsed_file_object: Final[object] = ( + return _FILE_OBJECT_ADAPTER.validate_python( json.loads(file_object) if isinstance(file_object, str) else file_object ) - except (json.JSONDecodeError, ValueError): + except (json.JSONDecodeError, ValueError, ValidationError): return None - if not isinstance(parsed_file_object, dict): + + @classmethod + def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]: + from litellm.types.utils import LiteLLMBatch + + parsed_file_object: Final = cls._parse_file_object(job) + if parsed_file_object is None: return None try: return LiteLLMBatch.model_validate(parsed_file_object).input_file_id except Exception: return None + @classmethod + def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]: + """The deployment that created the batch, written into file_object by the managed + files hook at create time. Older rows have no such key and fall back to matching + the bare model name against the router.""" + parsed_file_object: Final = cls._parse_file_object(job) + model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None + return model_id if isinstance(model_id, str) and model_id else None + async def _track_completed_batch_cost( self, job: "ManagedObjectRow", diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..0d46efbbb08 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -21,6 +21,7 @@ from typing import ( TypedDict, Union, cast, + runtime_checkable, ) from uuid import NAMESPACE_URL, uuid5 @@ -374,6 +375,12 @@ class _ManagedResourcePrismaClient(Protocol): def db(self) -> _ManagedResourceDatabase: ... +@runtime_checkable +class _WriterRoutedDatabase(Protocol): + @property + def writer(self) -> _ManagedResourceDatabase: ... + + class _SchedulerWithJobLookup(Protocol): def get_job(self, job_id: str) -> object: ... @@ -479,6 +486,11 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag return prisma_client.db.litellm_managedobjecttable +def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase: + """Reads that feed a write must see the create that just happened; a lagging read replica may not""" + return db.writer if isinstance(db, _WriterRoutedDatabase) else db + + _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00" @@ -512,6 +524,25 @@ def _without_model_output( ) +def _with_model_id( + file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse, + model_id: str | None, +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id}) + + +async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None: + row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id}) + if row is None: + return None + try: + parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object)) + except (json.JSONDecodeError, ValidationError): + return None + model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None + return model_id if isinstance(model_id, str) and model_id else None + + def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]: hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {} @@ -697,9 +728,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): persist_attribution: bool = False, create_if_missing: bool = True, batch_processed: bool = False, + creating_model_id: str | None = None, ) -> None: """Persist a managed object row, caching it and upserting it in the DB. + creating_model_id is the router deployment that created the object. It is kept + inside the stored file_object JSON as `model_id`, because the raw provider id an + unmanaged batch stores as unified_object_id says nothing about which deployment + made it, and the cost poller needs that deployment to read the result back. A + later observe of the row (create_if_missing=False) keeps the stored value rather + than overwriting it with whichever deployment served the retrieve. + persist_attribution is set only by the batch and background response creates, the callers that can speak for the creator; it gates the api_key, request_tags, and, for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or @@ -720,7 +759,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): It is written only in the upsert create branch. """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") - stored_object: Final = _without_model_output(file_object) + managed_database: Final = _managed_database(self.prisma_client) + managed_object_table: Final = managed_database.litellm_managedobjecttable + preserved_model_id: Final = ( + creating_model_id + if creating_model_id is not None or create_if_missing + else await _stored_model_id(_writer_database(managed_database).litellm_managedobjecttable, unified_object_id) + ) + stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id) litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -752,7 +798,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "status": file_object.status, "updated_by": user_api_key_dict.user_id, } - managed_object_table: Final = _managed_object_table(self.prisma_client) if not create_if_missing: await managed_object_table.update_many( where={"unified_object_id": unified_object_id}, @@ -2170,6 +2215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}), persist_attribution=is_batch_create, create_if_missing=is_batch_create, + creating_model_id=model_id if is_batch_create else None, ) if not is_batch_create: await self.store_batch_output_file_ownership( diff --git a/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py new file mode 100644 index 00000000000..91269e6557f --- /dev/null +++ b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py @@ -0,0 +1,181 @@ +"""Two Vertex deployments share a base model but live in different projects with different +buckets. An unmanaged gs:// batch created on the second one must be polled through that +deployment (LIT-10326). Before the fix the cost poller matched the base model name in the +gs:// path to the first deployment in the list and asked the wrong project for the job.""" + +import json +import re +from collections.abc import Callable +from contextlib import ExitStack +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.process import owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +LOCATION: Final = "us-central1" +BASE_MODEL: Final = "gemini-2.5-flash" +TEAMS: Final = ("a", "b") +_JOB_PATH: Final = re.compile( + r"^/v1/projects/(?P[^/]+)/locations/(?P[^/]+)/batchPredictionJobs(?:/(?P\d+))?$" +) + + +def _project(team: str) -> str: + return f"scripted-project-{team}" + + +def _deployment(team: str) -> str: + return f"{BASE_MODEL}-team-{team}" + + +def _input_file(team: str) -> str: + return f"gs://scripted-bucket-{team}/litellm-vertex-files/publishers/google/models/{BASE_MODEL}/{uuid4().hex}/input.jsonl" + + +def _job_number(team: str) -> str: + return f"{TEAMS.index(team) + 1}{uuid4().int % 10**17:017d}" + + +def _vertex_peer(jobs: dict[str, tuple[str, str]]) -> Callable[[Request], Reply]: + """A Vertex double that remembers which project each job was created in. A retrieve + from any other project answers the 404 real Vertex gives (the job does not exist + there), a retrieve from the right one reports the job still running.""" + + def respond(request: Request) -> Reply: + path: Final = _JOB_PATH.match(request.target.split("?", 1)[0]) + assert path is not None, request.target + assert request.headers.get("authorization") == "Bearer scripted-token", request.headers + project, location, job = path.group("project"), path.group("location"), path.group("job") + if request.method == "POST": + team: Final = project.rsplit("-", 1)[-1] + uri: Final = re.search(r'"(gs://[^"]+)"', request.body.decode()) + assert uri is not None and uri.group(1).startswith(f"gs://scripted-bucket-{team}/"), (project, request.body) + created, created_uri = jobs.setdefault(project, (_job_number(team), uri.group(1))) + return Reply(body=json.dumps(_job(project, location, created, created_uri, "JOB_STATE_PENDING")).encode()) + assert request.method == "GET" and job is not None, (request.method, request.target) + known: Final = jobs.get(project) + if known is None or known[0] != job: + return Reply( + status=404, + body=json.dumps( + { + "error": { + "code": 404, + "message": f"The BatchPredictionJob does not exist: {job}", + "status": "NOT_FOUND", + } + } + ).encode(), + ) + return Reply(body=json.dumps(_job(project, location, job, known[1], "JOB_STATE_RUNNING")).encode()) + + return respond + + +def _job(project: str, location: str, job: str, input_uri: str, state: str) -> dict[str, JsonValue]: + return { + "name": f"projects/{project}/locations/{location}/batchPredictionJobs/{job}", + "displayName": "litellm-vertex-batch-scripted", + "model": f"publishers/google/models/{BASE_MODEL}", + "inputConfig": {"gcsSource": {"uris": [input_uri]}, "instancesFormat": "jsonl"}, + "outputConfig": {"predictionsFormat": "jsonl", "gcsDestination": {"outputUriPrefix": input_uri.rsplit("/", 1)[0]}}, + "state": state, + "createTime": "2026-10-30T20:00:00.000000Z", + "updateTime": "2026-10-30T20:00:00.000000Z", + } + + +def _owned_config(path: Path, gateway: Gateway, wire_url: str) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": _deployment(team), + "litellm_params": { + "model": f"vertex_ai/{BASE_MODEL}", + "api_base": wire_url, + "vertex_project": _project(team), + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(_project(team), gateway.upstream_url.rstrip("/")), + "gcs_bucket_name": f"scripted-bucket-{team}", + }, + } + for team in TEAMS + ], + "general_settings": { + **object_value(config.get("general_settings") or {}), + "track_unmanaged_batch_cost": True, + }, + "litellm_settings": { + **object_value(config.get("litellm_settings") or {}), + "enable_loadbalancing_on_batch_endpoints": True, + "user_url_allowed_hosts": ["127.0.0.1"], + }, + } + ) + ) + return path + + +def _create_batch(gateway: Gateway, key: str, team: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": _input_file(team), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": _deployment(team), + }, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _retrieve_targets(requests: tuple[Request, ...], job: str) -> tuple[str, ...]: + return tuple(request.target for request in requests if request.method == "GET" and request.target.endswith(f"/{job}")) + + +@pytest.mark.timeout(240) +def test_unmanaged_batch_on_the_second_deployment_sharing_a_base_model_is_polled_through_its_own_project( + gateway: Gateway, tmp_path: Path +) -> None: + jobs: Final[dict[str, tuple[str, str]]] = {} # mutable-ok: the wire double records which project created each job + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_vertex_peer(jobs))) + config: Final = _owned_config(tmp_path / "two-vertex-teams.yaml", gateway, wire.url) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + key: Final = scenario.key(models=[_deployment(team) for team in TEAMS]) + + batch_a: Final = _create_batch(owned.gateway, key, "a") + batch_b: Final = _create_batch(owned.gateway, key, "b") + assert (batch_a, batch_b) == (jobs[_project("a")][0], jobs[_project("b")][0]), (batch_a, batch_b, jobs) + + seen: Final[list[Request]] = [] # mutable-ok: drain empties the wire queue, so samples accumulate across retries + + def polled_so_far() -> tuple[Request, ...]: + seen.extend(wire.drain()) + return tuple(seen) + + polled: Final = eventually( + polled_so_far, + lambda requests: bool(_retrieve_targets(requests, batch_b)) and bool(_retrieve_targets(requests, batch_a)), + seconds=90, + ) + for team, batch in (("a", batch_a), ("b", batch_b)): + assert all( + target.startswith(f"/v1/projects/{_project(team)}/") for target in _retrieve_targets(polled, batch) + ), (team, _retrieve_targets(polled, batch)) + diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index f2df14cc8a3..89fff365c3f 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1437,6 +1437,8 @@ def _make_object_store_instance(): mock_prisma = MagicMock() mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = mock_prisma.db return ( PROXY_LiteLLMManagedFiles( @@ -2161,3 +2163,108 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1") create_data = file_table.upsert.await_args.kwargs["data"]["create"] assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1") + + +@pytest.mark.asyncio +async def test_create_stores_the_creating_deployment_in_file_object(): + """Regression (LIT-10326): an unmanaged gs:// batch stores the raw Vertex job id, which + says nothing about which of two deployments sharing a base model created it. The create + must keep the router's deployment id inside the stored file_object so the cost poller + reads the result through that deployment instead of the first base-model match.""" + managed_files, mock_prisma = _make_object_store_instance() + creator: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="validating"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + create_if_missing=True, + creating_model_id="deploy-team-b", + ) + + created: Final = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"] + assert json.loads(created["file_object"])["model_id"] == "deploy-team-b" + + +@pytest.mark.asyncio +async def test_poll_refresh_keeps_the_stored_creating_deployment(): + """A later observe of the row (a GET through whichever deployment the client named) + rewrites file_object; it must carry the create-time model_id forward, not drop it or + replace it with the deployment that served the retrieve.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = json.dumps({**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"}) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]) + assert written["model_id"] == "deploy-team-b" + assert written["status"] == "completed" + + +@pytest.mark.asyncio +async def test_poll_refresh_reads_the_stored_creating_deployment_from_the_writer(): + """With a read replica configured, a retrieve that lands before the replica has the create + must still find the create-time model_id, so the lookup goes to the writer. Prisma may hand + the Json column back already decoded, so the row here is a dict rather than JSON text.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = {**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"} + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = MagicMock() + mock_prisma.db.writer.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]) + assert written["model_id"] == "deploy-team-b" + mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_id(): + """The post-call hook on a batch create hands the router's model_id to the object store; + a retrieve of the same batch hands None so it cannot overwrite the creator.""" + from litellm_enterprise.proxy.hooks.managed_files import BATCH_CREATE_HIDDEN_PARAM + + managed_files: Final = _make_managed_files_instance() + created: Final = _make_batch_response(model_id="deploy-team-b", output_file_id=None, status="validating") + created._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True + retrieved: Final = _make_batch_response(model_id="deploy-team-a", output_file_id=None, status="in_progress") + + with patch("litellm.proxy.proxy_server.llm_router", MagicMock()): + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=created + ) + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=retrieved + ) + + first, second = managed_files.store_unified_object_id.await_args_list + assert first.kwargs["creating_model_id"] == "deploy-team-b" + assert second.kwargs["creating_model_id"] is None diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 4861b2687de..5b37cee88c5 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1771,6 +1771,53 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-ft", "8823717160934178816") + def _two_team_deployments(self): + """Two Vertex deployments sharing the base model, team A listed first.""" + return [ + { + "model_name": f"gemini-2.5-flash-team-{team}", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "custom_llm_provider": "vertex_ai"}, + "model_info": {"id": f"deploy-team-{team}"}, + } + for team in ("a", "b") + ] + + def test_flag_on_prefers_the_deployment_stored_on_the_row_over_the_first_base_model_match(self): + """Regression (LIT-10326): two deployments share gemini-2.5-flash with different + buckets. A batch created on team B's deployment stores deploy-team-b in file_object, + so the poller must read it back through team B, not the first match (team A).""" + router: Final = MagicMock() + router.registry.return_value.resolve.return_value = None + router.registry.return_value.deployments.return_value = self._two_team_deployments() + router.registry.return_value.deployment = MagicMock(return_value=self._vertex_deployment()) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-team-b"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-b", "8823717160934178816") + router.registry.return_value.deployment.assert_called_once_with("deploy-team-b") + + def test_flag_on_falls_back_to_base_model_match_when_the_stored_deployment_is_gone(self): + """The stored deployment was removed from the config after the batch was created: + the row still routes through today's base-model lookup instead of being skipped.""" + router: Final = MagicMock() + router.registry.return_value.resolve.return_value = None + router.registry.return_value.deployments.return_value = self._two_team_deployments() + router.registry.return_value.deployment = MagicMock(return_value=None) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-removed"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-a", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" From 991ea7f7e8419feb9be44c1aa67fdc95289c1762 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 30 Oct 2026 15:47:06 -0700 Subject: [PATCH 07/31] refactor(ui): move policies, tool policies and guardrails monitor text onto i18n (#49343) Moves the page-authored English on /policies, /tool-policies and /guardrails-monitor into react-i18next t() calls backed by en.json, and adds a small data test for en.json. The rendered English is unchanged. --- .../_components/EvaluationSettingsModal.tsx | 36 +- .../_components/GuardrailDetail.tsx | 46 +- .../_components/GuardrailUsageBreakdown.tsx | 111 ++- .../_components/GuardrailsOverview.tsx | 91 +- .../_components/ScoreChart.tsx | 6 +- .../(dashboard)/guardrails-monitor/page.tsx | 4 +- .../policies/_components/AttachmentTable.tsx | 15 +- .../_components/AttachmentTableColumns.tsx | 76 +- .../policies/_components/PolicyTable.tsx | 15 +- .../_components/PolicyTableColumns.tsx | 65 +- .../policies/_components/TokenSelect.tsx | 6 +- .../_components/add_attachment_form.tsx | 240 ++--- .../policies/_components/add_policy_form.tsx | 215 +++-- .../_components/ai_suggestion_modal.tsx | 154 +-- .../_components/guardrail_selection_modal.tsx | 78 +- .../policies/_components/impact_popover.tsx | 28 +- .../_components/impact_preview_alert.tsx | 58 +- .../policies/_components/index.tsx | 157 +-- .../_components/pipeline_flow_builder.tsx | 530 +++++----- .../policies/_components/policy_info.tsx | 56 +- .../policies/_components/policy_templates.tsx | 31 +- .../_components/policy_test_panel.tsx | 99 +- .../_components/template_parameter_modal.tsx | 61 +- .../GuardrailsMonitor/CalcPopover.tsx | 7 +- .../GuardrailsMonitor/LogViewer.tsx | 38 +- .../GuardrailsMonitor/UnpricedNote.tsx | 26 +- .../GuardrailsMonitor/usageUnits.ts | 21 +- .../src/components/ToolDetail.tsx | 86 +- .../ToolPolicies/ToolPoliciesPanel.tsx | 51 +- .../ToolPolicies/ToolPoliciesTable.tsx | 82 +- .../ToolPolicies/ToolPoliciesTableColumns.tsx | 41 +- .../src/components/ToolPoliciesView.tsx | 6 +- .../policies/useDeletePolicyAttachment.ts | 6 +- ui/litellm-dashboard/src/i18n/locales.test.ts | 54 ++ ui/litellm-dashboard/src/i18n/locales/en.json | 911 ++++++++++++++++++ 35 files changed, 2455 insertions(+), 1052 deletions(-) create mode 100644 ui/litellm-dashboard/src/i18n/locales.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx index 78091584cde..12ff10bc0c5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx @@ -1,6 +1,7 @@ import { useTerminology } from "@/contexts/TerminologyContext"; import { Play } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { Button } from "@/components/ui/button"; @@ -48,6 +49,7 @@ export function EvaluationSettingsModal({ accessToken, onRunEvaluation, }: EvaluationSettingsModalProps) { + const { t } = useTranslation(); const term = useTerminology(); const [prompt, setPrompt] = useState(DEFAULT_PROMPT); const [schema, setSchema] = useState(DEFAULT_SCHEMA); @@ -94,11 +96,11 @@ export function EvaluationSettingsModal({ !nextOpen && onClose()}> - Evaluation Settings + {t("guardrailsMonitor.evaluationSettings.title")} {guardrailName - ? `${term("Configure AI evaluation for")} ${guardrailName}` - : term("Configure AI evaluation for re-running on logs")} + ? `${term(t("guardrailsMonitor.evaluationSettings.descriptionPrefix"))} ${guardrailName}` + : term(t("guardrailsMonitor.evaluationSettings.descriptionGeneric"))} @@ -106,10 +108,10 @@ export function EvaluationSettingsModal({