diff --git a/enterprise/litellm_enterprise/proxy/auth/route_checks.py b/enterprise/litellm_enterprise/proxy/auth/route_checks.py index fc57292a8d2..479537c4a90 100644 --- a/enterprise/litellm_enterprise/proxy/auth/route_checks.py +++ b/enterprise/litellm_enterprise/proxy/auth/route_checks.py @@ -2,6 +2,8 @@ import os from fastapi import HTTPException, status +from litellm.proxy._types import MODEL_LIST_ROUTES + class EnterpriseRouteChecks: @staticmethod @@ -43,7 +45,7 @@ class EnterpriseRouteChecks: # Routes that should remain accessible even when LLM API endpoints are disabled. # These are read-only model listing routes needed by the Admin UI. - LLM_API_EXEMPT_ROUTES = ["/models", "/v1/models"] + LLM_API_EXEMPT_ROUTES = MODEL_LIST_ROUTES @staticmethod def should_call_route(route: str): diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a0da0f6ddca..8331c68fe24 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _SERVICE_TIER_TO_COST_KEY_SUFFIX, BilledTokenRates, CostCalculatorUtils, + calculate_cache_writing_cost, calculate_cost_component, generic_cost_per_character, generic_cost_per_token, @@ -1550,6 +1551,8 @@ def completion_cost( """ try: call_type = infer_call_type(call_type, completion_response) or "completion" + if isinstance(completion_response, DecisionsResponse) and completion_response.usage is None: + raise ValueError("Cannot calculate Decisions cost: provider did not report usage") if pricing_provider is not None: custom_llm_provider = pricing_provider # rebind-ok: model_info.pricing_provider overrides the routing provider for cost-map lookups @@ -2969,7 +2972,12 @@ def batch_cost_calculator( total_completion_cost = 0.0 if batch_rates.input is not None: total_prompt_cost = _batch_prompt_cost( - usage, model_info, batch_rates.input, batch_rates.cache_read, batch_rates.cache_creation + usage, + model_info, + batch_rates.input, + batch_rates.cache_read, + batch_rates.cache_creation, + batch_rates.cache_creation_rates_by_ttl, ) elif input_cost_per_token: details: Final = parse_prompt_tokens_details(usage) @@ -3014,11 +3022,13 @@ def _batch_prompt_cost( input_rate: float, cache_read_rate: float | None, cache_creation_rate: float | None, + cache_creation_rates_by_ttl: Mapping[str, float], ) -> float: details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0 cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0 - written_tokens: Final = details["cache_creation_tokens"] if cache_creation_rate is not None else 0 + bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl) + written_tokens: Final = details["cache_creation_tokens"] if bills_cache_writes else 0 audio_tokens, image_tokens, video_tokens = ( details["audio_tokens"], details["image_tokens"], @@ -3035,7 +3045,16 @@ def _batch_prompt_cost( + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0) + cached_audio_tokens * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0) - + written_tokens * (cache_creation_rate or 0.0) + + ( + calculate_cache_writing_cost( + written_tokens, + details["cache_creation_token_details"], + input_rate if cache_creation_rate is None else cache_creation_rate, + cache_creation_rates_by_ttl, + ) + if bills_cache_writes + else 0.0 + ) ) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 4be5cabac04..00c6ce97e78 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -928,6 +928,71 @@ class _CarriesHiddenParams(Protocol): def hidden_params(self, hidden_params: dict[str, object]) -> None: ... # mutable-ok: API requires mutation +_USAGE_BLOCK_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _usage_field(usage: object, name: str) -> object: + if isinstance(usage, Mapping): + return _USAGE_BLOCK_ADAPTER.validate_python(usage).get(name) + return getattr(usage, name, None) + + +def _cost_number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + return float(value) + + +def _reported_cost_value(value: object) -> float | None: + """A reported cost is a number, or a breakdown whose total lives under ``total_cost``.""" + raw: Final[object] = value + if isinstance(value, Mapping): + return _cost_number(_usage_field(raw, "total_cost")) + return _cost_number(value) + + +def response_cost_in_hidden_params(response: _CarriesHiddenParams) -> float | None: + additional_headers: Final[object] = response.hidden_params.get("additional_headers") + if not isinstance(additional_headers, Mapping): + return None + return _cost_number(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) + + +def override_stamped_response_cost(response: _CarriesHiddenParams, cost: float) -> None: + """Replace the cost an earlier builder stamped so every reader of the hidden params agrees with the header.""" + set_response_cost_in_hidden_params(response, cost) + hidden_params: Final = response.hidden_params + hidden_params["response_cost"] = cost + + +def chunk_usage(chunk: object) -> object: + """Stream chunks reach the builders both as mappings and as ModelResponseStream objects.""" + if isinstance(chunk, Mapping): + return _USAGE_BLOCK_ADAPTER.validate_python(chunk).get("usage") + return getattr(chunk, "usage", None) + + +_MESSAGES_REPORTED_COST_PROVIDERS: Final = frozenset({"opencode_go"}) + + +def native_messages_reported_cost(usage: object, custom_llm_provider: object) -> float | None: + """Only a Messages upstream known to price its own responses is trusted, every other one keeps the price map.""" + if not isinstance(custom_llm_provider, str) or custom_llm_provider not in _MESSAGES_REPORTED_COST_PROVIDERS: + return None + return provider_reported_usage_cost(usage, custom_llm_provider) + + +def provider_reported_usage_cost(usage: object, custom_llm_provider: str | None) -> float | None: + """The request cost a provider states on its usage block, as a number or a breakdown with ``total_cost``. + + OpenCode Go reports the model's list price (peak hours and long-context tiers included) under + ``usage.opencode.list_cost``, while its ``usage.cost`` is what hit the prepaid balance, 0 inside the plan. + """ + if custom_llm_provider == "opencode_go": + return _reported_cost_value(_usage_field(_usage_field(usage, "opencode"), "list_cost")) + return _reported_cost_value(_usage_field(usage, "cost")) + + def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: float | None) -> None: """Record a provider-reported cost where the cost calculator looks before the price map.""" if cost is None: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fc8d1ecc8c1..9e0461e951d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -77,6 +77,8 @@ from litellm.litellm_core_utils.classifier_logging import ( from litellm.litellm_core_utils.core_helpers import ( get_provider_response_headers_from_hidden_params, is_expected_client_error, + native_messages_reported_cost, + provider_reported_usage_cost, proxy_stamped_kubernetes_pod_routing, proxy_stamped_used_client_oauth_token, reconstruct_model_name, @@ -4756,7 +4758,12 @@ class Logging(LiteLLMLoggingBaseClass): ): ## return unified Usage object if isinstance(result.response.usage, ResponseAPIUsage): - set_response_cost_in_hidden_params(result.response, result.response.usage.cost) + set_response_cost_in_hidden_params( + result.response, + provider_reported_usage_cost( + result.response.usage, self.model_call_details.get("custom_llm_provider") + ), + ) transformed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage( result.response.usage ) @@ -4822,8 +4829,11 @@ class Logging(LiteLLMLoggingBaseClass): bridge is replaced by the ModelResponse its spend row is built from, anything else is priced as is.""" served: Final[object] = result - if self.bridged_model_response is None or not isinstance(result, dict): - return result + call_type: Final[object] = getattr(self, "call_type", None) + if not isinstance(result, dict): + return served + if self.bridged_model_response is None and call_type != CallTypes.anthropic_messages.value: + return served return self._anthropic_messages_logged_response(result=served) def _anthropic_messages_logged_response(self, result: object) -> ModelResponse: @@ -4866,7 +4876,13 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(result, ResponsesAPIResponse): return self._translate_responses_api_response_to_model_response(result) - provider_response_id: Final = _provider_response_id(result) + raw_body: Final[object] = result + served_body: Final = _LITELLM_PARAMS.validate_python(raw_body) if isinstance(result, dict) else None + reported_cost: Final = native_messages_reported_cost( + served_body.get("usage") if served_body is not None else None, + self.model_call_details.get("custom_llm_provider"), + ) + provider_response_id: Final = _provider_response_id(raw_body) bridged: Final = self.bridged_model_response if bridged is not None and isinstance(result, dict): return bridged if provider_response_id is None else bridged.model_copy(update={"id": provider_response_id}) @@ -4891,7 +4907,7 @@ class Logging(LiteLLMLoggingBaseClass): else: from litellm.types.llms.anthropic import AnthropicResponse - pydantic_result: Final = AnthropicResponse.model_validate(result) + pydantic_result: Final = AnthropicResponse.model_validate({"stop_sequence": None, **(served_body or {})}) import httpx result = litellm.AnthropicConfig().transform_parsed_response( @@ -4904,6 +4920,7 @@ class Logging(LiteLLMLoggingBaseClass): json_mode=None, speed=self.optional_params.get("speed") if self.optional_params else None, ) + set_response_cost_in_hidden_params(result, reported_cost) return result def _translate_responses_api_response_to_model_response(self, result: ResponsesAPIResponse) -> ModelResponse: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 5d4b1a82759..3b49f239afe 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -104,7 +104,8 @@ _BATCH_RATE_PREFIXES: Final = ( "cache_read_input_token_cost", "cache_creation_input_token_cost", ) -_BATCH_TIER_KEY: Final = re.compile(rf"^({'|'.join(_BATCH_RATE_PREFIXES)})_above_(\d+k?)_tokens{_BATCH_KEY_SUFFIX}$") +_BATCH_TIER_PREFIXES: Final = (*_BATCH_RATE_PREFIXES, *_TTL_PRICED_CACHE_WRITE_RATE_KEYS.values()) +_BATCH_TIER_KEY: Final = re.compile(rf"^({'|'.join(_BATCH_TIER_PREFIXES)})_above_(\d+k?)_tokens{_BATCH_KEY_SUFFIX}$") _NON_STANDARD_THRESHOLD_SUFFIXES: Final = (*_SERVICE_TIER_SUFFIXES, _BATCH_KEY_SUFFIX) @@ -314,6 +315,7 @@ class BatchCostRates: output: float | None cache_read: float | None cache_creation: float | None + cache_creation_rates_by_ttl: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) def _batch_rate(model_info: ModelInfo, key: str) -> float | None: @@ -369,6 +371,13 @@ def get_batch_cost_rates(model_info: ModelInfo, usage: Usage, custom_llm_provide output=_batch_rate_for_prefix(model_info, "output_cost_per_token", usage, inclusive), cache_read=_batch_rate_for_prefix(model_info, "cache_read_input_token_cost", usage, inclusive), cache_creation=_batch_rate_for_prefix(model_info, "cache_creation_input_token_cost", usage, inclusive), + cache_creation_rates_by_ttl=MappingProxyType( + { + ttl: rate + for ttl, key in _TTL_PRICED_CACHE_WRITE_RATE_KEYS.items() + if (rate := _batch_rate_for_prefix(model_info, key, usage, inclusive)) is not None + } + ), ) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 50789406055..68cdc204af4 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict @@ -1352,14 +1353,27 @@ def stream_chunk_builder_text_completion( _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) -def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: +def _logged_provider(logging_obj: Optional["Logging"]) -> object: if logging_obj is None: - return False + return None provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider") - return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS + return provider -def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None: +def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: + return _logged_provider(logging_obj) in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS + + +def _stream_builder_response_cost( + response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object] +) -> float | None: + if _logged_provider(logging_obj) == LlmProviders.OPENCODE_GO.value: + # usage.opencode only survives on the chunks, the assembled Usage keeps declared fields + reported: Final = ( + provider_reported_usage_cost(chunk_usage(chunk), LlmProviders.OPENCODE_GO.value) + for chunk in reversed(chunks) + ) + return next((cost for cost in reported if cost is not None), None) usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj): return float(usage_cost) @@ -1397,8 +1411,10 @@ def _set_hidden_response_cost(response: ModelResponse, response_cost: float) -> hidden_params["response_cost"] = response_cost -def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None: - response_cost: Final = _stream_builder_response_cost(response, logging_obj) +def _set_stream_builder_response_cost( + response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object] +) -> None: + response_cost: Final = _stream_builder_response_cost(response, logging_obj, chunks) if response_cost is None: return _set_hidden_response_cost(response, response_cost) @@ -1550,7 +1566,11 @@ def stream_chunk_builder( break _stamp_streaming_usage_cost(usage, response, logging_obj) - _set_stream_builder_response_cost(response, logging_obj) + _set_stream_builder_response_cost( + response, + logging_obj, + chunks, # pyright: ignore[reportUnknownArgumentType] # legacy chunk list is untyped + ) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response @@ -1734,7 +1754,11 @@ def stream_chunk_builder( break _stamp_streaming_usage_cost(usage, response, logging_obj) - _set_stream_builder_response_cost(response, logging_obj) + _set_stream_builder_response_cost( + response, + logging_obj, + chunks, # pyright: ignore[reportUnknownArgumentType] # legacy chunk list is untyped + ) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fdcdca151be..2df470edc59 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -24,6 +24,7 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, @@ -72,7 +73,9 @@ _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None -_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) +_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset( + {LlmProviders.OPENROUTER.value, LlmProviders.OPENCODE_GO.value} +) def _next_sync_or_exhausted(it: Iterator[object]) -> object: @@ -1953,11 +1956,17 @@ class CustomStreamWrapper: def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", custom_llm_provider: str | None, + chunks: Sequence[object] = (), ) -> None: + """The builder's usage only keeps declared fields, so provider blocks like usage.opencode are read off the chunks.""" if custom_llm_provider not in _USAGE_COST_HEADER_PROVIDERS: return - _usage: Final[Usage | None] = getattr(response, "usage", None) - _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) + usages: Final = ( + getattr(response, "usage", None), + *(chunk_usage(chunk) for chunk in reversed(chunks)), + ) + costs: Final = (provider_reported_usage_cost(usage, custom_llm_provider) for usage in usages) + _cost: Final = next((cost for cost in costs if cost is not None), None) if _cost is not None: if "additional_headers" not in response.hidden_params: response.hidden_params["additional_headers"] = {} @@ -2099,7 +2108,9 @@ class CustomStreamWrapper: complete_streaming_response: Final = self._build_complete_streaming_response() response = self._new_usage_chunk() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) + self._propagate_usage_cost_to_hidden_params( + complete_streaming_response, self.custom_llm_provider, self.chunks + ) setattr( response, @@ -2318,7 +2329,9 @@ class CustomStreamWrapper: complete_streaming_response: Final = await self._async_build_complete_streaming_response() response: Final = self._new_usage_chunk() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) + self._propagate_usage_cost_to_hidden_params( + complete_streaming_response, self.custom_llm_provider, self.chunks + ) setattr( response, diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index dadd4d8558d..4eb53f35a03 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -66,6 +66,11 @@ class _ChunkStream(Protocol): def __aiter__(self) -> "AsyncIterator[ModelResponseStream]": ... +@runtime_checkable +class _AsyncCloseable(Protocol): + async def aclose(self) -> None: ... + + @runtime_checkable class _FailureLoggingWaiter(Protocol): @property @@ -197,6 +202,14 @@ class _CombinedChunkSplitter: self._buffer: deque[ModelResponseStream] = deque() self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({}) self._live_call_index: int | None = None + self._closed = False + + async def aclose(self) -> None: + if self._closed: + return + self._closed = True + if isinstance(self._stream, _AsyncCloseable): + await self._stream.aclose() @property def has_dispatched_failure_logging(self) -> bool: @@ -1239,6 +1252,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # For non-dict chunks, forward the original value unchanged yield chunk + async def aclose(self) -> None: + await self._combined_chunk_splitter.aclose() + async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]: """ Async version of anthropic_sse_wrapper. @@ -1481,4 +1497,7 @@ class AnthropicSSEStream(AsyncIterator[bytes]): return await self._byte_stream.__anext__() async def aclose(self) -> None: - await self._byte_stream.aclose() + try: + await self._byte_stream.aclose() + finally: + await self._anthropic_wrapper.aclose() diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index a122aeb3033..14c6a2d2a66 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -376,12 +376,14 @@ def _ir_answer(question: DecisionsIRQuestion, answer: DecisionAnswer | None) -> def systemone_response_to_ir(response: DecisionsResponse, request: DecisionsIRRequest) -> DecisionsIRResponse: - usage: Final = response.usage or DecisionsUsage() + usage: Final = response.usage keyed_questions: Final = zip(systemone_keys(request.questions), request.questions, strict=True) return DecisionsIRResponse( model=response.model, answers=tuple(_ir_answer(question, response.answers.get(key)) for key, question in keyed_questions), - usage=DecisionsIRUsage( + usage=None + if usage is None + else DecisionsIRUsage( input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, cached_tokens=usage.cached_tokens, @@ -485,7 +487,9 @@ def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIR **response.extra, "model": response.model, "answers": {key: answer for key, answer in keyed_answers if answer is not None}, - "usage": DecisionsUsage.model_validate( + "usage": None + if response.usage is None + else DecisionsUsage.model_validate( { **response.usage.extra, "input_tokens": response.usage.input_tokens, diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 0a76588f767..2bbfd2010cc 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -58,7 +58,7 @@ class CloudflareChatConfig(OpenAIGPTConfig): trimmed: Final = api_base.rstrip("/") if trimmed.endswith("/ai/run"): verbose_logger.warning( - "Cloudflare api_base ending in '/ai/run' is the legacy Workers AI path and no longer serves OpenAI-compatible requests; rewriting to the '/ai/v1' endpoint" + "Cloudflare chat completions require the OpenAI-compatible '/ai/v1' endpoint; rewriting api_base ending in '/ai/run' to '/ai/v1'. Decision models continue to use the native Workers AI '/ai/run/{model}' endpoint." ) return join_api_base(trimmed[: -len("/ai/run")], "ai/v1") return api_base diff --git a/litellm/llms/cloudflare/decisions/transformation.py b/litellm/llms/cloudflare/decisions/transformation.py index 7c776367886..7575e8bce21 100644 --- a/litellm/llms/cloudflare/decisions/transformation.py +++ b/litellm/llms/cloudflare/decisions/transformation.py @@ -20,6 +20,8 @@ class CloudflareDecisionsConfig(BaseDecisionsConfig): inlines_remote_images = True def get_default_api_base(self) -> str | None: + # Decision models (Clef) require native /ai/run/{model} requests with + # state and questions; /ai/v1/chat/completions only accepts chat requests. return workers_ai_default_api_base() def missing_api_base_message(self, custom_llm_provider: str) -> str: diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index a9a20101bb1..ecaa0a2ef16 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -17,6 +17,8 @@ from litellm.utils import supports_reasoning, supports_vision from ...openai.chat.gpt_transformation import OpenAIGPTConfig +SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "xhigh", "max", "ultra") + class DeepSeekChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list: @@ -33,7 +35,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): optional_params: dict, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """ Map OpenAI params to DeepSeek params. @@ -60,6 +62,13 @@ class DeepSeekChatConfig(OpenAIGPTConfig): elif reasoning_effort is not None: optional_params["thinking"] = {"type": "disabled" if reasoning_effort == "none" else "enabled"} + if ( + isinstance(reasoning_effort, str) + and reasoning_effort in SUPPORTED_REASONING_EFFORTS + and optional_params["thinking"] == {"type": "enabled"} + ): + return {**optional_params, "reasoning_effort": reasoning_effort} + return optional_params def _fill_reasoning_content(self, messages: list[AllMessageValues]) -> list[AllMessageValues]: diff --git a/litellm/llms/fireworks_ai/decisions/transformation.py b/litellm/llms/fireworks_ai/decisions/transformation.py index 87d402902be..bba4835060b 100644 --- a/litellm/llms/fireworks_ai/decisions/transformation.py +++ b/litellm/llms/fireworks_ai/decisions/transformation.py @@ -14,6 +14,8 @@ class FireworksAIDecisionsConfig(BaseDecisionsConfig): def parse_response(self, payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: response: Final = parse_systemone_response(self.unwrap_response(payload), request) + if response.usage is None: + return response cached: Final = response.usage.extra.get("cached_input_tokens") if isinstance(cached, bool) or not isinstance(cached, int): return response diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 349a4b4ab87..23f68335ea9 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -11,7 +11,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields +from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_internal_message_fields from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, ) @@ -154,7 +154,7 @@ class GroqChatConfig(OpenAILikeChatConfig): for k, v in _message.items(): if v is not None: new_message[k] = v - messages[idx] = strip_litellm_response_only_fields(new_message) + messages[idx] = strip_litellm_internal_message_fields(new_message) if is_async: return super()._transform_messages(messages=messages, model=model, is_async=True) diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 56f348881de..25d403a0dff 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -311,6 +311,7 @@ def _openai_answer(question: DecisionsIRQuestion, answer: DecisionsIRAnswer) -> def ir_to_openai_response( response: DecisionsIRResponse, request: DecisionsIRRequest, requested_model: str ) -> OpenAIDecisionResponse: + usage: Final = response.usage or DecisionsIRUsage() return OpenAIDecisionResponse( model=response.model or requested_model, answers=tuple( @@ -318,14 +319,14 @@ def ir_to_openai_response( for question, answer in zip(request.questions, response.answers, strict=True) ), usage=OpenAIDecisionUsage( - input_tokens=response.usage.input_tokens, + input_tokens=usage.input_tokens, input_tokens_details=OpenAIDecisionInputTokensDetails( - cached_tokens=response.usage.cached_tokens, - cache_write_tokens=response.usage.cache_write_tokens, + cached_tokens=usage.cached_tokens, + cache_write_tokens=usage.cache_write_tokens, ), - output_tokens=response.usage.output_tokens, - output_tokens_details=OpenAIDecisionOutputTokensDetails(reasoning_tokens=response.usage.reasoning_tokens), - total_tokens=response.usage.input_tokens + response.usage.output_tokens, + output_tokens=usage.output_tokens, + output_tokens_details=OpenAIDecisionOutputTokensDetails(reasoning_tokens=usage.reasoning_tokens), + total_tokens=usage.input_tokens + usage.output_tokens, ), ) diff --git a/litellm/llms/opencode/chat/transformation.py b/litellm/llms/opencode/chat/transformation.py index 40b99dbdd76..dcb30b5835c 100644 --- a/litellm/llms/opencode/chat/transformation.py +++ b/litellm/llms/opencode/chat/transformation.py @@ -1,6 +1,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Final +import httpx from pydantic import TypeAdapter import litellm @@ -9,7 +10,14 @@ from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig -from litellm.llms.opencode.common_utils import go_api_base, go_api_key, go_endpoint, go_headers, go_model_info +from litellm.llms.opencode.common_utils import ( + attach_go_reported_cost, + go_api_base, + go_api_key, + go_endpoint, + go_headers, + go_model_info, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -18,7 +26,46 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging -class OpenCodeGoChatConfig(OpenAILikeChatConfig): +class _GoReportedCostMixin: + """Record usage.opencode.list_cost so response_cost reflects OpenCode's price, not the static rows.""" + + go_usage_key: str = "usage" + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "Logging", + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + response: Final = super().transform_response( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # mixin precedes a BaseConfig subclass + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=dict(request_data), + messages=list(messages), + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + priced: Final = TypeAdapter(ModelResponse).validate_python(response) + attach_go_reported_cost( + priced, TypeAdapter(dict[str, object]).validate_json(raw_response.content).get(self.go_usage_key) + ) + return priced + + +class OpenCodeGoChatConfig(_GoReportedCostMixin, OpenAILikeChatConfig): @property def custom_llm_provider(self) -> str: return "opencode_go" @@ -60,7 +107,7 @@ class OpenCodeGoChatConfig(OpenAILikeChatConfig): return [f"opencode_go/{entry['id']}" for entry in models if isinstance(entry.get("id"), str)] -class OpenCodeGoMessagesChatConfig(AnthropicConfig): +class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): @property def custom_llm_provider(self) -> str: return "opencode_go" @@ -110,7 +157,9 @@ class OpenCodeGoMessagesChatConfig(AnthropicConfig): return TypeAdapter(dict[str, object]).validate_python(raw) -class OpenCodeGoGeminiChatConfig(GoogleAIStudioGeminiConfig): +class OpenCodeGoGeminiChatConfig(_GoReportedCostMixin, GoogleAIStudioGeminiConfig): + go_usage_key: str = "usageMetadata" + @property def supports_stream_param_in_request_body(self) -> bool: return False diff --git a/litellm/llms/opencode/common_utils.py b/litellm/llms/opencode/common_utils.py index 46e382d8e65..82cd957d23c 100644 --- a/litellm/llms/opencode/common_utils.py +++ b/litellm/llms/opencode/common_utils.py @@ -4,8 +4,11 @@ from typing import Final from pydantic import TypeAdapter +from litellm.litellm_core_utils.core_helpers import provider_reported_usage_cost, set_response_cost_in_hidden_params from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse def is_opencode_api_base(api_base: str | None) -> bool: @@ -67,3 +70,11 @@ def go_headers( params, "x-opencode-session", ) + + +def go_reported_cost(usage: object) -> float | None: + return provider_reported_usage_cost(usage, "opencode_go") + + +def attach_go_reported_cost(response: ModelResponse | ResponsesAPIResponse, usage: object) -> None: + set_response_cost_in_hidden_params(response, go_reported_cost(usage)) diff --git a/litellm/llms/opencode/responses/transformation.py b/litellm/llms/opencode/responses/transformation.py index ca75a319c7e..615ca60b45b 100644 --- a/litellm/llms/opencode/responses/transformation.py +++ b/litellm/llms/opencode/responses/transformation.py @@ -1,11 +1,24 @@ from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +import httpx from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig -from litellm.llms.opencode.common_utils import go_api_base, go_headers +from litellm.llms.opencode.common_utils import attach_go_reported_cost, go_api_base, go_headers +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + class OpenCodeGoResponsesConfig(OpenAIResponsesAPIConfig): @property @@ -25,3 +38,28 @@ class OpenCodeGoResponsesConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams | None, ) -> dict[str, object]: # mutable-ok: protocol handlers mutate the returned request and headers return go_headers(headers, litellm_params, litellm_params.api_key if litellm_params else None, model) + + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "Logging", + ) -> ResponsesAPIResponse: + response: Final = super().transform_response_api_response( + model=model, raw_response=raw_response, logging_obj=logging_obj + ) + attach_go_reported_cost(response, response.usage) + return response + + def transform_streaming_response( + self, + model: str, + parsed_chunk: Mapping[str, object], + logging_obj: "Logging", + ) -> ResponsesAPIStreamingResponse: + event: Final = super().transform_streaming_response( # pyright: ignore[reportUnknownMemberType] # base takes an untyped dict + model=model, parsed_chunk=dict(parsed_chunk), logging_obj=logging_obj + ) + if isinstance(event, ResponseCompletedEvent | ResponseFailedEvent | ResponseIncompleteEvent): + attach_go_reported_cost(event.response, event.response.usage) + return event diff --git a/litellm/llms/openrouter/decisions/transformation.py b/litellm/llms/openrouter/decisions/transformation.py index 0fa5c7ca7cb..30d33f09efd 100644 --- a/litellm/llms/openrouter/decisions/transformation.py +++ b/litellm/llms/openrouter/decisions/transformation.py @@ -13,7 +13,7 @@ class OpenRouterDecisionsConfig(BaseDecisionsConfig): return "https://openrouter.ai/api" def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: - cost: Final = response.usage.extra.get("cost") + cost: Final = response.usage.extra.get("cost") if response.usage is not None else None if not isinstance(cost, (int, float)): return None return float(cost) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c7dfce0712f..9a7d88b8eda 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -644,9 +644,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, @@ -4339,7 +4343,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -9232,7 +9236,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -9679,7 +9683,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5.5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -15659,6 +15663,28 @@ "/v1/images/generations" ] }, + "cerebras/llama-3.1-70b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, + "cerebras/llama-3.1-8b": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, "cerebras/llama-3.3-70b": { "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", @@ -33212,6 +33238,7 @@ "github_copilot/claude-opus-4.8-fast": { "cache_creation_input_token_cost": 1.25e-05, "cache_read_input_token_cost": 1e-06, + "deprecation_date": "2026-11-16", "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, @@ -37401,6 +37428,7 @@ }, "gpt-5.6-cyber": { "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_12hr": 3.125e-05, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, @@ -42698,6 +42726,29 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5e-07, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://docs.mistral.ai/models/zai-glm-5-3-flash", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/zai-glm-5": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -59514,6 +59565,7 @@ "supports_audio_input": true, "supports_audio_output": true, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -59709,6 +59761,7 @@ "tpm": 250000, "rpm": 10, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -75362,10 +75415,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-5-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 13700c6da35..81ff25ee31b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -219,6 +219,8 @@ from litellm.types.proxy.auth.user_roles import ( # noqa: E402 # re-export aft KeyManagementRoutes as KeyManagementRoutes, # noqa: PLC0414 # public re-export ) +MODEL_LIST_ROUTES: Final = ("/models", "/v1/models", "/claude_code_gateway/v1/models") + class LiteLLMRoutes(enum.Enum): openai_route_names = [ @@ -307,8 +309,7 @@ class LiteLLMRoutes(enum.Enum): "/threads/{thread_id}/runs", "/v1/threads/{thread_id}/runs", # models - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, # token counter "/utils/token_counter", "/utils/model_info", @@ -577,8 +578,7 @@ class LiteLLMRoutes(enum.Enum): "/health/services", "/key/list", "/user/filter/ui", - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, "/sso/get/ui_settings", "/get/user_banner", "/get/latest_release_info", @@ -762,8 +762,7 @@ class LiteLLMRoutes(enum.Enum): "/model/info", "/v2/model/info", "/v2/key/info", - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, "/global/spend", "/global/spend/logs", "/global/spend/keys", diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 4c4155379cd..ba3307bc1ad 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -58,6 +58,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg ) from litellm.proxy.management_endpoints.sso_helper_utils import CLI_SSO_SESSIONS_TARGET from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail +from litellm.proxy.proxy_server import model_list from litellm.types.llms.base import LiteLLMBaseModel GATEWAY_PREFIX: Final = "/claude_code_gateway" @@ -158,6 +159,13 @@ router: Final = APIRouter( _GATEWAY_ENABLED: Final = (Depends(ensure_gateway_enabled),) _AUTHENTICATED: Final = (Depends(user_api_key_auth),) +router.add_api_route( + "/v1/models", + model_list, + methods=["GET"], + dependencies=_GATEWAY_ENABLED, + include_in_schema=False, +) router.add_api_route( "/v1/messages", anthropic_response, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 2c6768b8423..79e385c376a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,6 +15,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from functools import partial from itertools import chain from types import MappingProxyType @@ -47,6 +48,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( + MODEL_LIST_ROUTES, RBAC_ROLES, UI_TEAM_ID, CallInfo, @@ -910,8 +912,7 @@ def _global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_che and not skip_budget_checks and global_proxy_spend is not None and RouteChecks.is_llm_api_route(route=route) - and route != "/v1/models" - and route != "/models" + and route not in MODEL_LIST_ROUTES ): if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget: raise litellm.BudgetExceededError( @@ -1028,8 +1029,7 @@ async def check_tools_allowlist( # budget cannot reach side-effectful routes like /health/services (Slack/email/webhook). See #27923. MODEL_DISCOVERY_ROUTES: Final = frozenset( { - "/v1/models", - "/models", + *MODEL_LIST_ROUTES, "/model/info", "/v1/model/info", "/v2/model/info", @@ -1800,6 +1800,20 @@ async def get_team_member_default_budget( return budget +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + if member_budget is not None and member_budget.max_budget is not None: + return member_budget.effective_max_budget(now) + if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: + return None + return team_default_budget.max_budget + ( + member_budget.active_temp_budget_increase(now) if member_budget is not None else 0.0 + ) + + async def resolve_default_end_user_budget( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, @@ -7035,29 +7049,24 @@ async def _check_team_member_budget( # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. - team_member_budget: float | None = None member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None now: Final = get_utc_datetime() - if member_budget_row is not None and member_budget_row.max_budget is not None: - team_member_budget = member_budget_row.effective_max_budget(now=now) - else: - default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") - if isinstance(default_budget_id, str): - default_budget: Final = await get_team_member_default_budget( - budget_id=default_budget_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - # Treat 0 on the team default as "no cap". - # Per-member rows still respect 0 as an explicit admin disable. - if ( - default_budget is not None - and default_budget.max_budget is not None - and default_budget.max_budget > 0 - ): - team_member_budget = default_budget.max_budget + ( - member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 - ) + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") + team_default_budget: Final = ( + await get_team_member_default_budget( + budget_id=default_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if (member_budget_row is None or member_budget_row.max_budget is None) + and isinstance(default_budget_id, str) + else None + ) + team_member_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget_row, + team_default_budget=team_default_budget, + now=now, + ) if team_member_budget is not None: team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 3ae671d40ac..c402f111563 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -7,6 +7,7 @@ from fastapi import HTTPException, Request, status from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( + MODEL_LIST_ROUTES, UI_TEAM_ID, CommonProxyErrors, KeyManagementRoutes, @@ -93,8 +94,7 @@ _SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( _UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( { - "/models", - "/v1/models", + *MODEL_LIST_ROUTES, "/utils/token_counter", "/utils/transform_request", "/responses/input_tokens", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 23faa964d5f..cec176ac559 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4821,10 +4821,15 @@ class ProxyBaseLLMRequestProcessing: model_info = metadata.get("model_info") or {} model_id = model_info.get("id", None) - # 3. Final fallback to self.data["litellm_metadata"] (for routes like /v1/responses that populate data before error) + # 3. Final fallback to the request metadata the router updates in place with the selected deployment if not model_id: - litellm_metadata: Final = self.data.get("litellm_metadata", {}) or {} - model_info = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", None) + model_id = self._metadata_model_id("litellm_metadata") or self._metadata_model_id("metadata") return model_id + + def _metadata_model_id(self, metadata_key: str) -> str | None: + match self.data.get(metadata_key): + case {"model_info": {"id": str() as request_model_id}}: + return request_model_id + case _: + return None diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index ccf2232e379..59fd00cec43 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -65,6 +65,7 @@ from litellm.repositories.table_repositories import ( from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, + SpendResetWrites, budget_cascade_unit_of_work, enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, @@ -235,7 +236,7 @@ def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: def _queue_budget_linked_resets( - writes: LinkedSpendResetWrites, + writes: SpendResetWrites, cascade: "_BudgetCascade", extra: Mapping[str, object] = MappingProxyType({}), ) -> None: @@ -908,6 +909,7 @@ class ResetBudgetJob: db_span("reset_budget_cascade", "LiteLLM_BudgetTable"), budget_cascade_unit_of_work(self._new_batch) as uow, ): + uow.team_memberships.queue_total_spend_repair(cascade.budget_ids) _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 22540866a77..7870ee15d2f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -472,6 +472,8 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS # recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma # types a raw array parameter from the first batch a connection sees, so after an all-$0 batch # (integers) every later fractional batch on that connection failed with "improper binary format". +# total_spend is floored at spend before adding: rows written before total_spend existed (migration +# 20260421135425, default 0) carry their history only in spend, so a plain increment keeps them behind forever. _TEAM_MEMBER_SPEND_SQL: Final = """ INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) SELECT member.user_id, member.team_id, member.cost, member.cost @@ -487,7 +489,7 @@ WHERE EXISTS ( ) ON CONFLICT (user_id, team_id) DO UPDATE SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, - total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend + total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend) + EXCLUDED.total_spend """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py new file mode 100644 index 00000000000..a15a372c462 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py @@ -0,0 +1,72 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, TypeAdapter + +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoOptionalParams + +from .reco import RecoGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +_CONFIG_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _optional_params_mapping(optional_params: object) -> Mapping[str, object]: + """optional_params arrives as a pydantic object from YAML, or as a plain dict when reconciled from the DB.""" + if optional_params is None: + return MappingProxyType({}) + raw: Final = ( + optional_params.model_dump(exclude_unset=True) if isinstance(optional_params, BaseModel) else optional_params + ) + return _CONFIG_MAPPING.validate_python(raw) + + +def _first_configured(field_name: str, sources: tuple[Mapping[str, object], ...]) -> object: + value: Final = next((source[field_name] for source in sources if source.get(field_name) is not None), None) + if isinstance(value, str) and value.startswith("os.environ/"): + return get_secret_str(value) + return value + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RecoGuardrail: + import litellm + + # Top-level fields only count when set explicitly: their defaults are fail-closed and would flip Reco's. + explicit_top_level: Final = _CONFIG_MAPPING.validate_python( + litellm_params.model_dump(include=litellm_params.model_fields_set & RecoOptionalParams.model_fields.keys()) + ) + sources: Final = (_optional_params_mapping(litellm_params.optional_params), explicit_top_level) + configured: Final = ( + (field_name, _first_configured(field_name, sources)) for field_name in RecoOptionalParams.model_fields + ) + reco_params: Final = RecoOptionalParams.model_validate( + {field_name: value for field_name, value in configured if value is not None} + ) + + _reco_callback: Final = RecoGuardrail( + guardrail_name=guardrail["guardrail_name"], + reco_tenant_id=reco_params.reco_tenant_id, + api_base=reco_params.api_base, + unreachable_fallback=reco_params.unreachable_fallback, + fail_on_error=reco_params.fail_on_error, + event_hook=litellm_params.mode, + default_on=bool(litellm_params.default_on), + ) + litellm.logging_callback_manager.add_litellm_callback(_reco_callback) + + return _reco_callback + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.RECO.value: initialize_guardrail, +} + + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.RECO.value: RecoGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/reco/reco.py b/litellm/proxy/guardrails/guardrail_hooks/reco/reco.py new file mode 100644 index 00000000000..c8998a20367 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/reco/reco.py @@ -0,0 +1,84 @@ +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, Literal, Optional + +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped +) +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( + GenericGuardrailAPI, +) +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.proxy.guardrails.guardrail_hooks.reco import validate_reco_tenant_id +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoOptionalParams + + +class RecoGuardrail(GenericGuardrailAPI): + """Reco guardrail integration for LiteLLM, built on the Generic Guardrail API wire contract.""" + + def __init__( + self, + reco_tenant_id: str | None, + api_base: str | None, + headers: Mapping[str, str] | None = None, + unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", + fail_on_error: bool = True, + guardrail_name: str | None = None, + event_hook: str | Sequence[str] | Mode | None = None, + default_on: bool = False, + ) -> None: + if not reco_tenant_id: + raise ValueError("reco_tenant_id is required for the Reco guardrail") + if not api_base: + raise ValueError("api_base is required for the Reco guardrail") + + validated_tenant_id: Final = validate_reco_tenant_id(reco_tenant_id) + + super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped extra_headers and **kwargs + headers={**(headers or {}), "X-Reco-Tenant-Id": validated_tenant_id}, + api_base=api_base, + unreachable_fallback=unreachable_fallback, + fail_on_error=fail_on_error, + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: must match GenericGuardrailAPI.apply_guardrail's own param type + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + # GenericGuardrailAPI raises with a hardcoded module-level guardrail_name on BLOCKED, + # not the guardrail's own configured name. Re-raise with the right one instead of + # patching the shared base class. + try: + return await super().apply_guardrail( # pyright: ignore[reportUnknownMemberType] # parent annotates request_data as a bare dict + inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj + ) + except GuardrailRaisedException as e: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=e.message, + should_wrap_with_default_message=False, + status_code=e.status_code, + blocked_content=e.blocked_content, + ) from e + + @staticmethod + def get_config_model() -> type["GuardrailConfigModel[RecoOptionalParams]"] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.reco import RecoConfigModel + + return RecoConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: fixed by base class + return [GuardrailEventHooks.pre_call] diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6801c8521e5..c11b59a809a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -38,6 +38,7 @@ from litellm.proxy.db.db_spend_update_writer import ( ) from litellm.proxy.hooks.search_free_tier import apply_search_free_tier from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route +from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( ObjectMapping, @@ -60,7 +61,12 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, should_store_prompts_and_responses_in_spend_logs, ) -from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend +from litellm.proxy.utils import ( + REJECTED_BEFORE_REQUEST_SETUP, + REJECTED_BEFORE_REQUEST_SETUP_KEY, + BilledFailedAttempt, + ProxyUpdateSpend, +) from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, @@ -396,6 +402,8 @@ class _ProxyDBLogger(CustomLogger): request_route: Final = user_api_key_dict.request_route if ( ProxyDBLogger._should_track_errors_in_db() is False + or ProxyDBLogger._skips_pre_routing_rejections() + and is_pre_routing_rejection(original_exception, _REQUEST_DATA.validate_python(request_data)) or request_route is not None and not ( RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) @@ -862,16 +870,40 @@ class _ProxyDBLogger(CustomLogger): If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings """ - from litellm.proxy.proxy_server import general_settings - - if general_settings.get("disable_error_logs") is True: + if _general_setting_enabled("disable_error_logs"): return False return + @staticmethod + def _skips_pre_routing_rejections() -> bool: + """``general_settings.disable_pre_routing_error_logs`` keeps requests the proxy rejected before + selecting a deployment out of the spend logs, so a client retrying a bad key or an exhausted + budget cannot flood ``LiteLLM_SpendLogs`` with zero-spend rows.""" + return _general_setting_enabled("disable_pre_routing_error_logs") + + +def _general_setting_enabled(name: str) -> bool: + from litellm.proxy.proxy_server import general_settings + + return name in general_settings and general_settings[name] is True + ProxyDBLogger: Final = _ProxyDBLogger +_REQUEST_DATA: Final = TypeAdapter(dict[str, object]) + + +def is_pre_routing_rejection(exception: Exception, request_data: Mapping[str, object]) -> bool: + """The pipeline sets up the logging object only after auth, body parsing and model resolution succeed, and + before guardrails or routing run. ``ProxyLogging._handle_logging_proxy_only_error`` marks a failure that + arrives without one with the ``REJECTED_BEFORE_REQUEST_SETUP`` sentinel, so that request never had a provider + considered. An unknown model is the one later rejection that still never reached a provider.""" + return request_data.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP or isinstance( + exception, ProxyModelNotFoundError + ) + + def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None} if not patch: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 361778ab2ea..1fa18e9b2b8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -52,6 +52,7 @@ from litellm.proxy._types import ( CommonProxyErrors, DeleteTeamRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, LiteLLM_DeletedTeamTable, Litellm_EntityType, LiteLLM_ManagementEndpoint_MetadataFields, @@ -114,6 +115,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo get_user_object, invalidate_team_member_spend_state, jwt_key_mapping_cache_key, + resolve_effective_team_member_budget, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, @@ -244,6 +246,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, TeamIdSearchMatch, + TeamListCallerMembership, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -254,6 +257,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( UpdateTeamMemberPermissionsRequest, ) from litellm.types.utils import BudgetConfig, GenericBudgetConfigType +from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import Prisma @@ -302,6 +306,34 @@ class _TeamIdRow(Protocol): def team_id(self) -> str: ... +class _TeamListBudgetRow(Protocol): + @property + def budget_id(self) -> str: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + + @property + def max_budget(self) -> float | None: ... + + @property + def temp_budget_expiry(self) -> datetime | None: ... + + @property + def temp_budget_increase(self) -> float | None: ... + + +class _TeamListMembershipRow(Protocol): + @property + def team_id(self) -> str: ... + + @property + def spend(self) -> float | None: ... + + @property + def litellm_budget_table(self) -> _TeamListBudgetRow | None: ... + + class _CacheableTeamRow(_TeamIdRow, _ModelDumpRow, Protocol): ... @@ -981,7 +1013,7 @@ async def get_all_team_memberships( returned_tm: Final[list[LiteLLM_TeamMembership]] = [] for tm in team_memberships: - returned_tm.append(LiteLLM_TeamMembership.model_validate(tm.model_dump())) + returned_tm.append(with_total_spend_floor(LiteLLM_TeamMembership.model_validate(tm.model_dump()))) return returned_tm @@ -4170,6 +4202,23 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.") +# One statement so the validated spend cannot go stale between the read and the write: the WHERE re-checks +# reset_to <= spend, and total_spend is floored at spend before spend drops so the history survives the reset. +_RESET_TEAM_MEMBER_SPEND_SQL: Final = """ +UPDATE "LiteLLM_TeamMembership" +SET total_spend = GREATEST(total_spend, spend), spend = $3::float8 +WHERE user_id = $1 AND team_id = $2 AND spend >= $3::float8 +""" + + +def with_total_spend_floor(membership: LiteLLM_TeamMembership) -> LiteLLM_TeamMembership: + """Rows written before total_spend existed (migration 20260421135425) hold their history only in spend.""" + spend: Final = membership.spend or 0.0 + if (membership.total_spend or 0.0) >= spend: + return membership + return membership.model_copy(update={"total_spend": spend}) + + def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn: detail: Final = {"error": message} raise HTTPException(status_code=status_code, detail=detail) @@ -4254,10 +4303,12 @@ async def reset_team_member_spend_fn( current_spend: Final = membership.spend or 0.0 reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership) - await _team_membership_db(prisma_client).update( - where=membership_where, - data={"spend": reset_to}, - ) + reset_rows: Final = await _execute_raw(prisma_client)(_RESET_TEAM_MEMBER_SPEND_SQL, user_id, team_id, reset_to) + if reset_rows == 0: + _raise_reset_spend_error( + status.HTTP_409_CONFLICT, + f"Member spend changed since it was read; reset_to ({reset_to}) must be <= current spend. Retry.", + ) await invalidate_team_member_spend_state( user_id=user_id, @@ -5467,16 +5518,15 @@ async def team_member_me( user_api_key_cache=user_api_key_cache, ) - caller_user_email: Final = user_api_key_dict.user_email - member_role: str | None = None - for m in team_table.members_with_roles: - # Match by user_id when present, else fall back to email — members - # added by email may have user_id=None on the stored entry. - if (m.user_id is not None and m.user_id == caller_user_id) or ( - m.user_email is not None and caller_user_email is not None and m.user_email == caller_user_email - ): - member_role = m.role - break + caller_identity: Final = _TeamListCallerIdentity( + user_id=caller_user_id, + user_email=user_api_key_dict.user_email, + ) + matched_member: Final = next( + (member for member in team_table.members_with_roles if _team_member_matches_caller(member, caller_identity)), + None, + ) + member_role: Final = matched_member.role if matched_member is not None else None if member_role is None: # Caller is not a member of this team. Even proxy admins get 404 here — @@ -5517,16 +5567,17 @@ async def team_member_me( litellm_budget_table=None, ) + floored: Final = with_total_spend_floor(membership) return TeamMemberInfoResponse( user_id=caller_user_id, team_id=team_id, team_alias=team_table.team_alias, role=member_role, user_email=user_email, - spend=membership.spend, - total_spend=membership.total_spend, - budget_id=membership.budget_id, - litellm_budget_table=membership.litellm_budget_table, + spend=floored.spend, + total_spend=floored.total_spend, + budget_id=floored.budget_id, + litellm_budget_table=floored.litellm_budget_table, ) @@ -5886,11 +5937,16 @@ def _convert_teams_to_response_models( members_count = len(members_with_roles) keys_count = counts.get(team_dict.get("team_id") or "", 0) team_list.append( - TeamListItem( - **team_dict, - members_count=members_count, - keys_count=keys_count, - access_group_skill_names=None, + TeamListItem.model_validate( + cast( # cast-ok: TeamListItem validates the Prisma model dump + object, + { + **team_dict, + "members_count": members_count, + "keys_count": keys_count, + "access_group_skill_names": None, + }, + ) ) ) return team_list @@ -5921,6 +5977,137 @@ async def _get_keys_count_by_team( return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")} +def _team_list_budget_model(budget_row: _TeamListBudgetRow | None) -> LiteLLM_BudgetTable | None: + if budget_row is None: + return None + return LiteLLM_BudgetTable( + max_budget=budget_row.max_budget, + temp_budget_expiry=budget_row.temp_budget_expiry, + temp_budget_increase=budget_row.temp_budget_increase, + ) + + +@dataclass(frozen=True, slots=True) +class _TeamListCallerIdentity: + user_id: str | None + user_email: str | None + + +def _team_member_matches_caller(member: Member, caller_identity: _TeamListCallerIdentity) -> bool: + return (member.user_id is not None and member.user_id == caller_identity.user_id) or ( + member.user_email is not None + and caller_identity.user_email is not None + and member.user_email == caller_identity.user_email + ) + + +def _caller_in_roster(team: TeamListItem, caller_identity: _TeamListCallerIdentity) -> bool: + return any(_team_member_matches_caller(member, caller_identity) for member in team.members_with_roles) + + +def _team_member_default_budget_id( + team: TeamListItem, + caller_identity: _TeamListCallerIdentity, + memberships_by_team: Mapping[str, _TeamListMembershipRow], +) -> str | None: + membership_row: Final = memberships_by_team.get(team.team_id) + if membership_row is None and not _caller_in_roster(team, caller_identity): + return None + member_budget: Final = _team_list_budget_model( + membership_row.litellm_budget_table if membership_row is not None else None + ) + if member_budget is not None and member_budget.max_budget is not None: + return None + metadata: Final = cast( # cast-ok: team metadata is stored as a JSON object + "dict[str, object] | None", team.metadata + ) + if metadata is None: + return None + budget_id: Final = metadata.get("team_member_budget_id") + return budget_id if isinstance(budget_id, str) else None + + +def _team_list_caller_membership( + team: TeamListItem, + caller_identity: _TeamListCallerIdentity, + memberships_by_team: Mapping[str, _TeamListMembershipRow], + default_budget_id: str | None, + budgets_by_id: Mapping[str, _TeamListBudgetRow], + now: datetime, +) -> TeamListCallerMembership | None: + membership_row: Final = memberships_by_team.get(team.team_id) + if membership_row is None and not _caller_in_roster(team, caller_identity): + return None + member_budget_row: Final = membership_row.litellm_budget_table if membership_row is not None else None + member_budget: Final = _team_list_budget_model(member_budget_row) + default_budget_row: Final = budgets_by_id.get(default_budget_id) if default_budget_id is not None else None + team_default_budget: Final = _team_list_budget_model(default_budget_row) + max_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget, + team_default_budget=team_default_budget, + now=now, + ) + budget_reset_at: Final = ( + member_budget_row.budget_reset_at if member_budget_row is not None and max_budget is not None else None + ) + return TeamListCallerMembership( + spend=(membership_row.spend if membership_row is not None else 0.0) or 0.0, + max_budget=max_budget, + budget_reset_at=budget_reset_at, + ) + + +async def _enrich_team_list_with_caller_membership( + prisma_client: PrismaClient, + teams: Sequence[TeamListItem], + caller_identity: _TeamListCallerIdentity, + now: datetime, +) -> tuple[TeamListItem, ...]: + caller_user_id: Final = caller_identity.user_id + if caller_user_id is None: + return tuple(team.model_copy(update={"caller_membership": None}) for team in teams) + if not teams: + return () + + page_team_ids: Final = tuple(team.team_id for team in teams) + membership_rows: Final[Sequence[_TeamListMembershipRow]] = await _team_membership_db(prisma_client).find_many( + # bounded-ok: page_team_ids is capped at 100 teams by list_team_v2 + where={"team_id": {"in": page_team_ids}, "user_id": caller_user_id}, + include={"litellm_budget_table": True}, + ) + memberships_by_team: Final = {row.team_id: row for row in membership_rows} + default_budget_id_by_team: Final = MappingProxyType( + {team.team_id: _team_member_default_budget_id(team, caller_identity, memberships_by_team) for team in teams} + ) + default_budget_ids: Final = tuple( + dict.fromkeys(budget_id for budget_id in default_budget_id_by_team.values() if budget_id is not None) + ) + default_budget_rows: Final[Sequence[_TeamListBudgetRow]] = ( + await _budget_db(prisma_client).find_many( + # bounded-ok: default_budget_ids is derived from this page of at most 100 teams + where={"budget_id": {"in": default_budget_ids}}, + ) + if default_budget_ids + else () + ) + budgets_by_id: Final = {row.budget_id: row for row in default_budget_rows} + return tuple( + team.model_copy( + update={ + "caller_membership": _team_list_caller_membership( + team=team, + caller_identity=caller_identity, + memberships_by_team=memberships_by_team, + default_budget_id=default_budget_id_by_team.get(team.team_id), + budgets_by_id=budgets_by_id, + now=now, + ) + } + ) + for team in teams + ) + + async def _enforce_list_team_v2_access( user_api_key_dict: UserAPIKeyAuth, user_id: str | None, @@ -6034,7 +6221,7 @@ async def list_team_v2( sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), status: str | None = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> TeamListResponse: """ Get a paginated list of teams with filtering and sorting options. @@ -6109,13 +6296,13 @@ async def list_team_v2( ) if where_conditions is None: - return { - "teams": [], - "total": 0, - "page": page, - "page_size": page_size, - "total_pages": 0, - } + return TeamListResponse( + teams=[], + total=0, + page=page, + page_size=page_size, + total_pages=0, + ) from litellm.proxy.management_helpers.team_list_sorting import ( COMPUTED_TEAM_SORTS, @@ -6227,13 +6414,27 @@ async def list_team_v2( team_groups, tuple(ag_id for ag_id in (team_item.access_group_ids or []) if ag_id not in ag_lookup) ) - return { - "teams": team_list, - "total": total_count, - "page": page, - "page_size": page_size, - "total_pages": total_pages, - } + caller_identity: Final = _TeamListCallerIdentity( + user_id=user_api_key_dict.user_id, + user_email=user_api_key_dict.user_email, + ) + team_list_with_caller_membership: Final = ( + await _enrich_team_list_with_caller_membership( + prisma_client=prisma_client, + teams=tuple(team for team in team_list if isinstance(team, TeamListItem)), + caller_identity=caller_identity, + now=get_utc_datetime(), + ) + if not use_deleted_table and team_list + else team_list + ) + return TeamListResponse( + teams=list(team_list_with_caller_membership), + total=total_count, + page=page, + page_size=page_size, + total_pages=total_pages, + ) async def _authorize_and_filter_teams( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 23ef89b3c5c..15f70073c7a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -1,15 +1,20 @@ import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE +from litellm.litellm_core_utils.core_helpers import ( + native_messages_reported_cost, + override_stamped_response_cost, + response_cost_in_hidden_params, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -53,6 +58,38 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +def _event_payload(event: str) -> Mapping[str, object] | None: + data_lines: Final = tuple( + line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + ) + if not data_lines or data_lines[0] in ("", "[DONE]"): + return None + try: + return _JSON_OBJECT.validate_json(data_lines[0]) + except ValidationError: + return None + + +def _message_delta_usages(all_chunks: Sequence[str | bytes]) -> Iterator[object]: + for chunk in reversed(all_chunks): + yield from ( + payload.get("usage") + for payload in map(_event_payload, reversed(split_sse_chunk_into_events(chunk))) + if payload is not None and payload.get("type") == "message_delta" + ) + + +def _stamp_reported_stream_cost( + response: ModelResponse, all_chunks: Sequence[str | bytes], logging_obj: LiteLLMLoggingObj +) -> None: + """The final ``message_delta`` is where a provider that prices its own responses reports the cost.""" + provider: Final = logging_obj.model_call_details.get("custom_llm_provider") + costs: Final = (native_messages_reported_cost(usage, provider) for usage in _message_delta_usages(all_chunks)) + reported: Final = next((cost for cost in costs if cost is not None), None) + if reported is not None: + override_stamped_response_cost(response, reported) + + class AnthropicPassthroughLoggingHandler: @staticmethod def anthropic_passthrough_handler( @@ -222,6 +259,7 @@ class AnthropicPassthroughLoggingHandler: ) if not isinstance(response, ModelResponse): return response + _stamp_reported_stream_cost(response, all_chunks, litellm_logging_obj) recovered_usage: Final = AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( response=response, all_chunks=all_chunks, model=model ) @@ -286,13 +324,21 @@ class AnthropicPassthroughLoggingHandler: if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/") else model ) + custom_pricing: Final = use_custom_pricing_for_model( + litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) + ) + reported_cost: Final = ( + response_cost_in_hidden_params(litellm_model_response) + if isinstance(litellm_model_response, ModelResponse) + else None + ) + if reported_cost is not None and not custom_pricing: + return reported_cost return litellm.completion_cost( completion_response=litellm_model_response, model=model_for_cost, custom_llm_provider=custom_llm_provider, - custom_pricing=use_custom_pricing_for_model( - litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) - ), + custom_pricing=custom_pricing, router_model_id=logging_obj.get_router_model_id(), ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8212e7cbe74..ce7e6afde24 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1476,6 +1476,13 @@ def _call_type_for_route(route: str | None) -> str | None: return call_types[0].value if len(operations) == 1 else None +REJECTED_BEFORE_REQUEST_SETUP_KEY: Final = "litellm_rejected_before_request_setup" +REJECTED_BEFORE_REQUEST_SETUP: Final = object() +"""Stored under ``REJECTED_BEFORE_REQUEST_SETUP_KEY`` on ``request_data`` when the failure handler finds no logging +object, meaning the request died (auth, body parsing, model resolution) before the pipeline set one up. A sentinel +object rather than ``True`` because ``request_data`` is the client body: JSON cannot forge this identity, so a +client can neither hide a guardrail block nor force a row by sending the key or a ``litellm_call_id`` itself.""" + _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data")) @@ -3888,6 +3895,7 @@ class ProxyLogging: **request_data, ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller + request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) if "metadata" not in request_data: request_data["metadata"] = {} request_data["metadata"].update(user_api_key_logged_metadata) diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index c42301a9316..63cfcbc7b45 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -171,4 +171,6 @@ class PrismaBatch(Protocol): @property def litellm_projecttable(self) -> BatchTable: ... + def execute_raw(self, query: str, *args: object) -> None: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 95f4f5e81be..f6aad883479 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -15,11 +15,11 @@ instead of full-model writes, which trip ``prisma.errors.DataError`` on rows carrying fields the update input type rejects (see #27730). """ -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime -from typing import Final +from typing import Final, Protocol, runtime_checkable from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch @@ -62,6 +62,13 @@ class TeamSpendResetWrites: ) +@runtime_checkable +class SpendResetWrites(Protocol): + def queue_spend_zero(self, where: Mapping[str, object]) -> None: ... + + def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: ... + + @dataclass(frozen=True, slots=True) class LinkedSpendResetWrites: table: BatchTable @@ -78,6 +85,30 @@ class LinkedSpendResetWrites: ) +# Membership rows written before total_spend existed (migration 20260421135425, default 0) hold their +# history only in spend. Zeroing spend would erase it, so the floor lands in the same batch, ahead of the reset. +_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL: Final = """ +UPDATE "LiteLLM_TeamMembership" +SET total_spend = spend +WHERE budget_id = ANY($1::text[]) AND spend > total_spend +""" + + +@dataclass(frozen=True, slots=True) +class TeamMembershipSpendResetWrites: + resets: LinkedSpendResetWrites + batch: PrismaBatch + + def queue_spend_zero(self, where: Mapping[str, object]) -> None: + self.resets.queue_spend_zero(where) + + def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: + self.resets.queue_spend_decrement(where, amount) + + def queue_total_spend_repair(self, budget_ids: Sequence[str]) -> None: + self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids)) + + @dataclass(frozen=True, slots=True) class BudgetWindowWrites: table: BatchTable @@ -106,7 +137,7 @@ class BudgetCascadeUnitOfWork: ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ - team_memberships: LinkedSpendResetWrites + team_memberships: TeamMembershipSpendResetWrites keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites @@ -132,7 +163,9 @@ async def budget_cascade_unit_of_work( ) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]: batch = new_batch() yield BudgetCascadeUnitOfWork( - team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership), + team_memberships=TeamMembershipSpendResetWrites( + resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch + ), keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken), organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 60318d6b184..8b00cdf4813 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -175,7 +175,7 @@ def responses_api_bridge_check( pinned_mode: Final = ( None if explicit_responses_prefix else deployment_routing_mode(deployment_model_info, custom_llm_provider) ) - if explicit_responses_prefix or pinned_mode == "responses": + if custom_llm_provider == LlmProviders.CHATGPT or explicit_responses_prefix or pinned_mode == "responses": model = model.replace("responses/", "") model_info["mode"] = "responses" elif pinned_mode == "chat" and backend_serves_chat_completions(model_info): diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 030911238ec..9e51e2533ea 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling @@ -289,6 +290,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] self.user_api_key_auth = user_api_key_auth self.original_request_params = original_request_params or {} + configured_rounds: Final = validated_max_agentic_loops( + self.original_request_params.get("max_agentic_loops"), field="litellm_params.max_agentic_loops" + ) + self.max_tool_call_rounds: int = MAX_MCP_TOOL_CALL_ROUNDS if configured_rounds is None else configured_rounds self.should_auto_execute = self._should_auto_execute_tools() # Streaming state management @@ -509,7 +514,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): 2. mcp_discovery - Emit MCP discovery events (after response.output_item.added) 3. continue_initial_response - Stream the current round's response (initial or follow-up). On completion, if auto-execute is on and the response contains tool calls, loops back - through tool_execution/follow-up instead of ending, up to MAX_MCP_TOOL_CALL_ROUNDS. + through tool_execution/follow-up instead of ending, up to self.max_tool_call_rounds. 4. tool_execution - Emit tool execution events 5. finished - End iteration """ @@ -1015,14 +1020,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Remove tool_choice to avoid forcing more tool calls follow_up_params.pop("tool_choice", None) - if self.tool_call_round >= MAX_MCP_TOOL_CALL_ROUNDS: + if self.tool_call_round >= self.max_tool_call_rounds: # Hit the round cap: drop tools entirely so the model must # answer in text instead of emitting another (unexecuted) # tool call that would otherwise end the stream in silence. follow_up_params.pop("tools", None) verbose_logger.warning( - "MCP auto-execute hit MAX_MCP_TOOL_CALL_ROUNDS=%s; forcing a text-only follow-up.", - MAX_MCP_TOOL_CALL_ROUNDS, + "MCP auto-execute hit max_agentic_loops=%s; forcing a text-only follow-up.", + self.max_tool_call_rounds, ) follow_up_response: Final = await aresponses(**follow_up_params) diff --git a/litellm/router.py b/litellm/router.py index 96102b678c3..7993b5b921c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14315,6 +14315,30 @@ class Router: return healthy_deployments return cast(_RouteDeploymentsT, filtered_deployments) # cast-ok: same list shape as the input candidates + @staticmethod + def _without_declared_text_only_deployments( + healthy_deployments: _RouteDeploymentsT, + group: str, + messages: list[dict[str, str]] | None, # mutable-ok: handed to _request_carries_image as-is + request_kwargs: Mapping[str, object] | None, + ) -> _RouteDeploymentsT: + if isinstance(healthy_deployments, Mapping) or not Router._request_carries_image( + messages, request_kwargs or {} + ): + return healthy_deployments + + vision_capable_deployments: Final = [ + deployment + for deployment in healthy_deployments + if not Router._deployment_declares_no_vision( + cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape + group, + ) + ] + if not vision_capable_deployments: + return healthy_deployments + return cast(_RouteDeploymentsT, vision_capable_deployments) # cast-ok: same list shape as the input candidates + @staticmethod def _request_carries_image( messages: list[dict[str, str]] | None, # mutable-ok: handed to resolve_structured_messages as-is @@ -14355,7 +14379,6 @@ class Router: - model context window < message length. For azure openai models, requires 'base_model' is set. - https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models - filter models above rpm limits - if region given, filter out models not in that region / unknown region - - request carries an image and the deployment declares supports_vision false, unless every candidate does - [TODO] function call and model doesn't support function calling """ @@ -14515,20 +14538,9 @@ class Router: model=model, llm_provider="", ) - remaining: Final = tuple( - (idx, deployment) for idx, deployment in param_candidates if idx not in excluded_indices - ) - text_only_indices: Final[frozenset[int]] = ( - frozenset(idx for idx, deployment in remaining if self._deployment_declares_no_vision(deployment, model)) - if remaining and self._request_carries_image(messages, request_settings) - else frozenset() - ) - dropped_indices: Final[frozenset[int]] = excluded_indices | ( - text_only_indices if len(text_only_indices) < len(remaining) else frozenset() - ) - if len(dropped_indices) > 0: + if len(excluded_indices) > 0: # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) - _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in dropped_indices] + _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in excluded_indices] return _returned_deployments @@ -15051,8 +15063,11 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts - list[dict[str, object]], healthy_deployments + batch_candidates: Final = self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( batch_candidates, @@ -16291,8 +16306,11 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts - list[dict[str, object]], healthy_deployments + batch_candidates: Final = self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( batch_candidates, diff --git a/litellm/types/decisions.py b/litellm/types/decisions.py index 5ad6e45e6d2..33b2c677d31 100644 --- a/litellm/types/decisions.py +++ b/litellm/types/decisions.py @@ -650,7 +650,7 @@ class DecisionsIRUsage: class DecisionsIRResponse: model: str | None answers: tuple[DecisionsIRAnswer, ...] - usage: DecisionsIRUsage + usage: DecisionsIRUsage | None extra: Mapping[str, object] = field(default_factory=lambda: _NO_EXTRA) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ff063d2e233..e1b891caaaf 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -59,6 +59,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.reco import ( + RecoConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import ( RepelloAIGuardrailConfigModel, ) @@ -153,6 +156,7 @@ class SupportedGuardrailIntegrations(Enum): TYPESAFE = "typesafe" STRAIKER = "straiker" ALICE = "alice" + RECO = "reco" AGENT_365 = "agent_365" LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" @@ -1398,6 +1402,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, SingulrGuardrailConfigModel, + RecoConfigModel, Agent365GuardrailConfigModel, WingbackGuardrailConfigModel, ): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/reco.py b/litellm/types/proxy/guardrails/guardrail_hooks/reco.py new file mode 100644 index 00000000000..ef3b3935340 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/reco.py @@ -0,0 +1,48 @@ +import uuid +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from .base import GuardrailConfigModel + + +def validate_reco_tenant_id(value: str) -> str: + try: + uuid.UUID(value) + except ValueError as e: + raise ValueError(f"reco_tenant_id must be a valid UUID, got {value!r}") from e + return value + + +class RecoOptionalParams(BaseModel): + """Configuration parameters for the Reco guardrail""" + + model_config = ConfigDict(frozen=True) + + reco_tenant_id: str = Field( + description="Tenant identifier for the Reco account, as a UUID. Sent as the X-Reco-Tenant-Id header on every guardrail request.", + ) + api_base: str = Field( + description="Base URL for the Reco guardrail API. Reco's endpoint is per-region and per-silo, so this must be set explicitly rather than defaulted.", + ) + unreachable_fallback: Literal["fail_closed", "fail_open"] = Field( + default="fail_closed", + description="Behavior when Reco is unreachable. 'fail_closed' (default) blocks the request; 'fail_open' lets it through uninspected. Only applies while fail_on_error is true.", + ) + fail_on_error: bool = Field( + default=True, + description="Behavior on any Reco error, including unreachability. True (default) blocks the request; False lets it through uninspected.", + ) + + @field_validator("reco_tenant_id") + @classmethod + def _check_reco_tenant_id(cls, v: str) -> str: + return validate_reco_tenant_id(v) + + +class RecoConfigModel(GuardrailConfigModel[RecoOptionalParams]): + """Configuration parameters for the Reco guardrail""" + + @staticmethod + def ui_friendly_name() -> str: + return "Reco" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 51c48e6ebf9..742842bd028 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from datetime import datetime from typing import Final, Literal from pydantic import ConfigDict, Field, field_validator, model_validator @@ -83,11 +84,22 @@ class BulkUpdateTeamMemberPermissionsResponse(LiteLLMBaseModel): permissions_appended: list[str] | None = None +class TeamListCallerMembership(LiteLLMBaseModel): + """The calling user's own spend and effective member budget in a team.""" + + spend: float = 0.0 + max_budget: float | None = None + budget_reset_at: datetime | None = None + + model_config = ConfigDict(frozen=True) + + class TeamListItem(LiteLLM_TeamTable): """A team item in the paginated list response, enriched with computed fields.""" members_count: int = 0 keys_count: int = 0 + caller_membership: TeamListCallerMembership | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ee8683f8ae8..cc7ef851111 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -328,6 +328,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_above_12hr_above_272k_tokens: ReadOnly[float | None] cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: ReadOnly[float | None] cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_batches: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_ultrafast: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_creation_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing @@ -3973,6 +3977,10 @@ class MirroredPricingParams(LiteLLMBaseModel): cache_creation_input_token_cost_above_12hr_above_272k_tokens: float | None = None cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: float | None = None cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: float | None = None + cache_creation_input_token_cost_above_12hr_batches: float | None = None + cache_creation_input_token_cost_above_12hr_ultrafast: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: float | None = None tiered_pricing: list[dict[str, Any]] | None = None diff --git a/litellm/utils.py b/litellm/utils.py index acf4594c6b9..e2820f9c682 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6842,6 +6842,18 @@ def get_model_info_helper( cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority=_model_info.get( "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority", None ), + cache_creation_input_token_cost_above_12hr_batches=_model_info.get( + "cache_creation_input_token_cost_above_12hr_batches", None + ), + cache_creation_input_token_cost_above_12hr_ultrafast=_model_info.get( + "cache_creation_input_token_cost_above_12hr_ultrafast", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast", None + ), off_peak_pricing=_model_info.get("off_peak_pricing", None), input_cost_per_character=_model_info.get("input_cost_per_character", None), input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c7dfce0712f..9a7d88b8eda 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -644,9 +644,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, @@ -4339,7 +4343,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -9232,7 +9236,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -9679,7 +9683,7 @@ "deprecation_date": "2027-05-18", "input_cost_per_token": 5.5e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, + "max_input_tokens": 128000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -15659,6 +15663,28 @@ "/v1/images/generations" ] }, + "cerebras/llama-3.1-70b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, + "cerebras/llama-3.1-8b": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-07, + "supports_function_calling": true, + "supports_tool_choice": true + }, "cerebras/llama-3.3-70b": { "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", @@ -33212,6 +33238,7 @@ "github_copilot/claude-opus-4.8-fast": { "cache_creation_input_token_cost": 1.25e-05, "cache_read_input_token_cost": 1e-06, + "deprecation_date": "2026-11-16", "input_cost_per_token": 1e-05, "litellm_provider": "github_copilot", "max_input_tokens": 200000, @@ -37401,6 +37428,7 @@ }, "gpt-5.6-cyber": { "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_12hr": 3.125e-05, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, @@ -42698,6 +42726,29 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5e-07, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://docs.mistral.ai/models/zai-glm-5-3-flash", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/zai-glm-5": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -59514,6 +59565,7 @@ "supports_audio_input": true, "supports_audio_output": true, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -59709,6 +59761,7 @@ "tpm": 250000, "rpm": 10, "gemini_native_audio": true, + "input_cost_per_video_token": 3e-06, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -75362,10 +75415,13 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 3.3e-07, "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2.75e-06, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-amazon-nova-2-5-sonic.md", "supports_audio_input": true, "supports_audio_output": true }, diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py index 384bf3a2918..a80eeacad46 100644 --- a/tests/integration/authorization/test_claude_code_gateway_oauth.py +++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py @@ -15,6 +15,7 @@ from urllib.parse import parse_qs, urlencode, urlsplit import httpx import pytest import yaml +from anthropic import Anthropic from anthropic.types import Message from integration._support import claude_code as cc from integration._support.client import Gateway, eventually, gateway_from_environment @@ -397,7 +398,31 @@ def test_device_authorization_polls_complete_sso_and_mints_a_model_restricted_be "anthropic-beta": cc.CLI_BETA, "x-claude-code-session-id": f"oauth-{user_id}", } - request_body: Final = {**request_template, "model": model_a} + standard_listing: Final = enabled_gateway.client.get("/v1/models", headers=bearer_headers) + gateway_listing: Final = enabled_gateway.client.get( + "/claude_code_gateway/v1/models", headers=bearer_headers + ) + assert standard_listing.status_code == 200, standard_listing.text + assert gateway_listing.status_code == 200, gateway_listing.text + assert gateway_listing.json() == standard_listing.json() + with Anthropic( + base_url=f"{str(enabled_gateway.client.base_url).rstrip('/')}/claude_code_gateway", + auth_token=access_token.access_token, + ) as sdk: + discovered_models: Final = tuple(sdk.models.list()) + assert {model.id for model in discovered_models} == {model_a, model_c} + discovered_model: Final = next(model for model in discovered_models if model.id == model_a) + openai_listing: Final = enabled_gateway.client.get( + "/claude_code_gateway/v1/models", + headers={"authorization": f"Bearer {access_token.access_token}"}, + ) + assert openai_listing.status_code == 200, openai_listing.text + assert openai_listing.json()["object"] == "list" + assert {model["id"] for model in openai_listing.json()["data"]} == {model_a, model_c} + unauthenticated: Final = enabled_gateway.client.get("/claude_code_gateway/v1/models") + assert unauthenticated.status_code == 401, unauthenticated.text + assert upstream.drain() == () + request_body: Final = {**request_template, "model": discovered_model.id} expected_message: Final = _MESSAGE.validate_python( { **_JSON_OBJECT.validate_json( diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index e978a9a316a..95363e92366 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -91,6 +91,10 @@ class CostMapEntry(BaseModel): cache_creation_input_token_cost_above_12hr_above_272k_tokens: float | None = None cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex: float | None = None cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority: float | None = None + cache_creation_input_token_cost_above_12hr_batches: float | None = None + cache_creation_input_token_cost_above_12hr_ultrafast: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: float | None = None + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: float | None = None cache_creation_input_token_cost_above_1hr_above_200k_tokens: float | None = None cache_read_input_token_cost_above_200k_tokens: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None diff --git a/tests/integration/management/test_team_list_v2.py b/tests/integration/management/test_team_list_v2.py index d608d054daa..4c9171310a6 100644 --- a/tests/integration/management/test_team_list_v2.py +++ b/tests/integration/management/test_team_list_v2.py @@ -1,8 +1,8 @@ -from typing import FrozenSet, Optional +from typing import Final, FrozenSet, Optional import pytest -from tests.integration._support.actors import Actor +from tests.integration._support.actors import Actor, ActorSet, GatewayClient pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -139,3 +139,31 @@ async def test_team_list_v2_invalid_status_is_400(gateway_client, actors): headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"}, ) assert resp.status_code == 400, resp.text + + +async def test_team_list_v2_returns_caller_membership_only_for_members( + gateway_client: GatewayClient, + actors: ActorSet, +) -> None: + member: Final = actors.keys[Actor.TEAM_ADMIN] + member_response: Final = await gateway_client.get( + f"/v2/team/list?team_id={actors.team_alpha_id}&user_id={member.user_id}", + headers={"Authorization": f"Bearer {member.cleartext}"}, + ) + + assert member_response.status_code == 200, member_response.text + member_team: Final = member_response.json()["teams"][0] + assert member_team["caller_membership"] == { + "spend": 0.0, + "max_budget": None, + "budget_reset_at": None, + } + + admin: Final = actors.keys[Actor.PROXY_ADMIN] + admin_response: Final = await gateway_client.get( + f"/v2/team/list?team_id={actors.team_alpha_id}", + headers={"Authorization": f"Bearer {admin.cleartext}"}, + ) + + assert admin_response.status_code == 200, admin_response.text + assert admin_response.json()["teams"][0]["caller_membership"] is None diff --git a/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py b/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py new file mode 100644 index 00000000000..100e760d5be --- /dev/null +++ b/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py @@ -0,0 +1,232 @@ +import asyncio +import json +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import ExitStack, asynccontextmanager +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import httpx +import pytest +import pytest_asyncio +import yaml +from aiohttp import web + +from tests.integration._support.anthropic_sse import LIFECYCLE, delta_text, event_types, parse_sse, user_prompt +from tests.integration._support.client import JSON_OBJECT, Gateway +from tests.integration._support.openai_wire import chat_stream +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import OwnedProxy, owned_proxy_process, paused + +_MODEL: Final = "disconnect-repro" +_KEY: Final = "sk-integration-master-key" +_DEADLINE: Final = 10 +_RACE_ATTEMPTS: Final = 20 + + +class _Mode(Enum): + OCCUPY = "occupy" + QUEUED = "queued" + CONTENT = "content" + COMPLETE = "complete" + + +@dataclass(frozen=True, slots=True) +class _Request: + mode: _Mode + headers_allowed: asyncio.Event = field(default_factory=asyncio.Event) + headers_sent: asyncio.Event = field(default_factory=asyncio.Event) + running: asyncio.Event = field(default_factory=asyncio.Event) + released: asyncio.Event = field(default_factory=asyncio.Event) + + +@dataclass(frozen=True, slots=True) +class _Backend: + url: str + requests: asyncio.Queue[_Request] + runner: web.AppRunner + + +@pytest_asyncio.fixture(loop_scope="function") +async def backend() -> AsyncIterator[_Backend]: + requests: Final[asyncio.Queue[_Request]] = asyncio.Queue() + worker: Final = asyncio.Semaphore(1) + + async def complete(request: web.Request) -> web.StreamResponse: + body: Final = JSON_OBJECT.validate_json(await request.read()) + assert body["model"] == "test-model" and body["stream"] is True, body + assert request.headers["Authorization"] == "Bearer integration-provider-key" + pending: Final = _Request(_Mode(user_prompt(body))) + await requests.put(pending) + response: Final = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) + try: + await pending.headers_allowed.wait() + await response.prepare(request) + pending.headers_sent.set() + async with worker: + pending.running.set() + chunks: Final = chat_stream("chatcmpl-disconnect", "test-model", "hello") + if pending.mode in (_Mode.CONTENT, _Mode.COMPLETE): + for chunk in chunks[:2]: + await response.write(chunk) + if pending.mode is _Mode.COMPLETE: + await response.write(chunks[2]) + await response.write_eof() + else: + await asyncio.Event().wait() + finally: + pending.released.set() + return response + + app: Final = web.Application() + app.router.add_post("/v1/chat/completions", complete) + runner: Final = web.AppRunner(app, handler_cancellation=True, shutdown_timeout=0.1) + await runner.setup() + port: Final = allocate_port() + await web.TCPSite(runner, "127.0.0.1", port).start() + try: + yield _Backend(f"http://127.0.0.1:{port}", requests, runner) + finally: + await runner.cleanup() + + +@pytest_asyncio.fixture(loop_scope="function") +async def proxy(backend: _Backend, tmp_path: Path) -> AsyncIterator[OwnedProxy]: + config: Final = tmp_path / "disconnect.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": "custom_openai/test-model", + "api_base": backend.url + "/v1", + "api_key": "integration-provider-key", + }, + } + ], + "litellm_settings": {"anthropic_sse_ping_interval_seconds": 0}, + "router_settings": {"num_retries": 0}, + "general_settings": {"dangerously_permit_weak_or_unset_master_key": True}, + } + ) + ) + with httpx.Client() as client: + template: Final = Gateway(client, _KEY, backend.url) + with owned_proxy_process( + template, + tmp_path, + {"STORE_MODEL_IN_DB": "False"}, + config=config, + database_setup=(), + remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "REDIS_HOST", "REDIS_PORT"), + ) as owned: + try: + yield owned + finally: + await backend.runner.cleanup() + + +def _body(mode: _Mode) -> bytes: + return json.dumps( + {"model": _MODEL, "stream": True, "max_tokens": 16, "messages": [{"role": "user", "content": mode.value}]} + ).encode() + + +@asynccontextmanager +async def _request( + proxy: OwnedProxy, mode: _Mode, endpoint: str = "/v1/messages" +) -> AsyncGenerator[asyncio.StreamReader, None]: + address: Final = urlsplit(str(proxy.gateway.client.base_url)) + reader, writer = await asyncio.open_connection(address.hostname, address.port) + body: Final = _body(mode) + writer.write( + f"POST {endpoint} HTTP/1.1\r\nHost: localhost\r\nAuthorization: Bearer {_KEY}\r\n" + f"Content-Type: application/json\r\nContent-Length: {len(body)}\r\n\r\n".encode() + + body + ) + await writer.drain() + try: + yield reader + finally: + writer.close() + await writer.wait_closed() + + +async def _next(backend: _Backend, mode: _Mode) -> _Request: + pending: Final = await asyncio.wait_for(backend.requests.get(), _DEADLINE) + assert pending.mode is mode + return pending + + +async def _headers(pending: _Request) -> None: + pending.headers_allowed.set() + await asyncio.wait_for(pending.headers_sent.wait(), _DEADLINE) + + +async def _released(pending: _Request) -> None: + try: + await asyncio.wait_for(pending.released.wait(), _DEADLINE) + except TimeoutError: + pytest.fail(f"Backend {pending.mode.value} request still active after client disconnect", pytrace=False) + + +async def _disconnect_queued_request(backend: _Backend, proxy: OwnedProxy) -> None: + with ExitStack() as scheduling: + async with _request(proxy, _Mode.QUEUED): + queued: Final = await _next(backend, _Mode.QUEUED) + scheduling.enter_context(paused(proxy.process)) + await _headers(queued) + assert not queued.running.is_set() + await _released(queued) + assert backend.requests.empty(), "Disconnect must not trigger another upstream request" + + +async def test_disconnect_with_upstream_headers_pending_releases_queued_request( + backend: _Backend, proxy: OwnedProxy +) -> None: + """Hold headers and the client FIN ready together without patching proxy code.""" + async with _request(proxy, _Mode.OCCUPY): + occupied: Final = await _next(backend, _Mode.OCCUPY) + await _headers(occupied) + await asyncio.wait_for(occupied.running.wait(), _DEADLINE) + for attempt in range(_RACE_ATTEMPTS): + await _disconnect_queued_request(backend, proxy) + assert not occupied.released.is_set(), attempt + await _released(occupied) + + +@pytest.mark.parametrize("endpoint", ["/v1/messages", "/v1/chat/completions"]) +async def test_disconnect_after_content_releases_upstream(backend: _Backend, proxy: OwnedProxy, endpoint: str) -> None: + async with _request(proxy, _Mode.CONTENT, endpoint) as response: + pending: Final = await _next(backend, _Mode.CONTENT) + await _headers(pending) + received: Final = await asyncio.wait_for(response.readuntil(b'"hello"'), _DEADLINE) + assert b"200 OK" in received + if endpoint == "/v1/messages": + assert b"content_block_delta" in received + await _released(pending) + assert backend.requests.empty(), "Disconnect must not trigger another upstream request" + + +async def test_messages_stream_completes_with_unchanged_events(backend: _Backend, proxy: OwnedProxy) -> None: + async with httpx.AsyncClient(base_url=str(proxy.gateway.client.base_url), trust_env=False) as client: + response_task: Final = asyncio.create_task( + client.post( + "/v1/messages", + headers={"Authorization": f"Bearer {_KEY}", "Content-Type": "application/json"}, + content=_body(_Mode.COMPLETE), + ) + ) + pending: Final = await _next(backend, _Mode.COMPLETE) + await _headers(pending) + response: Final = await asyncio.wait_for(response_task, _DEADLINE) + assert response.status_code == 200, response.text + events: Final = parse_sse(response.text) + assert event_types(events) == LIFECYCLE + assert delta_text(events) == "hello" + assert pending.released.is_set() + assert backend.requests.empty(), "Successful streaming must not trigger another upstream request" diff --git a/tests/integration/pricing/test_opencode_go_reported_cost_wire.py b/tests/integration/pricing/test_opencode_go_reported_cost_wire.py new file mode 100644 index 00000000000..6ac3bdccf27 --- /dev/null +++ b/tests/integration/pricing/test_opencode_go_reported_cost_wire.py @@ -0,0 +1,198 @@ +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_LIST_COST: Final = 0.00412 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + +# OpenCode's documented usage block (opencode.ai/v2/docs/console/api/inference, "Cost", read 2026-10-24): +# usage.cost is what hit the prepaid balance (0 inside the Go plan), usage.opencode.list_cost is the model price +_GO_BLOCK: Final[Mapping[str, JsonValue]] = { + "cost": 0, + "opencode": {"billing_source": "go", "list_cost": _LIST_COST, "allowance": _LIST_COST}, +} +_CHAT_USAGE: Final[Mapping[str, JsonValue]] = { + "prompt_tokens": 1200, + "completion_tokens": 300, + "total_tokens": 1500, + **_GO_BLOCK, +} +_RESPONSES_USAGE: Final[Mapping[str, JsonValue]] = { + "input_tokens": 1200, + "output_tokens": 300, + "total_tokens": 1500, + **_GO_BLOCK, +} +_MESSAGES_USAGE: Final[Mapping[str, JsonValue]] = {"input_tokens": 1200, "output_tokens": 300, **_GO_BLOCK} + + +def _approx(value: float) -> object: + return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs + + +def _data_frame(event: Mapping[str, JsonValue]) -> bytes: + return f"data: {json.dumps(event)}\n\n".encode() + + +def _event_frame(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _chat_completion() -> bytes: + return json.dumps( + { + "id": "chatcmpl-go", + "object": "chat.completion", + "created": 1, + "model": "glm-5.3", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "ok"}}], + "usage": dict(_CHAT_USAGE), + } + ).encode() + + +def _chat_stream() -> tuple[bytes, ...]: + text: Final = { + "id": "chatcmpl-go", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-5.3", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}], + } + finish: Final = {**text, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]} + final_usage: Final = {**text, "choices": [], "usage": dict(_CHAT_USAGE)} + return (*(_data_frame(event) for event in (text, finish, final_usage)), b"data: [DONE]\n\n") + + +def _responses_response() -> Mapping[str, JsonValue]: + return { + "id": "resp_go", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-6-luna", + "output": [ + { + "id": "msg_go", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": dict(_RESPONSES_USAGE), + } + + +def _responses_stream() -> tuple[bytes, ...]: + response: Final = _responses_response() + return ( + _event_frame( + {"type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress"}} + ), + _event_frame( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_go", + "output_index": 0, + "content_index": 0, + "delta": "ok", + } + ), + _event_frame({"type": "response.completed", "sequence_number": 2, "response": response}), + ) + + +def _anthropic_message() -> Mapping[str, JsonValue]: + return { + "id": "msg_go", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": dict(_MESSAGES_USAGE), + } + + +def _anthropic_stream() -> tuple[bytes, ...]: + start: Final = {**_anthropic_message(), "content": [], "stop_reason": None} + return ( + _event_frame( + {"type": "message_start", "message": {**start, "usage": {"input_tokens": 1200, "output_tokens": 0}}} + ), + _event_frame({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _event_frame({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}), + _event_frame({"type": "content_block_stop", "index": 0}), + _event_frame( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": dict(_MESSAGES_USAGE), + } + ), + _event_frame({"type": "message_stop"}), + ) + + +def _respond(request: Request) -> Reply: + assert request.method == "POST", request.method + streaming: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + if request.target.endswith("/chat/completions"): + return ( + Reply(content_type="text/event-stream", chunks=_chat_stream()) + if streaming + else Reply(body=_chat_completion()) + ) + if request.target.endswith("/responses"): + if streaming: + return Reply(content_type="text/event-stream", chunks=_responses_stream()) + return Reply(body=json.dumps(_responses_response()).encode()) + assert request.target.endswith("/v1/messages"), request.target + if streaming: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream()) + return Reply(body=json.dumps(_anthropic_message()).encode()) + + +def _logged_spend(call_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return float(rows[0]["spend"]) + + +def _body(path: str, model: str, stream: bool) -> Mapping[str, JsonValue]: + if path == "/v1/responses": + return {"model": model, "input": "hi", "stream": stream} + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}], "stream": stream} + + +@pytest.mark.parametrize("stream", (False, True), ids=("non-stream", "stream")) +@pytest.mark.parametrize( + ("path", "model"), + ( + ("/v1/chat/completions", "opencode_go/glm-5.3"), + ("/v1/responses", "opencode_go/gpt-6-luna"), + ("/v1/messages", "opencode_go/claude-haiku-5-5"), + ), + ids=("chat", "responses", "messages"), +) +def test_request_is_billed_at_the_reported_list_cost(gateway: Gateway, path: str, model: str, stream: bool) -> None: + with wire_server(_respond) as wire, gateway.scenario() as scenario: + deployment: Final = scenario.model(model=model, api_base=wire.url, api_key="synthetic-go-key") + response: Final = gateway.request("POST", path, _body(path, deployment, stream)) + assert response.status_code == 200, response.text + assert "ok" in response.text + if not stream: + assert float(response.headers["x-litellm-response-cost"]) == _approx(_LIST_COST) + assert _logged_spend(response.headers["x-litellm-call-id"]) == _approx(_LIST_COST) diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index ea29914ba65..89ca674ef31 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -660,3 +660,22 @@ def test_a_deployment_whose_provider_has_no_decisions_support_is_refused_naming_ for provider in _PROVIDERS: assert provider.name in response.text, response.text assert _upstream_calls(gateway, handle) == [] + + +def test_typesafe_gateway_preserves_missing_usage(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + body = {"model": "jev-1.13.0", "answers": {"exp": {"type": "noul", "noul": 0.8}}} + handle = _register(scenario, body) + model = _deployment(scenario, handle, _PROVIDERS[1]) + response = gateway.request( + "POST", + "/v1/systemone", + { + "model": model, + "state": {"finding": "x"}, + "questions": {"exp": {"type": "noul", "instructions": "Is it exploitable?"}}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["usage"] is None + assert response.json()["answers"] == body["answers"] diff --git a/tests/integration/spend/test_team_member_total_spend_reconciliation.py b/tests/integration/spend/test_team_member_total_spend_reconciliation.py new file mode 100644 index 00000000000..f5ccae0d12b --- /dev/null +++ b/tests/integration/spend/test_team_member_total_spend_reconciliation.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import asyncio +import json +import os +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.utils import PrismaClient, ProxyLogging +from tests.integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value +from tests.integration._support.database import read_rows, write_rows +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + +_CALL_COST: Final = 0.02 +JSON_VALUE: Final = TypeAdapter(JsonValue) +_LEGACY_SPEND: Final = 612.4 +_LEGACY_TOTAL: Final = 148.1 +_HEALTHY_SPEND: Final = 12.25 +_HEALTHY_TOTAL: Final = 97.8 + + +def _completion() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +def _priced_model(scenario: Scenario) -> str: + return scenario.model(api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002) + + +def _member_with_budget(gateway: Gateway, scenario: Scenario, team_id: str, *, max_budget: float) -> str: + user_id: Final = scenario.user(user_role="internal_user") + gateway.post( + "/team/member_add", + {"team_id": team_id, "member": {"role": "user", "user_id": user_id}, "max_budget_in_team": max_budget}, + ) + return user_id + + +def _seed_membership(user_id: str, team_id: str, *, spend: float, total_spend: float) -> None: + write_rows( + 'UPDATE "LiteLLM_TeamMembership" SET spend = %s::float8, total_spend = %s::float8 ' + "WHERE user_id = %s AND team_id = %s", + (str(spend), str(total_spend), user_id, team_id), + ) + + +def _membership(user_id: str, team_id: str) -> tuple[float, float]: + rows: Final = read_rows( + 'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s', + (user_id, team_id), + ) + assert len(rows) == 1, rows + return float(str(rows[0]["spend"])), float(str(rows[0]["total_spend"])) + + +def _membership_reaches(user_id: str, team_id: str, *, total_spend: float) -> tuple[float, float]: + return eventually( + lambda: _membership(user_id, team_id), + lambda row: row[1] == pytest.approx(total_spend), + seconds=70, + ) + + +def _served(gateway: Gateway, provider: SharedProvider, model: str, key: str) -> None: + provider.expect(_completion()) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"reconcile {uuid.uuid4().hex}"}]}, + key=key, + ) + assert response.status_code == 200, response.text + assert len(provider.received()) == 1 + + +def _reset_spend(gateway: Gateway, team_id: str, user_id: str, reset_to: float) -> httpx.Response: + return gateway.request("POST", f"/team/{team_id}/member/{user_id}/reset_spend", {"reset_to": reset_to}) + + +def _membership_figures(row: JsonValue) -> tuple[str, tuple[float, float]]: + membership: Final = object_value(row) + return string_value(membership["user_id"]), ( + float(str(membership["spend"])), + float(str(membership["total_spend"])), + ) + + +def _reported_memberships(gateway: Gateway, path: str, params: dict[str, str]) -> dict[str, tuple[float, float]]: + """/team/info answers with one team object, /team/list with a list of them.""" + response: Final = gateway.request("GET", path, params=params) + assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}" + body: Final = JSON_VALUE.validate_json(response.content) + teams: Final = [object_value(body)] if isinstance(body, dict) else [object_value(team) for team in list_value(body)] + rows: Final = [ + row for team in teams for row in list_value(team["team_memberships"]) + ] # comprehension-ok: flatten memberships across teams + return dict(_membership_figures(row) for row in rows) + + +def test_a_flush_lifts_a_legacy_total_to_the_cycle_spend_and_leaves_a_healthy_row_on_its_baseline( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team_id: Final = scenario.team(models=[model]) + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + _seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL) + legacy_key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model]) + healthy_key: Final = scenario.key(user_id=healthy, team_id=team_id, models=[model]) + + _served(gateway, provider, model, legacy_key) + _served(gateway, provider, model, healthy_key) + + legacy_row: Final = _membership_reaches( + legacy, team_id, total_spend=max(_LEGACY_SPEND, _LEGACY_TOTAL) + _CALL_COST + ) + assert legacy_row[0] == pytest.approx(_LEGACY_SPEND + _CALL_COST) + healthy_row: Final = _membership_reaches(healthy, team_id, total_spend=_HEALTHY_TOTAL + _CALL_COST) + assert healthy_row[0] == pytest.approx(_HEALTHY_SPEND + _CALL_COST) + + +def test_team_info_and_team_list_report_a_legacy_total_no_lower_than_the_cycle_spend(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + _seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL) + + info: Final = _reported_memberships(gateway, "/team/info", {"team_id": team_id}) + listed: Final = _reported_memberships(gateway, "/team/list", {"team_id": team_id}) + + for reported in (info, listed): + assert reported[legacy] == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_SPEND)) + assert reported[healthy] == (pytest.approx(_HEALTHY_SPEND), pytest.approx(_HEALTHY_TOTAL)) + assert _membership(legacy, team_id) == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_TOTAL)) + + +@pytest.mark.parametrize("reset_to", [0.0, 100.0], ids=["to_zero", "to_nonzero"]) +def test_an_admin_reset_keeps_the_legacy_cycle_spend_as_the_lifetime_floor(gateway: Gateway, reset_to: float) -> None: + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + + response: Final = _reset_spend(gateway, team_id, legacy, reset_to) + + assert response.status_code == 200, response.text + assert _membership(legacy, team_id) == (pytest.approx(reset_to), pytest.approx(_LEGACY_SPEND)) + + +def test_concurrent_admin_resets_leave_one_winner_and_a_lifetime_floor_at_the_old_cycle_spend( + gateway: Gateway, +) -> None: + """Two admins reset the same member at once. Each UPDATE re-checks reset_to <= spend, so + the higher reset either lands first or is refused with 409 once the lower one has landed. + Whatever the order, the row ends at the lower value with the old cycle spend preserved.""" + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + + with ThreadPoolExecutor(max_workers=2) as pool: + higher, lower = pool.map(lambda value: _reset_spend(gateway, team_id, legacy, value), (600.0, 500.0)) + + assert lower.status_code == 200, lower.text + assert higher.status_code in (200, 409), higher.text + assert _membership(legacy, team_id) == (pytest.approx(500.0), pytest.approx(_LEGACY_SPEND)) + + +def test_an_admin_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team_id: Final = scenario.team(models=[model]) + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model]) + + _served(gateway, provider, model, key) + reset: Final = _reset_spend(gateway, team_id, legacy, 0.0) + assert reset.status_code == 200, reset.text + + spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST) + assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total) + assert total >= spend + + +def _expire_member_budget(user_id: str, team_id: str, *, budget_duration: str) -> None: + write_rows( + "UPDATE \"LiteLLM_BudgetTable\" SET budget_duration = %s, budget_reset_at = NOW() - INTERVAL '1 hour' " + 'WHERE budget_id = (SELECT budget_id FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s)', + (budget_duration, user_id, team_id), + ) + + +async def _run_scheduled_reset() -> None: + client: Final = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()) + ) + await client.connect() + try: + await ResetBudgetJob( + ProxyLogging(user_api_key_cache=DualCache()), client + ).reset_budget_for_litellm_budget_table() + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize( + "spend, max_budget, carried", + [(8.0, 10.0, 0.0), (10.0, 10.0, 0.0), (15.0, 10.0, 5.0)], + ids=["below_cap", "at_cap", "above_cap_carries_overage"], +) +def test_a_scheduled_reset_lifts_the_legacy_total_before_zeroing_or_carrying_the_cycle_spend( + gateway: Gateway, monkeypatch: pytest.MonkeyPatch, spend: float, max_budget: float, carried: float +) -> None: + monkeypatch.setattr(litellm, "budget_rollover", True) + with gateway.scenario() as scenario: + team_id: Final = scenario.team() + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=max_budget) + _seed_membership(legacy, team_id, spend=spend, total_spend=3.0) + _expire_member_budget(legacy, team_id, budget_duration="1d") + + asyncio.run(_run_scheduled_reset()) + + assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend)) + asyncio.run(_run_scheduled_reset()) + assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend)) + + +def test_a_scheduled_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost( + gateway: Gateway, provider: SharedProvider +) -> None: + with gateway.scenario() as scenario: + model: Final = _priced_model(scenario) + team_id: Final = scenario.team(models=[model]) + legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000) + _seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL) + _expire_member_budget(legacy, team_id, budget_duration="1d") + key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model]) + + _served(gateway, provider, model, key) + asyncio.run(_run_scheduled_reset()) + + spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST) + assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total) + assert total >= spend diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 84e7a970195..e2385a2e215 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -3799,3 +3799,35 @@ def test_nace_upstream_failures_preserve_status_and_message( ) assert getattr(error.value, "status_code") == status assert upstream.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + "usage", [None, {"input_tokens": 0, "output_tokens": 0}, {"input_tokens": 12, "output_tokens": 1}] +) +async def test_native_decisions_distinguishes_missing_usage_from_reported_zero( + is_async: bool, usage: dict[str, int] | None, respx_mock: respx.MockRouter +) -> None: + body = { + **{key: value for key, value in _RESPONSE.items() if key != "usage"}, + "model": "jev-1.13.0", + **({"usage": usage} if usage is not None else {}), + } + respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=body) + params = {"model": "typesafe/jev-1.13.0", "state": "review", "questions": _QUESTIONS, "api_key": "synthetic-key"} + response = await litellm.adecisions(**params) if is_async else litellm.decisions(**params) + + assert isinstance(response, DecisionsResponse) + if usage is None: + assert response.usage is None + assert response._hidden_params.get("response_cost") is None + with pytest.raises(ValueError, match="provider did not report usage"): + litellm.completion_cost(completion_response=response) + else: + assert response.usage is not None + assert response.usage.input_tokens == usage["input_tokens"] + assert response.usage.output_tokens == usage["output_tokens"] + expected = litellm.completion_cost(completion_response=response) + assert response._hidden_params["response_cost"] == pytest.approx(expected) + assert (expected == 0) == (usage["input_tokens"] == usage["output_tokens"] == 0) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 1729db94ed8..21b527250c4 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2179,6 +2179,48 @@ def test_generic_cost_per_token_falls_back_to_default_for_missing_12h_rate() -> assert result[0] == pytest.approx(100 * model_info["cache_creation_input_token_cost"]) +@pytest.mark.parametrize(("prompt_tokens", "twelve_hour_rate"), [(10_000, 7e-6), (300_000, 14e-6)]) +def test_generic_cost_per_token_prices_ultrafast_12h_cache_writes_at_the_ultrafast_rate( + prompt_tokens: int, twelve_hour_rate: float +) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_token_ultrafast": 5e-6, + "cache_creation_input_token_cost": 2e-6, + "cache_creation_input_token_cost_ultrafast": 3e-6, + "cache_creation_input_token_cost_above_12hr": 4e-6, + "cache_creation_input_token_cost_above_12hr_ultrafast": 7e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 8e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 14e-6, + } + usage: Final = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=0, + total_tokens=prompt_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cache_creation_tokens=5_000, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=3_000, ephemeral_12h_input_tokens=2_000 + ), + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="test-ultrafast-12h-cache-write", + usage=usage, + custom_llm_provider="openai", + service_tier="ultrafast", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + (prompt_tokens - 5_000) * model_info["input_cost_per_token_ultrafast"] + + 3_000 * model_info["cache_creation_input_token_cost_ultrafast"] + + 2_000 * twelve_hour_rate + ) + + def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): """ Regression test: when cache_creation_tokens is 0 but cache_creation_token_details @@ -4079,6 +4121,44 @@ def test_get_batch_cost_rates_has_no_cache_write_rate_without_a_cache_write_batc assert rates.cache_creation is None +@pytest.mark.parametrize(("prompt_tokens", "expected"), [(1_000, 2.5e-7), (272_000, 2.5e-7), (300_000, 5e-7)]) +def test_get_batch_cost_rates_reads_the_12h_batch_cache_write_rate_for_the_crossed_tier( + prompt_tokens: int, expected: float +) -> None: + from litellm.litellm_core_utils.llm_cost_calc.utils import get_batch_cost_rates + + rates: Final = get_batch_cost_rates( + _batch_rates_model_info( + input_cost_per_token_batches=1e-7, + cache_creation_input_token_cost_batches=1.25e-7, + cache_creation_input_token_cost_above_12hr_batches=2.5e-7, + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches=5e-7, + ), + Usage(prompt_tokens=prompt_tokens, completion_tokens=1, total_tokens=prompt_tokens + 1), + "openai", + ) + + assert rates.cache_creation == 1.25e-7 + assert rates.cache_creation_rates_by_ttl == {"12hr": expected} + + +def test_get_batch_cost_rates_has_no_12h_cache_write_rate_without_a_12h_batch_key() -> None: + from litellm.litellm_core_utils.llm_cost_calc.utils import get_batch_cost_rates + + rates: Final = get_batch_cost_rates( + _batch_rates_model_info( + input_cost_per_token_batches=1e-7, + cache_creation_input_token_cost_batches=1.25e-7, + cache_creation_input_token_cost_above_12hr=2.5e-7, + cache_creation_input_token_cost_above_12hr_ultrafast=5e-7, + ), + Usage(prompt_tokens=300_000, completion_tokens=1, total_tokens=300_001), + "openai", + ) + + assert rates.cache_creation_rates_by_ttl == {} + + @pytest.mark.parametrize("model_base", ["gpt-6-sol", "gpt-6-luna"]) def test_azure_gpt_6_foundry_price_sheet(_local_model_cost_map, model_base): """Azure Foundry hosts gpt-6-sol and gpt-6-luna at OpenAI's Global rates, with the diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 8509b1d6f7d..561ac9b2b33 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -76,6 +76,7 @@ import inspect from typing import List, Optional from litellm.types.utils import LiteLLMCommonStrings + @pytest.fixture def logging_obj(): return LitellmLogging( @@ -346,8 +347,10 @@ async def test_mcp_native_structured_replacement_must_match_returned_content( ) logging_obj.dynamic_success_callbacks = [NativeReplacement()] returned = await logging_obj.async_post_mcp_tool_call_hook( - kwargs={"original_response": result}, response_obj=result, - start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), + kwargs={"original_response": result}, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), ) assert returned is result assert result.content == [TextContent(type="text", text="native-safe" if same_content else "final-safe")] @@ -356,7 +359,6 @@ async def test_mcp_native_structured_replacement_must_match_returned_content( assert "SECRET-1234" not in result.model_dump_json() - @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["none", "wrapper", "exception"]) @pytest.mark.parametrize("structured", [False, True]) @@ -374,8 +376,10 @@ async def test_mcp_direct_content_edit_invalidates_stale_structured_data(logging ) logging_obj.dynamic_success_callbacks = [DirectRedactor()] returned = await logging_obj.async_post_mcp_tool_call_hook( - kwargs={"original_response": result}, response_obj=result, - start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), + kwargs={"original_response": result}, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), ) assert returned is result assert result.content == [TextContent(type="text", text="[REDACTED]")] @@ -2620,7 +2624,6 @@ async def test_completed_success_selects_one_implemented_hook( assert (thread_id == loop_thread) is (kind != "sync") - @pytest.mark.parametrize("for_async_request,split_registries", [(False, False), (True, False), (True, True)]) @pytest.mark.parametrize("private_first", [False, True]) @pytest.mark.parametrize("response_kind", ["responses", "completion"]) @@ -3312,8 +3315,13 @@ async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_g class RecordingShadowLogger(ShadowEvalLogger): async def async_log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: object, - end_time: object, *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: object, + end_time: object, + *, + guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: shadow_snapshots.append(guardrail_snapshot) await super().async_log_success_event( @@ -3322,13 +3330,20 @@ async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_g class RecordingLogger(CustomLogger): async def async_log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object, + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: object, + end_time: object, ) -> None: other_payloads.append(kwargs) class LoggingGuardrail(CustomGuardrail): async def async_logging_hook( - self, kwargs: dict[str, object], result: object, call_type: str, + self, + kwargs: dict[str, object], + result: object, + call_type: str, ) -> tuple[dict[str, object], object]: hook_snapshots.append(logging_obj.shadow_eval_request_snapshot) if hook_mode == "raises": @@ -3340,26 +3355,36 @@ async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_g "user_api_key_hash": "test-key", } snapshot: Final = GuardrailRequestSnapshot.capture( - {"messages": [{"role": "user", "content": "snapshot-only"}]}, metadata, + {"messages": [{"role": "user", "content": "snapshot-only"}]}, + metadata, ) assert snapshot is not None shadow: Final = RecordingShadowLogger(prisma_provider=no_prisma, jobs_cache=InMemoryCache()) guardrail: Final = LoggingGuardrail( - guardrail_name="late-mask", default_on=True, + guardrail_name="late-mask", + default_on=True, event_hook=GuardrailEventHooks.pre_call if hook_mode == "disabled" else GuardrailEventHooks.logging_only, ) monkeypatch.setattr(litellm, "_async_success_callback", []) logging_obj: Final = LitellmLogging( - model="test-model", messages=[], stream=stream, call_type="anthropic_messages", - start_time=datetime.datetime.now(), litellm_call_id="private-snapshot", function_id="private-snapshot", + model="test-model", + messages=[], + stream=stream, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="private-snapshot", + function_id="private-snapshot", dynamic_async_success_callbacks=[shadow, RecordingLogger(), guardrail], ) logging_obj.update_messages([{"role": "user", "content": "logged input"}]) logging_obj.update_environment_variables(litellm_params={"metadata": metadata}, optional_params={}) logging_obj.shadow_eval_request_snapshot = snapshot payload: Final = { - "id": "private-snapshot", "call_type": "anthropic_messages", "metadata": metadata, - "model_group": "test-model", "model_parameters": {}, + "id": "private-snapshot", + "call_type": "anthropic_messages", + "metadata": metadata, + "model_group": "test-model", + "model_parameters": {}, } await logging_obj.async_success_handler(result=ModelResponse(), standard_logging_object=payload) @@ -3816,7 +3841,9 @@ def test_sentry_event_scrubber_initialization(monkeypatch): call_args = mock_init.call_args[1] assert call_args["send_default_pii"] is False assert call_args["event_scrubber"].recursive is True - assert {name.lower() for name in SENTRY_PII_DENYLIST} <= {name.lower() for name in call_args["event_scrubber"].denylist} + assert {name.lower() for name in SENTRY_PII_DENYLIST} <= { + name.lower() for name in call_args["event_scrubber"].denylist + } assert call_args["before_send"] is call_args["before_send_transaction"] @@ -3833,7 +3860,9 @@ def test_sentry_send_default_pii_opt_in(monkeypatch): call_args = mock_init.call_args[1] assert call_args["send_default_pii"] is True - assert not {name.lower() for name in SENTRY_PII_DENYLIST} & {name.lower() for name in call_args["event_scrubber"].denylist} + assert not {name.lower() for name in SENTRY_PII_DENYLIST} & { + name.lower() for name in call_args["event_scrubber"].denylist + } def test_get_masked_values(): @@ -5103,12 +5132,10 @@ def test_reported_cost_respects_configured_pricing_and_additional_charges( else {} ), } - } - } - ) - logging_obj.model_call_details.update( - {"additional_response_cost": additional_cost, "cache_hit": cache_hit} + }, + }, ) + logging_obj.model_call_details.update({"additional_response_cost": additional_cost, "cache_hit": cache_hit}) logging_obj._provider_reported_cost = reported_cost response: Final = ModelResponse( model=pricing_model, @@ -7055,8 +7082,11 @@ async def test_background_interaction_completion_rebills_after_in_progress_succe class SyncRecorder(CustomLogger): def log_success_event( - self, kwargs: dict[str, object], response_obj: object, - start_time: dt.datetime, end_time: dt.datetime, + self, + kwargs: dict[str, object], + response_obj: object, + start_time: dt.datetime, + end_time: dt.datetime, ) -> None: events.put((response_obj, kwargs.get("response_cost"))) logged.set() @@ -7089,11 +7119,11 @@ async def test_background_interaction_completion_rebills_after_in_progress_succe assert logging_obj.model_call_details["response_cost"] > 0 assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 - assert events.qsize() == 2 assert events.get_nowait() == (in_progress, None) assert events.get_nowait() == (completed, logging_obj.model_call_details["response_cost"]) + @pytest.mark.asyncio async def test_background_interaction_completion_prices_the_settled_body_itself(): """ @@ -8127,7 +8157,9 @@ def test_debugging_log_with_json_logs_tolerates_missing_headers(logging_obj, mon logging_obj.litellm_request_debug = True with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning: - logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers=None, additional_args={}) + logging_obj._print_llm_call_debugging_log( + api_base="https://api.openai.com/v1", headers=None, additional_args={} + ) assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"] @@ -10378,8 +10410,13 @@ class TestAzurePTUSpilloverCost: assert payload is not None model_map_value: Final = payload["model_map_information"]["model_map_value"] assert model_map_value is not None - assert model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"] - assert model_map_value["output_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"] + assert ( + model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"] + ) + assert ( + model_map_value["output_cost_per_token"] + == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"] + ) assert response_cost == pytest.approx( 100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"] ) @@ -10494,8 +10531,6 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() assert logging_obj.response_cost_calculator(result=assembled) == 0.0042 - - def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): logging_obj: Final = _responses_stream_logging_obj() inner_response: Final = ResponsesAPIResponse( @@ -10725,9 +10760,7 @@ def test_standard_logging_retries(): ] ) - with patch.object( - customHandler, "log_failure_event", new=MagicMock() - ) as mock_client: + with patch.object(customHandler, "log_failure_event", new=MagicMock()) as mock_client: try: router.completion( model="gpt-3.5-turbo", @@ -10739,19 +10772,10 @@ def test_standard_logging_retries(): pass assert mock_client.call_count == 2 + assert mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"]["trace_id"] is not None assert ( - mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"][ - "trace_id" - ] - is not None - ) - assert ( - mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"][ - "trace_id" - ] - == mock_client.call_args_list[1].kwargs["kwargs"][ - "standard_logging_object" - ]["trace_id"] + mock_client.call_args_list[0].kwargs["kwargs"]["standard_logging_object"]["trace_id"] + == mock_client.call_args_list[1].kwargs["kwargs"]["standard_logging_object"]["trace_id"] ) @@ -11018,9 +11042,7 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch): custom_logger_init_args={}, ) assert created is None - assert not [ - cb for cb in logging_module._in_memory_loggers if getattr(cb, "callback_name", None) == "signoz" - ] + assert not [cb for cb in logging_module._in_memory_loggers if getattr(cb, "callback_name", None) == "signoz"] finally: logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) @@ -11162,7 +11184,11 @@ def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_b } } logging_obj.update_environment_variables( - model="gpt-5.4-nano", user="", optional_params={}, litellm_params=litellm_params, custom_llm_provider="openai" + model="gpt-5.4-nano", + user="", + optional_params={}, + litellm_params=litellm_params, + custom_llm_provider="openai", ) response: Final = ModelResponse( model="gpt-5.4-nano", @@ -11474,10 +11500,7 @@ def test_enterprise_alerting_loggers_falls_back_without_disabling_callback_contr ) logging_module._enterprise_alerting_loggers.cache_clear() for module_name in tuple(sys.modules): - if any( - module_name == prefix or module_name.startswith(f"{prefix}.") - for prefix in module_prefixes - ): + if any(module_name == prefix or module_name.startswith(f"{prefix}.") for prefix in module_prefixes): monkeypatch.delitem(sys.modules, module_name, raising=False) monkeypatch.setattr(sys, "meta_path", [BlockPagerDutyFinder(), *sys.meta_path]) @@ -11597,13 +11620,16 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest_asyncio.fixture(loop_scope="function") async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]: yield await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0 + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -11641,6 +11667,7 @@ def isolate_litellm_state(): if attr in _DEFAULTS: setattr(litellm, attr, _DEFAULTS[attr]) + _LIST_ATTRS = ( "callbacks", "success_callback", @@ -11668,6 +11695,7 @@ _SCALAR_ATTRS = ( _DEFAULTS: dict = {} + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -11690,6 +11718,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "response_obj,expected_values", @@ -11737,6 +11766,7 @@ def test_get_usage(response_obj, expected_values): assert usage.completion_tokens == expected_values[1] assert usage.total_tokens == expected_values[2] + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_usage_from_image_generation_response(): """ @@ -11784,6 +11814,7 @@ def test_get_usage_from_image_generation_response(): assert usage.completion_tokens_details.image_tokens == 272 assert usage.completion_tokens_details.text_tokens == 100 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_additional_headers(): additional_headers = { @@ -11824,10 +11855,12 @@ def test_get_additional_headers(): assert additional_logging_headers.get("llm_provider-request-id") == "req_01F6CycZZPSHKRCCctcS1Vto" assert additional_logging_headers.get("llm_provider-anthropic-ratelimit-requests-reset") == "2024-10-29T23:57:40Z" + def all_fields_present(standard_logging_metadata: StandardLoggingMetadata): for field in StandardLoggingMetadata.__annotations__.keys(): assert field in standard_logging_metadata + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "metadata_key, metadata_value", @@ -11859,6 +11892,7 @@ def test_get_standard_logging_metadata(metadata_key, metadata_value): # Assert that the specific metadata field is set correctly assert standard_logging_metadata[metadata_key] == metadata_value + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_user_api_key_hash(): valid_hash = "a" * 64 # 64 character string @@ -11866,6 +11900,7 @@ def test_get_standard_logging_metadata_user_api_key_hash(): result = StandardLoggingPayloadSetup.get_standard_logging_metadata(metadata) assert result["user_api_key_hash"] == valid_hash + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_invalid_user_api_key(): invalid_hash = "not_a_valid_hash" @@ -11874,6 +11909,7 @@ def test_get_standard_logging_metadata_invalid_user_api_key(): all_fields_present(result) assert result["user_api_key_hash"] is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_non_string_user_api_key(): """Non-string user_api_key should not be set as user_api_key_hash.""" @@ -11882,6 +11918,7 @@ def test_get_standard_logging_metadata_non_string_user_api_key(): all_fields_present(result) assert result["user_api_key_hash"] is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_none_user_api_key(): """None user_api_key should not be set as user_api_key_hash.""" @@ -11890,6 +11927,7 @@ def test_get_standard_logging_metadata_none_user_api_key(): all_fields_present(result) assert result["user_api_key_hash"] is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_standard_logging_metadata_invalid_keys(): metadata = { @@ -11903,6 +11941,7 @@ def test_get_standard_logging_metadata_invalid_keys(): assert "invalid_key" not in result assert "another_invalid_key" not in result + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_cleanup_timestamps(): """Test cleanup_timestamps with different input types""" @@ -11937,6 +11976,7 @@ def test_cleanup_timestamps(): with pytest.raises(ValueError, match="start_time is required, got=invalid of type "): StandardLoggingPayloadSetup.cleanup_timestamps("invalid", end_float, completion_float) + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_model_cost_information(): """Test get_model_cost_information with different inputs""" @@ -11969,6 +12009,7 @@ def test_get_model_cost_information(): # assert all fields in StandardLoggingModelInformation are present assert all(field in result for field in StandardLoggingModelInformation.__annotations__) + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_model_cost_information_custom_pricing_uses_base_model(): result = StandardLoggingPayloadSetup.get_model_cost_information( @@ -11980,6 +12021,7 @@ def test_get_model_cost_information_custom_pricing_uses_base_model(): assert result["model_map_value"] is not None assert result["model_map_key"] != "invoke_test_claude" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_deployment_when_no_base_model(): """metadata["deployment"] is used for cost-map lookup when base_model is not set.""" @@ -12037,6 +12079,7 @@ def test_standard_logging_payload_uses_deployment_when_no_base_model(): assert payload["model_map_information"]["model_map_value"] is not None assert payload["model_map_information"]["model_map_key"] != "invoke_test_claude" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_hidden_params(): """Test get_hidden_params with different inputs""" @@ -12072,6 +12115,7 @@ def test_get_hidden_params(): # assert all fields in StandardLoggingHiddenParams are present assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_final_response_obj(): """Test get_final_response_obj with different input types and redaction scenarios""" @@ -12103,6 +12147,7 @@ def test_get_final_response_obj(): # Reset litellm.turn_off_message_logging to its original value litellm.turn_off_message_logging = False + def testget_standard_logging_payload_trace_id(): """Test get_standard_logging_payload_trace_id with different input scenarios""" # Test case 1: When litellm_trace_id is provided in litellm_params @@ -12140,6 +12185,7 @@ def testget_standard_logging_payload_trace_id(): assert result == "12345" assert isinstance(result, str) + def testget_standard_logging_payload_trace_id_prioritizes_trace_id_when_flag_on(monkeypatch): """With request_correlation_in_logs on, an explicit litellm_trace_id wins over litellm_session_id.""" from unittest.mock import MagicMock @@ -12154,6 +12200,7 @@ def testget_standard_logging_payload_trace_id_prioritizes_trace_id_when_flag_on( ) assert result == "the-trace-id" + def testget_standard_logging_payload_trace_id_prioritizes_session_id_when_flag_off(monkeypatch): """With request_correlation_in_logs off (default), legacy behavior is preserved: litellm_session_id still wins over litellm_trace_id.""" @@ -12169,6 +12216,7 @@ def testget_standard_logging_payload_trace_id_prioritizes_session_id_when_flag_o ) assert result == "the-session-id" + def testget_standard_logging_payload_session_id_when_flag_on(monkeypatch): """Test get_standard_logging_payload_session_id with different input scenarios, flag enabled""" from unittest.mock import MagicMock @@ -12221,6 +12269,7 @@ def testget_standard_logging_payload_session_id_when_flag_on(monkeypatch): ) assert result == "" + def testget_standard_logging_payload_session_id_empty_when_flag_off(monkeypatch): """When request_correlation_in_logs is off (default), session_id is always empty, even if litellm_session_id was explicitly supplied - preserves the pre-existing @@ -12237,6 +12286,7 @@ def testget_standard_logging_payload_session_id_empty_when_flag_off(monkeypatch) ) assert result == "" + @pytest.fixture def restore_verbose_logger_level() -> Iterator[None]: original_level: Final = verbose_logger.level @@ -12279,6 +12329,7 @@ def test_truncate_standard_logging_payload(): assert len(str(truncated["response"])) < 10_500 assert len(str(truncated["error_str"])) < 10_500 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_truncate_standard_logging_payload_keeps_a_partial_payload_intact(): """A payload built with only some of its fields comes back with exactly those keys and values""" @@ -12287,12 +12338,14 @@ def test_truncate_standard_logging_payload_keeps_a_partial_payload_intact(): assert _custom_logger.truncate_standard_logging_payload_content(partial_payload) == partial_payload + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_strip_trailing_slash(): common_api_base = "https://api.test.com" assert StandardLoggingPayloadSetup.strip_trailing_slash(common_api_base + "/") == common_api_base assert StandardLoggingPayloadSetup.strip_trailing_slash(common_api_base) == common_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_error_information(): """Test get_error_information with different types of exceptions""" @@ -12329,6 +12382,7 @@ def test_get_error_information(): assert result["llm_provider"] == "openai" assert result["error_message"] == "litellm.RateLimitError: Test error" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_response_time(): """Test get_response_time with different streaming scenarios""" @@ -12364,6 +12418,7 @@ def test_get_response_time(): # For streaming, should return completion_start_time - start_time assert response_time == 2.0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "metadata, expected_requester_metadata", @@ -12386,6 +12441,7 @@ def test_standard_logging_metadata_requester_metadata(metadata, expected_request result = StandardLoggingPayloadSetup.get_standard_logging_metadata(metadata) assert result["requester_metadata"] == expected_requester_metadata + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_cost_breakdown_in_standard_logging_payload(): """ @@ -12475,6 +12531,7 @@ def test_cost_breakdown_in_standard_logging_payload(): print("✅ Cost breakdown test passed!") + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_cost_breakdown_missing_in_standard_logging_payload(): """ @@ -12535,6 +12592,7 @@ def test_cost_breakdown_missing_in_standard_logging_payload(): print("✅ Cost breakdown missing test passed!") + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "use_combined_usage_object", @@ -12622,6 +12680,7 @@ def test_usage_dict_roundtrip_in_payload(use_combined_usage_object): assert usage_obj["completion_tokens"] == 58 assert usage_obj["total_tokens"] == 100 + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_actual_model_for_azure_router(): from litellm.litellm_core_utils.litellm_logging import ( @@ -12670,6 +12729,7 @@ def test_standard_logging_payload_uses_actual_model_for_azure_router(): assert payload is not None assert payload["model"] == "azure_ai/gpt-5-nano-2025-08-07" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_standard_logging_payload_uses_actual_model_for_azure_router_with_underscore(): from litellm.litellm_core_utils.litellm_logging import ( @@ -12718,6 +12778,7 @@ def test_standard_logging_payload_uses_actual_model_for_azure_router_with_unders assert payload is not None assert payload["model"] == "azure_ai/gpt-5-nano-2025-08-07" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_basic(): """ @@ -12749,6 +12810,7 @@ def test_merge_litellm_metadata_basic(): assert result["model_info"] == {"id": "model-123"} assert result["tags"] == ["tag1", "tag2"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_precedence(): """ @@ -12775,6 +12837,7 @@ def test_merge_litellm_metadata_precedence(): # litellm_metadata values should only be included if not in metadata assert result["model_group"] == "gpt-4-group" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_skip_non_serializable(): """ @@ -12807,6 +12870,7 @@ def test_merge_litellm_metadata_skip_non_serializable(): assert result["safe_field"] == "safe_value" assert result["model_group"] == "gpt-4-group" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_empty_params(): """ @@ -12842,6 +12906,7 @@ def test_merge_litellm_metadata_empty_params(): result = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) assert result == {} + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_merge_litellm_metadata_bedrock_passthrough_scenario(): """ @@ -12901,8 +12966,10 @@ def test_merge_litellm_metadata_bedrock_passthrough_scenario(): # Verify total number of fields (9 user fields + 4 model fields = 13) assert len(result) == 13 + service_logger = ServiceLogging() + def setup_logging(): return Logging( model="gpt-5.5", @@ -12992,6 +13059,7 @@ def test_get_callback_name(): # Test string callback assert logging._get_callback_name("callback_string") == "callback_string" + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_is_internal_litellm_proxy_callback(): """ @@ -13012,6 +13080,7 @@ def test_is_internal_litellm_proxy_callback(): # Test string callback assert logging._is_internal_litellm_proxy_callback("callback_string") == False + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_should_run_sync_callbacks_for_async_calls(): """ @@ -13037,6 +13106,7 @@ def test_should_run_sync_callbacks_for_async_calls(): litellm.success_callback = [PROXY_MaxIterationsHandler] assert logging._should_run_sync_callbacks_for_async_calls() == False + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_remove_internal_litellm_callbacks(): logging = setup_logging() @@ -13058,6 +13128,7 @@ def test_remove_internal_litellm_callbacks(): assert PROXY_MaxIterationsHandler not in filtered assert PROXY_CacheControlCheck not in filtered + @pytest.mark.asyncio async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch): """ @@ -13160,14 +13231,9 @@ class CompletionCustomHandler( assert metadata_value is None or isinstance(metadata_value, dict) if metadata_value is not None: if litellm.turn_off_message_logging is True: - assert ( - metadata_value["raw_request"] - is LiteLLMCommonStrings.redacted_by_litellm.value - ) + assert metadata_value["raw_request"] is LiteLLMCommonStrings.redacted_by_litellm.value else: - assert "raw_request" not in metadata_value or isinstance( - metadata_value["raw_request"], str - ) + assert "raw_request" not in metadata_value or isinstance(metadata_value["raw_request"], str) except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -13216,23 +13282,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -13265,9 +13326,7 @@ class CompletionCustomHandler( ) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["litellm_params"]["api_base"], str) @@ -13277,10 +13336,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["user"], (str, type(None))) assert ( isinstance(kwargs["input"], list) - and ( - isinstance(kwargs["input"][0], dict) - or isinstance(kwargs["input"][0], str) - ) + and (isinstance(kwargs["input"][0], dict) or isinstance(kwargs["input"][0], str)) ) or isinstance(kwargs["input"], (dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert isinstance( @@ -13308,9 +13364,7 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) @@ -13318,15 +13372,12 @@ class CompletionCustomHandler( assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -13344,9 +13395,7 @@ class CompletionCustomHandler( assert isinstance(messages, list) and isinstance(messages[0], dict) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) @@ -13358,9 +13407,7 @@ class CompletionCustomHandler( async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - print( - "in async_log_success_event", kwargs, response_obj, start_time, end_time - ) + print("in async_log_success_event", kwargs, response_obj, start_time, end_time) self.states.append("async_success") ## START TIME assert isinstance(start_time, datetime) @@ -13389,9 +13436,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -13423,9 +13468,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -13566,3 +13609,36 @@ def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_a assert s3_writer.log_event.call_count == expected_s3_writes if expected_s3_writes: assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response + + +_GO_PRICED_MESSAGE: Final = { + "id": "msg_go", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1200, "output_tokens": 300, "cost": 0, "opencode": {"list_cost": 0.00412}}, +} + + +@pytest.mark.parametrize( + ("usage", "expected_header"), + ( + ({"input_tokens": 1200, "output_tokens": 300, "cost": 0, "opencode": {"list_cost": 0.00412}}, 0.00412), + ({"input_tokens": 1200, "output_tokens": 300}, None), + ), + ids=("list-cost-reported", "no-list-cost-keeps-the-price-map"), +) +def test_native_messages_body_carries_the_reported_list_cost_into_the_priced_response(usage, expected_header): + logging_obj = _anthropic_messages_logging_obj() + logging_obj.stream = False + logging_obj.optional_params = {} + logging_obj.model_call_details["custom_llm_provider"] = "opencode_go" + + priced = logging_obj.billable_result({**_GO_PRICED_MESSAGE, "usage": usage}) + + assert isinstance(priced, ModelResponse) + headers = priced._hidden_params.get("additional_headers") or {} + assert headers.get("llm_provider-x-litellm-response-cost") == expected_header diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py index fbbbc579d94..ec57636a17b 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_sse_stream.py @@ -8,8 +8,11 @@ wrapper's chunks/messages/model so disconnect-time partial billing can read them from typing import Final from unittest.mock import MagicMock +import httpx import pytest +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicSSEStream, AnthropicStreamWrapper, @@ -85,3 +88,67 @@ def test_sse_stream_exposes_chunks_messages_and_model(): assert stream.messages == [{"role": "user", "content": "hi"}] chunks: Final = stream.chunks assert isinstance(chunks, list) and len(chunks) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("started", [False, True]) +async def test_sse_stream_close_releases_http_response_before_upstream_body(started: bool) -> None: + response: Final = httpx.Response(200, stream=httpx.ByteStream(b"unread upstream body")) + upstream: Final = CustomStreamWrapper( + completion_stream=response, + model="test-model", + custom_llm_provider="openai", + logging_obj=Logging( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=1.0, + litellm_call_id="disconnect-test", + function_id="test", + ), + ) + stream: Final = AnthropicSSEStream(AnthropicStreamWrapper(completion_stream=upstream, model="test-model")) + if started: + assert (await anext(stream)).startswith(b"event: message_start") + assert not response.is_stream_consumed + assert not response.is_closed + + await stream.aclose() + await stream.aclose() + + assert response.is_closed + assert not response.is_stream_consumed + with pytest.raises(StopAsyncIteration): + await anext(stream) + + +class _CloseableAsyncStream(_AsyncStream): + def __init__(self, items: list[MagicMock]) -> None: + super().__init__(items) + self.close_calls = 0 + + async def aclose(self) -> None: + self.close_calls += 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("finish_stream", [False, True]) +async def test_sse_stream_releases_upstream_once_after_content(finish_stream: bool) -> None: + upstream: Final = _CloseableAsyncStream( + [_make_chunk(Delta(content="hello")), _make_chunk(Delta(), finish_reason="stop")] + ) + stream: Final = AnthropicSSEStream(AnthropicStreamWrapper(completion_stream=upstream, model="test-model")) + events: Final[list[bytes]] = [] + async for event in stream: + events.append(event) + if not finish_stream and b"event: content_block_delta" in event: + break + assert any(b'"text":"hello"' in event.replace(b" ", b"") for event in events) + if finish_stream: + assert events[-1].startswith(b"event: message_stop") + + await stream.aclose() + await stream.aclose() + + assert upstream.close_calls == 1 diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index dd99ad8bf91..3fbb12881da 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -511,6 +511,54 @@ class TestDeepSeekVisionMultimodalContent: assert any(block.get("type") == "image_url" for block in content) +@pytest.mark.usefixtures("_deepseek_httpx_transport") +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("model", ("deepseek-flash", "deepseek-v4-pro", "deepseek-chat", "deepseek-reasoner")) +@pytest.mark.parametrize("effort", ("minimal", "low", "medium", "high", "xhigh", "max", "ultra")) +@pytest.mark.parametrize("thinking", (None, "enabled")) +@pytest.mark.respx(assert_all_called=True) +async def test_deepseek_forwards_native_effort( + respx_mock: respx.MockRouter, model: str, effort: str, thinking: str | None, stream: bool +) -> None: + route: Final = respx_mock.post(DEEPSEEK_CHAT_COMPLETIONS_URL).mock(return_value=_deepseek_response(stream)) + response: Final = await litellm.acompletion( + model=f"deepseek/{model}", + messages=[{"role": "user", "content": "Hello"}], + api_key=DEEPSEEK_API_KEY, + reasoning_effort=effort, + stream=stream, + **({"thinking": {"type": thinking}} if thinking is not None else {}), + ) + if stream: + async for _ in response: + pass + body: Final = json.loads(route.calls[0].request.content) + assert body["reasoning_effort"] == effort + assert body["thinking"] == {"type": "enabled"} + + +@pytest.mark.parametrize( + "thinking,effort,expected_thinking", + ( + (None, None, None), + (None, "none", "disabled"), + ("enabled", "none", "enabled"), + ("disabled", "max", "disabled"), + ("enabled", None, "enabled"), + (None, "default", "enabled"), + (None, "bogus", "enabled"), + (None, {"effort": "high"}, "enabled"), + ), +) +def test_deepseek_preserves_toggle_without_forwarding_unsupported_effort( + thinking: str | None, effort: str | dict[str, str] | None, expected_thinking: str | None +) -> None: + params: Final = {"reasoning_effort": effort, **({"thinking": {"type": thinking}} if thinking is not None else {})} + result: Final = DeepSeekChatConfig().map_openai_params(params, {}, "deepseek-flash", False) + assert "reasoning_effort" not in result + assert result.get("thinking") == ({"type": expected_thinking} if expected_thinking else None) + + class TestDeepSeekThinkingParams: """Test thinking and reasoning_effort parameter handling for DeepSeek.""" @@ -735,9 +783,7 @@ async def test_deepseek_provider_async_completion(stream): messages = [{"role": "user", "content": "Hello, world!"}] # Mock AsyncHTTPHandler.post method for async test - with patch( - "litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post") as mock_post: mock_response_data = litellm.ModelResponse( choices=[ litellm.Choices( @@ -775,9 +821,7 @@ async def test_deepseek_provider_async_completion(stream): # Check request body request_body = json.loads(call_args.kwargs["data"]) assert call_args.kwargs["url"] == "https://api.deepseek.com/beta/chat/completions" - assert ( - request_body["model"] == "deepseek-reasoner" - ) # Model name should be stripped of provider prefix + assert request_body["model"] == "deepseek-reasoner" # Model name should be stripped of provider prefix assert request_body["messages"] == messages assert request_body["stream"] == stream diff --git a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py index 95611e344dc..7aa8c65215e 100644 --- a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -1,4 +1,5 @@ import logging +from typing import Final from unittest.mock import patch import httpx @@ -523,3 +524,68 @@ def test_tool_search_reply_sent_to_groq_drops_litellm_response_fields(): } ], } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("with_tool_call", [False, True]) +async def test_streamed_reasoning_reply_can_be_replayed(is_async: bool, with_tool_call: bool) -> None: + tool_calls: Final = [ + { + "index": 0, + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ] + handler: Final = GroqChatCompletionStreamingHandler(streaming_response=iter(()), sync_stream=True) + chunk: Final = handler.chunk_parser( + { + "id": "chatcmpl-reasoning", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "openai/gpt-oss-20b", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "reasoning": "I should check the weather.", + "content": "" if with_tool_call else "Hello there!", + **({"tool_calls": tool_calls} if with_tool_call else {}), + }, + "finish_reason": "tool_calls" if with_tool_call else "stop", + } + ], + } + ) + reply: Final = litellm.stream_chunk_builder([chunk]).choices[0].message + assert reply.reasoning_content == "I should check the weather." + history: Final = [reply.model_dump(exclude_none=True)] + config: Final = GroqChatConfig() + + if is_async: + replayed = await config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=True) + else: + replayed = config._transform_messages(messages=history, model="openai/gpt-oss-20b", is_async=False) + + assert replayed == [ + { + "role": "assistant", + "content": "" if with_tool_call else "Hello there!", + **( + { + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ] + } + if with_tool_call + else {} + ), + } + ] + assert reply.reasoning_content == "I should check the weather." diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py index ce8f46f3da0..0a5e49ca1b2 100644 --- a/tests/unit/llms/opencode/chat/test_transformation.py +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -1,3 +1,4 @@ +import asyncio import json from collections.abc import Mapping from datetime import datetime @@ -6,8 +7,11 @@ from typing import Final import httpx import pytest import respx +from pydantic import TypeAdapter import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.opencode.common_utils import go_headers @@ -400,3 +404,179 @@ def test_gemini_sync_stream_preserves_tool_calls_and_usage(go_catalog: None) -> assert json.loads(tool.function.arguments) == {"city": "Paris"} assert any(c.choices[0].finish_reason == "tool_calls" for c in chunks if c.choices) assert chunks[-1].usage.total_tokens == 15 + + +GO_REPORTED_USAGE: Final = { + "cost": 0, + "opencode": {"billing_source": "go", "list_cost": 0.00412, "allowance": 0.00412}, +} +ROW_PRICE_FOR_FIXTURE_USAGE: Final = 10 * 0.000002 + 5 * 0.000005 + + +def go_priced_reply(route: str, model: str, usage_extra: Mapping[str, object]) -> Mapping[str, object]: + reply: Final = upstream_reply(route, model) + key: Final = "usageMetadata" if route == "gemini" else "usage" + usage: Final = TypeAdapter(Mapping[str, object]).validate_python(reply[key]) + return {**reply, key: {**usage, **usage_extra}} + + +class CostRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.costs: list[float | None] = [] # mutable-ok: test probe accumulates logged costs + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + self.costs.append(kwargs.get("response_cost")) + + +async def logged_cost(recorder: CostRecorder) -> float | None: + for _ in range(50): + if recorder.costs: + return recorder.costs[-1] + await asyncio.sleep(0.05) + raise AssertionError("success callback never ran") + + +@pytest.fixture +def cost_recorder(monkeypatch: pytest.MonkeyPatch) -> CostRecorder: + recorder: Final = CostRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + return recorder + + +@pytest.mark.asyncio +@pytest.mark.parametrize("upstream", ["chat/completions", "messages", "responses", "gemini"]) +@pytest.mark.parametrize("ingress", ["chat", "responses"]) +async def test_reported_list_cost_is_billed_instead_of_rows( + go_catalog: None, cost_recorder: CostRecorder, upstream: str, ingress: str +) -> None: + model: Final = f"fixture-{upstream.split('/')[0]}" + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response(200, json=go_priced_reply(upstream, model, GO_REPORTED_USAGE)) + ) + kwargs: Final = {"model": f"opencode_go/{model}", "api_key": "fake-go-key"} + if ingress == "chat": + response: Final = await litellm.acompletion(messages=[{"role": "user", "content": "hello"}], **kwargs) + else: + response = await litellm.aresponses(input="hello", **kwargs) + assert response._hidden_params["additional_headers"][RESPONSE_COST_HEADER] == pytest.approx(0.00412) + assert await logged_cost(cost_recorder) == pytest.approx(0.00412) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("usage_extra", [{}, {"cost": 0}, {"opencode": {"billing_source": "go"}}]) +async def test_rows_price_the_request_when_no_list_cost_is_reported( + go_catalog: None, cost_recorder: CostRecorder, usage_extra: Mapping[str, object] +) -> None: + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response(200, json=go_priced_reply("chat/completions", "fixture-chat", usage_extra)) + ) + response: Final = await litellm.acompletion( + model="opencode_go/fixture-chat", api_key="fake-go-key", messages=[{"role": "user", "content": "hello"}] + ) + assert RESPONSE_COST_HEADER not in response._hidden_params.get("additional_headers", {}) + assert await logged_cost(cost_recorder) == pytest.approx(ROW_PRICE_FOR_FIXTURE_USAGE) + + +def sse(*events: Mapping[str, object]) -> str: + return "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_chat_stream_bills_the_final_usage_chunks_list_cost( + go_catalog: None, cost_recorder: CostRecorder +) -> None: + chunk: Final = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "fixture-chat", + } + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=sse( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hello"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + **chunk, + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, **GO_REPORTED_USAGE}, + }, + ), + ) + ) + stream: Final = await litellm.acompletion( + model="opencode_go/fixture-chat", + api_key="fake-go-key", + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [piece async for piece in stream] + assert "".join(piece.choices[0].delta.content or "" for piece in chunks if piece.choices) == "hello" + assert await logged_cost(cost_recorder) == pytest.approx(0.00412) + + +@pytest.mark.asyncio +async def test_responses_stream_bills_the_completed_events_list_cost( + go_catalog: None, cost_recorder: CostRecorder +) -> None: + completed: Final = go_priced_reply("responses", "fixture-responses", GO_REPORTED_USAGE) + with respx.mock(using="httpx") as transport: + transport.post().mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=sse( + { + "type": "response.output_text.delta", + "delta": "hello", + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + }, + {"type": "response.completed", "response": completed, "sequence_number": 2}, + ), + ) + ) + stream: Final = await litellm.aresponses( + model="opencode_go/fixture-responses", api_key="fake-go-key", input="hello", stream=True + ) + events: Final = [event async for event in stream] + assert events[-1].type == "response.completed" + assert events[-1].response._hidden_params["additional_headers"][RESPONSE_COST_HEADER] == pytest.approx(0.00412) + assert await logged_cost(cost_recorder) == pytest.approx(0.00412) + + +def test_stream_chunk_builder_reads_list_cost_from_dict_chunks(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="opencode_go/glm-5.3", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="dict-chunks", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "opencode_go" + chunk = { + "id": "chatcmpl-go", + "object": "chat.completion.chunk", + "created": 1, + "model": "glm-5.3", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + } + final_usage = {**chunk, "choices": [], "usage": dict(GO_REPORTED_USAGE)} + + assembled = litellm.stream_chunk_builder(chunks=[chunk, final_usage], logging_obj=logging_obj) + + assert assembled is not None + assert assembled._hidden_params["response_cost"] == GO_REPORTED_USAGE["opencode"]["list_cost"] diff --git a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py index fb676ae9a61..2b682fda982 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -588,6 +588,6 @@ def test_registered_inference_routes_have_an_explicit_managed_access_decision(ro "/vector_stores", "/vector_store/", "/search", "/containers", "/skills", "/claude-code/", "/interactions", "/agents", "/responses/{", "/responses/input_tokens", "/realtime/client_secrets", "/realtime/calls", "/realtime/transcription_sessions", - )) or normalized in ("/models", "/cursor/models", "/cursor/v1/models") + )) or normalized in ("/models", "/cursor/models", "/cursor/v1/models", "/claude_code_gateway/v1/models") concrete: Final = route.split("?")[0].replace("{model}", "model").replace("{model_name:path}", "model") assert managed_agent_route_allowed(concrete, None) is not unsupported, route diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py index 8ac0e2a6f6d..8c2d8e4f157 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -813,3 +813,95 @@ def test_messages_gated_by_enable_flag(): with _gateway_env(enabled=False) as (client, _): resp = client.post("/claude_code_gateway/v1/messages", json={"model": "claude-sonnet-4-5", "messages": []}) assert resp.status_code == 404 + + +@pytest.mark.parametrize("headers", [{}, {"anthropic-version": "2023-06-01"}]) +def test_gateway_model_discovery_preserves_grants_and_format( + monkeypatch: pytest.MonkeyPatch, headers: dict[str, str] +) -> None: + from litellm import Router + from litellm.proxy import proxy_server + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + model_router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"openai/{name}", "api_key": "local-placeholder"}, + } + for name in ("discovery-example", "restricted-example") + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", model_router) + monkeypatch.setattr(proxy_server, "llm_model_list", model_router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_claude_code_gateway": True}) + caller: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + models=["discovery-example"], + ) + app: Final = FastAPI() + app.include_router(proxy_server.router) + app.include_router(gateway_endpoints.router) + app.dependency_overrides[gateway_endpoints.user_api_key_auth] = lambda: caller + with TestClient(app) as client: + standard: Final = client.get("/v1/models", headers=headers) + gateway: Final = client.get("/claude_code_gateway/v1/models", headers=headers) + assert standard.status_code == 200, standard.text + assert gateway.status_code == 200, gateway.text + assert gateway.json() == standard.json() + assert [model["id"] for model in gateway.json()["data"]] == ["discovery-example"] + invalid_scope: Final = client.get("/claude_code_gateway/v1/models?scope=invalid", headers=headers) + assert invalid_scope.status_code == 400 + + +def test_gateway_model_discovery_requires_authentication() -> None: + with _gateway_env(real_auth=True) as (client, _), pytest.raises(ProxyException) as exc_info: + client.get("/claude_code_gateway/v1/models") + assert exc_info.value.code == "401" + + +@pytest.mark.parametrize("settings", [{}, {"enable_claude_code_gateway": False}]) +def test_gateway_model_discovery_disabled(monkeypatch: pytest.MonkeyPatch, settings: dict[str, bool]) -> None: + from litellm.proxy import proxy_server + + with _gateway_env(real_auth=True) as (client, _): + monkeypatch.setattr(proxy_server, "general_settings", settings) + response: Final = client.get( + "/claude_code_gateway/v1/models", headers={"Authorization": f"Bearer {_MASTER_KEY}"} + ) + assert response.status_code == 404, response.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/models", "/v1/models", "/claude_code_gateway/v1/models"]) +async def test_gateway_model_discovery_preserves_route_exemptions(monkeypatch: pytest.MonkeyPatch, route: str) -> None: + from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + + import litellm + from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.auth.route_checks import RouteChecks + + monkeypatch.setattr(litellm, "max_budget", 1.0) + result: Final = await common_checks( + request_body={}, + team_object=LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0), + user_object=None, + end_user_object=None, + global_proxy_spend=2.0, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=UserAPIKeyAuth(team_id="test-team"), + request=MagicMock(), + ) + assert result is True + assert RouteChecks.is_unbilled_route(route) + assert RouteChecks.is_llm_api_route(route) + monkeypatch.setenv("DISABLE_LLM_API_ENDPOINTS", "true") + monkeypatch.setenv("DISABLE_ADMIN_ENDPOINTS", "true") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + EnterpriseRouteChecks.should_call_route(route) diff --git a/tests/unit/proxy/auth/test_team_member_budget.py b/tests/unit/proxy/auth/test_team_member_budget.py index b38a953d189..6055e4f0356 100644 --- a/tests/unit/proxy/auth/test_team_member_budget.py +++ b/tests/unit/proxy/auth/test_team_member_budget.py @@ -3,8 +3,11 @@ Unit tests for team member budget checks in common_checks. These tests verify the team member budget enforcement without requiring a proxy server. """ -import pytest +from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch + +import pytest from fastapi import Request import litellm @@ -15,7 +18,10 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks import common_checks, get_team_membership +from litellm.proxy.auth.auth_checks import ( + common_checks, + resolve_effective_team_member_budget, +) @pytest.mark.asyncio @@ -436,3 +442,48 @@ async def test_team_member_budget_check_personal_key_not_team(): # Should pass and get_team_membership should not be called assert result is True mock_get_team_membership.assert_not_called() + + +def test_resolve_effective_team_member_budget_preserves_member_and_default_precedence() -> None: + now: Final = datetime(2030, 1, 1, tzinfo=timezone.utc) + member_budget: Final = LiteLLM_BudgetTable( + max_budget=100, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + expired_member_budget: Final = LiteLLM_BudgetTable( + max_budget=100, + temp_budget_increase=10, + temp_budget_expiry=datetime(2029, 12, 31, tzinfo=timezone.utc), + ) + active_inherited_increase: Final = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + expired_inherited_increase: Final = LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=10, + temp_budget_expiry=datetime(2029, 12, 31, tzinfo=timezone.utc), + ) + + assert resolve_effective_team_member_budget(member_budget, None, now) == 110 + assert resolve_effective_team_member_budget(expired_member_budget, None, now) == 100 + assert resolve_effective_team_member_budget( + active_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 60 + assert resolve_effective_team_member_budget( + expired_inherited_increase, + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 50 + assert resolve_effective_team_member_budget( + LiteLLM_BudgetTable(max_budget=0), + LiteLLM_BudgetTable(max_budget=50), + now, + ) == 0 + assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=50), now) == 50 + assert resolve_effective_team_member_budget(None, LiteLLM_BudgetTable(max_budget=0), now) is None + assert resolve_effective_team_member_budget(None, None, now) is None diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 70f7774e691..183dd7919c1 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -113,6 +113,9 @@ class MockBatcher: self.litellm_projecttable = _Table("project", self) self.litellm_endusertable = _Table("enduser", self) + def execute_raw(self, query: str, *args: object) -> None: + self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args}) + async def commit(self): self.committed = True return self.calls @@ -2112,6 +2115,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")} assert batcher.committed is True assert {(call["table"], call["op"]) for call in batcher.calls} == { + ("raw", "execute_raw"), ("team_membership", "update_many"), ("key", "update_many"), ("org", "update_many"), @@ -3393,6 +3397,35 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled( counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1") +def test_budget_cascade_floors_membership_total_spend_before_zeroing_spend( + rollover_enabled: None, + reset_budget_job: ResetBudgetJob, + mock_prisma_client: MockPrismaClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Membership rows written before total_spend existed hold their history only in + spend. The cascade lifts total_spend to spend for the expiring tier's rows, in the + same batch and ahead of the zero/decrement, so the reset cannot erase it.""" + _make_counter_invalidation_job(monkeypatch) + budget: Final = _budget_row(budget_id="budget-legacy", budget_duration="7d", max_budget=10.0) + mock_prisma_client.data["budget"] = [budget] + membership: Final = type( + "Membership", + (), + {"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-legacy"}, + ) + mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership]) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + writes: Final = mock_prisma_client.db.batch_calls + repair_index: Final = next(i for i, call in enumerate(writes) if call["op"] == "execute_raw") + first_membership_reset: Final = next(i for i, call in enumerate(writes) if call["table"] == "team_membership") + assert repair_index < first_membership_reset, writes + assert writes[repair_index]["args"] == (["budget-legacy"],) + assert mock_prisma_client.db.batchers[-1].committed + + def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch ): diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index ac235902925..531e7fb787c 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1128,7 +1128,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement - assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement + assert ( + 'total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend)' + " + EXCLUDED.total_spend" + ) in statement @pytest.mark.asyncio diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py new file mode 100644 index 00000000000..9a7670dbe97 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py @@ -0,0 +1,285 @@ +""" +Tests for the Reco guardrail integration. +""" + +from collections.abc import Mapping +from typing import Final +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.reco import ( + RecoGuardrail, + initialize_guardrail, +) +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.proxy.guardrails.guardrail_hooks.reco import ( + RecoConfigModel, + RecoOptionalParams, +) +from litellm.types.utils import GenericGuardrailAPIInputs + + +@pytest.fixture +def reco_guardrail() -> RecoGuardrail: + return RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + guardrail_name="my-reco-guardrail", + event_hook="pre_call", + default_on=True, + ) + + +class TestRecoGuardrailConfiguration: + def test_tenant_id_forwarded_as_header(self, reco_guardrail: RecoGuardrail) -> None: + assert reco_guardrail.headers == {"X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111"} + + def test_tenant_header_merged_with_existing_headers(self) -> None: + guardrail: Final = RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + headers={"X-Custom": "value"}, + ) + assert guardrail.headers == { + "X-Custom": "value", + "X-Reco-Tenant-Id": "11111111-1111-1111-1111-111111111111", + } + + def test_missing_reco_tenant_id_raises(self) -> None: + with pytest.raises(ValueError, match="reco_tenant_id"): + RecoGuardrail(reco_tenant_id=None, api_base="https://edge1.us.reco.ai") + + def test_missing_api_base_raises(self) -> None: + with pytest.raises(ValueError, match="api_base"): + RecoGuardrail(reco_tenant_id="11111111-1111-1111-1111-111111111111", api_base=None) + + def test_non_uuid_reco_tenant_id_raises(self) -> None: + with pytest.raises(ValueError, match="UUID"): + RecoGuardrail(reco_tenant_id="tenant-123", api_base="https://edge1.us.reco.ai") + + def test_optional_params_rejects_non_uuid_tenant_id(self) -> None: + with pytest.raises(ValidationError, match="UUID"): + RecoOptionalParams(reco_tenant_id="tenant-123", api_base="https://edge1.us.reco.ai") + + def test_get_supported_event_hooks_returns_pre_call_only(self) -> None: + assert RecoGuardrail.get_supported_event_hooks() == [GuardrailEventHooks.pre_call] + + def test_get_config_model_returns_reco_config_model(self) -> None: + assert RecoGuardrail.get_config_model() is RecoConfigModel + + def test_config_model_ui_friendly_name(self) -> None: + assert RecoConfigModel.ui_friendly_name() == "Reco" + + def test_litellm_params_stay_assignable_for_registry_secret_resolution(self) -> None: + litellm_params: Final = LitellmParams(guardrail="reco", mode="pre_call") + + litellm_params.api_base = "https://edge1.us.reco.ai" + + assert litellm_params.api_base == "https://edge1.us.reco.ai" + + +class TestRecoGuardrailInitializer: + def test_initialize_guardrail_wires_optional_params(self) -> None: + litellm_params: Final = LitellmParams( + guardrail="reco", + mode="pre_call", + optional_params={ + "reco_tenant_id": "22222222-2222-2222-2222-222222222222", + "api_base": "https://edge2.eu.reco.ai", + }, + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + + assert isinstance(guardrail, RecoGuardrail) + assert guardrail.guardrail_name == "prod-reco" + assert guardrail.headers == {"X-Reco-Tenant-Id": "22222222-2222-2222-2222-222222222222"} + assert guardrail.api_base == "https://edge2.eu.reco.ai/beta/litellm_basic_guardrail_api" + + def test_initialize_guardrail_reads_plain_dict_optional_params(self) -> None: + """Guardrails reconciled from the DB can carry `optional_params` as a + plain dict rather than a pydantic object; attribute-only access would + silently read None for every field.""" + litellm_params: Final = LitellmParams.model_construct( + guardrail="reco", + optional_params={ + "reco_tenant_id": "33333333-3333-3333-3333-333333333333", + "api_base": "https://edge3.eu.reco.ai", + }, + mode="pre_call", + default_on=True, + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + + assert guardrail.headers == {"X-Reco-Tenant-Id": "33333333-3333-3333-3333-333333333333"} + assert guardrail.api_base == "https://edge3.eu.reco.ai/beta/litellm_basic_guardrail_api" + + def test_initialize_guardrail_resolves_env_var_reference(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("RECO_TEST_TENANT_ID", "44444444-4444-4444-4444-444444444444") + litellm_params: Final = LitellmParams( + guardrail="reco", + optional_params={ + "reco_tenant_id": "os.environ/RECO_TEST_TENANT_ID", + "api_base": "https://edge4.eu.reco.ai", + }, + mode="pre_call", + default_on=True, + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + + assert guardrail.headers == {"X-Reco-Tenant-Id": "44444444-4444-4444-4444-444444444444"} + + +def _litellm_params(top_level: Mapping[str, object], nested: Mapping[str, object]) -> LitellmParams: + return LitellmParams( + guardrail="reco", + mode="pre_call", + optional_params={ + "reco_tenant_id": "22222222-2222-2222-2222-222222222222", + "api_base": "https://edge2.eu.reco.ai", + **nested, + }, + **top_level, + ) + + +async def _apply_with_reco_unreachable(guardrail: RecoGuardrail) -> GenericGuardrailAPIInputs: + with patch.object(guardrail.async_handler, "post", side_effect=httpx.ConnectError("connection refused")): + return await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + +class TestRecoGuardrailFailureModes: + @pytest.mark.asyncio + async def test_unreachable_reco_blocks_request_by_default(self, reco_guardrail: RecoGuardrail) -> None: + with pytest.raises(Exception, match="connection refused"): + await _apply_with_reco_unreachable(reco_guardrail) + + @pytest.mark.asyncio + async def test_fail_open_lets_request_through_when_reco_unreachable(self) -> None: + guardrail: Final = RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + fail_on_error=False, + ) + + result: Final = await _apply_with_reco_unreachable(guardrail) + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_fail_closed_blocks_request_when_reco_unreachable(self) -> None: + guardrail: Final = RecoGuardrail( + reco_tenant_id="11111111-1111-1111-1111-111111111111", + api_base="https://edge1.us.reco.ai", + unreachable_fallback="fail_closed", + fail_on_error=True, + ) + + with pytest.raises(Exception, match="connection refused"): + await _apply_with_reco_unreachable(guardrail) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("top_level", "nested"), + [ + pytest.param({}, {}, id="defaults"), + pytest.param({}, {"unreachable_fallback": "fail_closed", "fail_on_error": True}, id="optional-params"), + pytest.param({"unreachable_fallback": "fail_closed", "fail_on_error": True}, {}, id="top-level"), + ], + ) + async def test_initialize_guardrail_fails_closed_when_configured( + self, top_level: Mapping[str, object], nested: Mapping[str, object] + ) -> None: + guardrail: Final = initialize_guardrail(_litellm_params(top_level, nested), {"guardrail_name": "prod-reco"}) + + with pytest.raises(Exception, match="connection refused"): + await _apply_with_reco_unreachable(guardrail) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("top_level", "nested"), + [ + pytest.param({}, {"fail_on_error": False}, id="optional-params"), + pytest.param({"fail_on_error": False}, {}, id="top-level"), + pytest.param( + {"unreachable_fallback": "fail_closed", "fail_on_error": True}, + {"fail_on_error": False}, + id="optional-params-wins-over-top-level", + ), + ], + ) + async def test_initialize_guardrail_fails_open_when_configured( + self, top_level: Mapping[str, object], nested: Mapping[str, object] + ) -> None: + guardrail: Final = initialize_guardrail(_litellm_params(top_level, nested), {"guardrail_name": "prod-reco"}) + + result: Final = await _apply_with_reco_unreachable(guardrail) + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_initialize_guardrail_keeps_top_level_fail_open_when_nested_params_are_a_model(self) -> None: + nested: Final = RecoOptionalParams( + reco_tenant_id="22222222-2222-2222-2222-222222222222", api_base="https://edge2.eu.reco.ai" + ) + litellm_params: Final = LitellmParams.model_construct( + guardrail="reco", mode="pre_call", optional_params=nested, fail_on_error=False + ) + + guardrail: Final = initialize_guardrail(litellm_params, {"guardrail_name": "prod-reco"}) + result: Final = await _apply_with_reco_unreachable(guardrail) + + assert result["texts"] == ["hello"] + + +class TestRecoGuardrailBlocking: + @pytest.mark.asyncio + async def test_action_blocked_raises_exception_with_configured_guardrail_name( + self, reco_guardrail: RecoGuardrail + ) -> None: + mock_response: Final = MagicMock() + mock_response.json.return_value = { + "action": "BLOCKED", + "blocked_reason": "Sensitive data detected", + } + mock_response.raise_for_status = MagicMock() + + with patch.object(reco_guardrail.async_handler, "post", return_value=mock_response) as mock_post: + with pytest.raises(GuardrailRaisedException) as exc_info: + await reco_guardrail.apply_guardrail( + inputs={"texts": ["some sensitive prompt"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + assert exc_info.value.guardrail_name == "my-reco-guardrail" + assert str(exc_info.value) == "Sensitive data detected" + assert exc_info.value.blocked_content is True + + _, call_kwargs = mock_post.call_args + assert call_kwargs["headers"]["X-Reco-Tenant-Id"] == "11111111-1111-1111-1111-111111111111" + + @pytest.mark.asyncio + async def test_action_none_allows_content(self, reco_guardrail: RecoGuardrail) -> None: + mock_response: Final = MagicMock() + mock_response.json.return_value = {"action": "NONE", "texts": ["hello"]} + mock_response.raise_for_status = MagicMock() + + with patch.object(reco_guardrail.async_handler, "post", return_value=mock_response): + result: Final = await reco_guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + assert result["texts"] == ["hello"] diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 76ad1e08e28..79193d37b04 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,35 +1,42 @@ import asyncio import json import logging +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch import httpx import pytest +from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY -from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth +from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( - _get_budget_reservation_from_metadata, ProxyDBLogger, + _get_budget_reservation_from_metadata, _should_track_cost_callback, _update_database_and_spend_counters, + is_pre_routing_rejection, run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError -from litellm.proxy.utils import ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import ( + REJECTED_BEFORE_REQUEST_SETUP, + REJECTED_BEFORE_REQUEST_SETUP_KEY, + ProxyUpdateSpend, +) from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage @@ -3559,3 +3566,103 @@ async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(m assert sum(counter.totals.values()) == 5 assert outcome.failed is True assert outcome.retryable_failure is False + + +_ROUTED_REQUEST: Final[Mapping[str, object]] = { + "model": "anthropic-haiku-4-5", + "messages": [{"role": "user", "content": "hi"}], + "proxy_server_request": {"request_id": "req-1"}, + "litellm_call_id": "call-1", +} +_PRE_ROUTING_REQUEST: Final[Mapping[str, object]] = { + **_ROUTED_REQUEST, + REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP, +} +_SPOOFED_REQUEST: Final[Mapping[str, object]] = {**_ROUTED_REQUEST, REJECTED_BEFORE_REQUEST_SETUP_KEY: True} +_INVALID_REQUEST_400: Final = ProxyException( + message="Violated guardrail policy", type="invalid_request_error", param=None, code=400 +) +_PROVIDER_AUTH_ERROR: Final = litellm.AuthenticationError( + message="bad key", llm_provider="anthropic", model="claude-haiku-4-5" +) + + +@pytest.mark.parametrize( + ("exception", "request_data", "expected"), + [ + (_INVALID_REQUEST_400, _PRE_ROUTING_REQUEST, True), + (litellm.BudgetExceededError(current_cost=1.0, max_budget=0.0), _PRE_ROUTING_REQUEST, True), + (ProxyModelNotFoundError(route="/chat/completions", model_name="4.6"), _ROUTED_REQUEST, True), + (_INVALID_REQUEST_400, _ROUTED_REQUEST, False), + (_INVALID_REQUEST_400, _SPOOFED_REQUEST, False), + (_PROVIDER_AUTH_ERROR, _ROUTED_REQUEST, False), + ( + litellm.RateLimitError(message="slow down", llm_provider="anthropic", model="claude-haiku-4-5"), + _ROUTED_REQUEST, + False, + ), + (HTTPException(status_code=429, detail="Max parallel request limit reached"), _ROUTED_REQUEST, False), + ], + ids=[ + "bad-body", + "budget", + "unknown-model", + "guardrail-block", + "guardrail-block-with-client-spoofed-marker", + "provider-auth", + "provider-429", + "proxy-429", + ], +) +def test_is_pre_routing_rejection_is_decided_by_pipeline_stage_not_exception_type( + exception: Exception, request_data: Mapping[str, object], expected: bool +) -> None: + assert is_pre_routing_rejection(exception, request_data) is expected + + +class _RecordingSpendWriter: + def __init__(self) -> None: + self.rows: list[Mapping[str, object]] = [] # mutable-ok: test double records the spend rows written + + async def update_database(self, **kwargs: object) -> None: + self.rows.append(kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("flag_enabled", "exception", "request_data", "rows_written"), + [ + (True, _INVALID_REQUEST_400, _PRE_ROUTING_REQUEST, 0), + (True, litellm.BudgetExceededError(current_cost=1.0, max_budget=0.0), _PRE_ROUTING_REQUEST, 0), + (True, ProxyModelNotFoundError(route="/chat/completions", model_name="4.6"), _ROUTED_REQUEST, 0), + (True, _INVALID_REQUEST_400, _ROUTED_REQUEST, 1), + (True, _INVALID_REQUEST_400, _SPOOFED_REQUEST, 1), + (True, _PROVIDER_AUTH_ERROR, _ROUTED_REQUEST, 1), + (False, _INVALID_REQUEST_400, _PRE_ROUTING_REQUEST, 1), + (False, litellm.BudgetExceededError(current_cost=1.0, max_budget=0.0), _PRE_ROUTING_REQUEST, 1), + ], + ids=[ + "flag-bad-body", + "flag-budget", + "flag-unknown-model", + "flag-guardrail-block-kept", + "flag-guardrail-block-with-spoofed-marker-kept", + "flag-provider-error-kept", + "off-bad-body", + "off-budget", + ], +) +async def test_async_post_call_failure_hook_disable_pre_routing_error_logs( + flag_enabled: bool, exception: Exception, request_data: Mapping[str, object], rows_written: int +) -> None: + writer: Final = _RecordingSpendWriter() + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="test_api_key", user_id="u", team_id="t", request_route="/chat/completions" + ) + with patch("litellm.proxy.proxy_server.general_settings", {"disable_pre_routing_error_logs": flag_enabled}): + await ProxyDBLogger(spend_writer=lambda: cast(DBSpendUpdateWriter, writer)).async_post_call_failure_hook( + request_data={**request_data, "metadata": dict[str, object]()}, + original_exception=exception, + user_api_key_dict=user_api_key_dict, + ) + assert len(writer.rows) == rows_written diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 37028558af3..f727448f3ac 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -54,6 +54,8 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _LOCK_DELETED_TEAM_MEMBERS_SQL, _STRIP_DELETED_TEAMS_FROM_USERS_SQL, DeletedTeams, + _enrich_team_list_with_caller_membership, + _TeamListCallerIdentity, GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, _build_team_list_where_conditions, @@ -91,6 +93,8 @@ from litellm.router import Router from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddRequest, BulkTeamMemberAddResponse, + TeamListCallerMembership, + TeamListItem, TeamMemberAddResult, ) from litellm.types.utils import StandardAuditLogPayload @@ -4172,6 +4176,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) with patch( @@ -4191,9 +4196,8 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): ) # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4229,6 +4233,7 @@ async def test_list_team_v2_security_check_admin_user(): ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) # Should NOT raise an exception @@ -4242,9 +4247,8 @@ async def test_list_team_v2_security_check_admin_user(): ) # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4310,10 +4314,8 @@ async def test_list_team_v2_with_status_deleted(): mock_db.litellm_teamtable.count.assert_not_called() # Should return results without error - assert "teams" in result - assert "total" in result - assert result["total"] == 2 - assert len(result["teams"]) == 2 + assert len(result.teams) == 2 + assert result.total == 2 @pytest.mark.asyncio @@ -4372,6 +4374,7 @@ async def test_list_team_v2_includes_litellm_model_table(): side_effect=lambda **kw: [_team_row("team_1", kw.get("include"))] ) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( @@ -4383,8 +4386,8 @@ async def test_list_team_v2_includes_litellm_model_table(): status=None, ) - assert result["teams"][0].litellm_model_table is not None - assert result["teams"][0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} + assert result.teams[0].litellm_model_table is not None + assert result.teams[0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} mock_db.litellm_deletedteamtable.find_many = AsyncMock( side_effect=lambda **kw: [_team_row("team_2", kw.get("include"))] @@ -4465,6 +4468,7 @@ async def test_list_team_v2_org_admin_sees_org_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) @@ -4482,9 +4486,9 @@ async def test_list_team_v2_org_admin_sees_org_teams(): status=None, ) - assert result["total"] == 1 - assert len(result["teams"]) == 1 - assert result["teams"][0].members_count == 1 + assert result.total == 1 + assert len(result.teams) == 1 + assert result.teams[0].members_count == 1 # Verify org-scoped where clause where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] @@ -4565,6 +4569,7 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): return_value=[mock_team_1, mock_team_2] ) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) @@ -4583,8 +4588,8 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): status=None, ) - assert result["total"] == 2 - assert len(result["teams"]) == 2 + assert result.total == 2 + assert len(result.teams) == 2 # Verify the where clause scopes by org OR own membership — no # top-level team_id filter that would hide org teams they aren't in @@ -4689,6 +4694,7 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( prisma_client = MagicMock() prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) prisma_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=LiteLLM_UserTable( @@ -4721,8 +4727,8 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( sort_order="asc", status=None, ) - assert result["total"] == len(result["teams"]) - return [t.team_id for t in result["teams"]] + assert result.total == len(result.teams) + return [t.team_id for t in result.teams] own_view = ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] assert await list_teams("org_admin_user") == own_view @@ -4950,6 +4956,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): } mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) result = await list_team_v2( @@ -4966,7 +4973,7 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): status=None, ) - assert result["total"] == 1 + assert result.total == 1 # Verify the where clause filters by user's teams AND org scope where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] @@ -5219,6 +5226,7 @@ async def test_list_team_v2_populates_keys_count(): mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) + mock_db.litellm_teammembership.find_many = AsyncMock(return_value=[]) mock_db.litellm_verificationtoken.group_by = AsyncMock( return_value=[ {"team_id": "team_a", "_count": {"team_id": 3}}, @@ -5235,16 +5243,391 @@ async def test_list_team_v2_populates_keys_count(): status=None, ) - assert result["total"] == 2 - by_id = {t.team_id: t for t in result["teams"]} + assert result.total == 2 + by_id = {t.team_id: t for t in result.teams} assert by_id["team_a"].keys_count == 3 assert by_id["team_b"].keys_count == 0 - # The aggregate is one batched query, filtered by the page's team IDs. group_by_kwargs = mock_db.litellm_verificationtoken.group_by.call_args.kwargs assert group_by_kwargs["by"] == ["team_id"] assert group_by_kwargs["where"] == {"team_id": {"in": ["team_a", "team_b"]}} assert group_by_kwargs["count"] == {"team_id": True} + membership_where: Final = mock_db.litellm_teammembership.find_many.call_args.kwargs["where"] + assert membership_where == { + "team_id": {"in": ("team_a", "team_b")}, + "user_id": "admin_user_123", + } + + +def _team_list_prisma( + membership_rows: Sequence[LiteLLM_TeamMembership], + budget_rows: Sequence[LiteLLM_BudgetTableFull], +) -> tuple[MagicMock, AsyncMock, AsyncMock]: + prisma_client: Final = MagicMock() + membership_query: Final = AsyncMock(return_value=tuple(membership_rows)) + budget_query: Final = AsyncMock(return_value=tuple(budget_rows)) + prisma_client.db.litellm_teammembership.find_many = membership_query + prisma_client.db.litellm_budgettable.find_many = budget_query + return prisma_client, membership_query, budget_query + + +def _team_list_member( + team_id: str, + user_id: str = "caller", + default_budget_id: str | None = None, +) -> TeamListItem: + metadata: Final = {"team_member_budget_id": default_budget_id} if default_budget_id is not None else None + return TeamListItem( + team_id=team_id, + metadata=metadata, + members_with_roles=[Member(user_id=user_id, role="user")], + ) + + +def _team_list_caller_identity( + user_id: str | None = "caller", + user_email: str | None = None, +) -> _TeamListCallerIdentity: + return _TeamListCallerIdentity(user_id=user_id, user_email=user_email) + + +def _budget_table( + budget_id: str, + max_budget: float | None, + budget_reset_at: datetime | None = None, + temp_budget_increase: float | None = None, + temp_budget_expiry: datetime | None = None, +) -> LiteLLM_BudgetTableFull: + return LiteLLM_BudgetTableFull( + budget_id=budget_id, + max_budget=max_budget, + budget_reset_at=budget_reset_at, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + temp_budget_increase=temp_budget_increase, + temp_budget_expiry=temp_budget_expiry, + ) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_member_spend_budget_and_reset_date() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + member_budget: Final = _budget_table("member-budget", 100, reset_at) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + spend=50, + litellm_budget_table=member_budget, + ) + prisma_client, _, budget_query = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=50, + max_budget=100, + budget_reset_at=reset_at, + ) + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_includes_roster_members_without_membership_rows() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=None, + ) + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_matches_email_only_roster_member() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id=None, user_email="caller@x.com", role="user")], + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + _team_list_caller_identity(user_email="caller@x.com"), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=None, + ) + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() + + +@pytest.mark.parametrize("caller_user_email", ("other@x.com", None)) +@pytest.mark.asyncio +async def test_team_list_caller_membership_ignores_email_only_roster_member_without_matching_email( + caller_user_email: str | None, +) -> None: + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id=None, user_email="caller@x.com", role="user")], + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + (), + [_budget_table("team-budget", 50)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + _team_list_caller_identity(user_email=caller_user_email), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_team_default_without_member_reset_date() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + membership: Final = LiteLLM_TeamMembership(user_id="caller", team_id="team-a", spend=12) + prisma_client, _, budget_query = _team_list_prisma( + [membership], + [_budget_table("team-budget", 50, reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=12, + max_budget=50, + budget_reset_at=None, + ) + budget_query.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_reset_date_from_member_linked_default_budget() -> None: + reset_at: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) + member_budget: Final = _budget_table("team-budget", 50, reset_at) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, budget_query = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=50, + budget_reset_at=reset_at, + ) + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_ignores_zero_team_default_budget() -> None: + membership: Final = LiteLLM_TeamMembership(user_id="caller", team_id="team-a") + prisma_client, _, _ = _team_list_prisma( + [membership], + [_budget_table("team-budget", 0, datetime(2026, 10, 20, tzinfo=timezone.utc))], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership(spend=0.0, max_budget=None, budget_reset_at=None) + + +@pytest.mark.parametrize( + ("expiry", "expected_budget"), + ( + (datetime(2030, 1, 2, tzinfo=timezone.utc), 110), + (datetime(2029, 12, 31, tzinfo=timezone.utc), 100), + ), +) +@pytest.mark.asyncio +async def test_team_list_caller_membership_applies_only_active_temporary_increases( + expiry: datetime, + expected_budget: float, +) -> None: + member_budget: Final = _budget_table( + "member-budget", + 100, + datetime(2030, 2, 1, tzinfo=timezone.utc), + temp_budget_increase=10, + temp_budget_expiry=expiry, + ) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, _ = _team_list_prisma([membership], []) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + _team_list_caller_identity(), + datetime(2030, 1, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is not None + assert teams[0].caller_membership.max_budget == expected_budget + assert teams[0].caller_membership.budget_reset_at == datetime(2030, 2, 1, tzinfo=timezone.utc) + + +@pytest.mark.parametrize( + "member_reset_at", + ( + datetime(2030, 1, 15, tzinfo=timezone.utc), + None, + ), +) +@pytest.mark.asyncio +async def test_team_list_caller_membership_uses_private_row_reset_with_inherited_budget( + member_reset_at: datetime | None, +) -> None: + team_reset_at: Final = datetime(2030, 2, 1, tzinfo=timezone.utc) + member_budget: Final = _budget_table( + "member-budget", + None, + member_reset_at, + temp_budget_increase=10, + temp_budget_expiry=datetime(2030, 1, 2, tzinfo=timezone.utc), + ) + membership: Final = LiteLLM_TeamMembership( + user_id="caller", + team_id="team-a", + litellm_budget_table=member_budget, + ) + prisma_client, _, _ = _team_list_prisma( + [membership], + [_budget_table("team-budget", 50, team_reset_at)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a", default_budget_id="team-budget")], + _team_list_caller_identity(), + datetime(2030, 1, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership == TeamListCallerMembership( + spend=0.0, + max_budget=60, + budget_reset_at=member_reset_at, + ) + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_is_null_for_nonmember_admins() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma([], []) + team: Final = TeamListItem( + team_id="team-a", + metadata={"team_member_budget_id": "team-budget"}, + members_with_roles=[Member(user_id="another-user", role="user")], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [team], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_awaited_once() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_skips_queries_without_caller_user_id() -> None: + prisma_client, membership_query, budget_query = _team_list_prisma((), ()) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [_team_list_member("team-a")], + _team_list_caller_identity(user_id=None), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert teams[0].caller_membership is None + membership_query.assert_not_awaited() + budget_query.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_list_caller_membership_batches_page_membership_and_budget_queries() -> None: + memberships: Final = ( + LiteLLM_TeamMembership(user_id="caller", team_id="team-a"), + LiteLLM_TeamMembership(user_id="caller", team_id="team-b"), + ) + prisma_client, membership_query, budget_query = _team_list_prisma( + memberships, + [_budget_table("budget-a", 50), _budget_table("budget-b", 75)], + ) + + teams: Final = await _enrich_team_list_with_caller_membership( + prisma_client, + [ + _team_list_member("team-a", default_budget_id="budget-a"), + _team_list_member("team-b", default_budget_id="budget-b"), + ], + _team_list_caller_identity(), + datetime(2026, 10, 1, tzinfo=timezone.utc), + ) + + assert [team.caller_membership for team in teams] == [ + TeamListCallerMembership(spend=0.0, max_budget=50, budget_reset_at=None), + TeamListCallerMembership(spend=0.0, max_budget=75, budget_reset_at=None), + ] + membership_query.assert_awaited_once() + budget_query.assert_awaited_once() @pytest.mark.asyncio @@ -5282,8 +5665,8 @@ async def test_list_team_v2_keys_count_skipped_for_empty_page(): status=None, ) - assert result["total"] == 0 - assert result["teams"] == [] + assert result.total == 0 + assert result.teams == [] mock_db.litellm_verificationtoken.group_by.assert_not_called() @@ -5332,7 +5715,7 @@ async def test_list_team_v2_keys_count_skipped_for_deleted_status(): status="deleted", ) - assert result["total"] == 1 + assert result.total == 1 mock_db.litellm_verificationtoken.group_by.assert_not_called() @@ -12178,6 +12561,36 @@ async def test_team_member_me_returns_caller_membership(mock_db_client): assert call_kwargs["team_id"] == team_id +@pytest.mark.asyncio +@pytest.mark.parametrize( + "spend, stored_total, reported_total", + [(612.4, 148.1, 612.4), (12.25, 97.8, 97.8)], + ids=["legacy_row_floors_to_cycle_spend", "healthy_row_keeps_its_total"], +) +async def test_team_member_me_reports_total_spend_no_lower_than_cycle_spend( + mock_db_client, spend, stored_total, reported_total +): + from fastapi import Request + + from litellm.proxy.management_endpoints.team_endpoints import team_member_me + + team_id = "team-me-floor" + caller_id = "alice@example.com" + team = _build_team_for_me(team_id, [{"user_id": caller_id, "user_email": None, "role": "user"}]) + membership = _build_membership_for_me(caller_id, team_id, spend=spend).model_copy( + update={"total_spend": stored_total} + ) + p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=None) + with p_team, p_membership, p_user: + response = await team_member_me( + http_request=MagicMock(spec=Request), + team_id=team_id, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id), + ) + + assert (response.spend, response.total_spend) == (spend, reported_total) + + @pytest.mark.asyncio async def test_team_member_me_matches_email_only_member(mock_db_client): """ @@ -16225,7 +16638,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch): litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0), ) mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=1) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) @@ -16248,15 +16661,76 @@ async def test_reset_team_member_spend_fn_success(monkeypatch): assert response["spend"] == 0.0 assert response["previous_spend"] == 10.0 assert response["max_budget"] == 50.0 - mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, - data={"spend": 0.0}, - ) + statement, reset_user_id, reset_team_id, reset_to = mock_prisma_client.db.execute_raw.await_args.args + assert (reset_user_id, reset_team_id, reset_to) == ("member-1", "team-1", 0.0) + assert "total_spend = GREATEST(total_spend, spend)" in statement + assert "spend >= $3" in statement assert await real_cache.async_get_cache(key="team-1_member-1") is None assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0 +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_caches(monkeypatch): + """The reset validates reset_to against a spend it read a moment ago. If the row moved + underneath it (a flush landed, another admin reset first) the guarded UPDATE matches no + row, and the endpoint must say so instead of reporting a reset that never happened.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="cached-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ), pytest.raises(HTTPException) as exc_info: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=5.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert exc_info.value.status_code == 409 + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") == "cached-membership" + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_get_all_team_memberships_floors_total_spend_at_current_cycle_spend(): + """Rows written before total_spend existed carry their history only in spend. /team/info + and /team/list both read through this helper, so a legacy row reports its lifetime spend + as at least the current cycle, while a healthy row keeps its own total.""" + from litellm.proxy.management_endpoints.team_endpoints import get_all_team_memberships + + legacy_row = MagicMock() + legacy_row.model_dump = lambda: {"user_id": "legacy", "team_id": "team-1", "spend": 612.4, "total_spend": 148.1} + healthy_row = MagicMock() + healthy_row.model_dump = lambda: {"user_id": "healthy", "team_id": "team-1", "spend": 12.25, "total_spend": 97.8} + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[legacy_row, healthy_row]) + + legacy, healthy = await get_all_team_memberships(mock_prisma_client, ["team-1"]) + + assert (legacy.spend, legacy.total_spend) == (612.4, 612.4) + assert (healthy.spend, healthy.total_spend) == (12.25, 97.8) + + @pytest.mark.asyncio async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch): mock_prisma_client = MagicMock() @@ -16372,7 +16846,7 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0) mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=1) with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests "litellm.proxy.management_endpoints.team_endpoints.get_team_object", @@ -17611,6 +18085,7 @@ async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an disable_audit_logging_for_mocked_team, ): import contextlib + from litellm.models.budget import LiteLLM_BudgetTable with contextlib.ExitStack() as stack: diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index cc701e1ebd4..7530703a05a 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -145,9 +145,7 @@ class TestAnthropicLoggingHandlerModelFallback: '{"type": "message_stop"}', ] - def _create_mock_logging_obj( - self, model_in_details: str = None - ) -> LiteLLMLoggingObj: + def _create_mock_logging_obj(self, model_in_details: str = None) -> LiteLLMLoggingObj: """Create a mock logging object with optional model in model_call_details""" mock_logging_obj = MagicMock() @@ -166,21 +164,13 @@ class TestAnthropicLoggingHandlerModelFallback: mock_handler = MagicMock() return mock_handler - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) - @patch.object( - AnthropicPassthroughLoggingHandler, "_create_anthropic_response_logging_payload" - ) - def test_model_from_request_body_used_when_present( - self, mock_create_payload, mock_build_response - ): + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") + @patch.object(AnthropicPassthroughLoggingHandler, "_create_anthropic_response_logging_payload") + def test_model_from_request_body_used_when_present(self, mock_create_payload, mock_build_response): """Test that model from request_body is used when present""" # Arrange request_body = {"model": "claude-3-sonnet-20240229"} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-haiku-20240307" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-haiku-20240307") passthrough_handler = self._create_mock_passthrough_handler() # Mock successful response building @@ -204,57 +194,37 @@ class TestAnthropicLoggingHandlerModelFallback: # Verify that _build_complete_streaming_response was called with the request_body model mock_build_response.assert_called_once() call_args = mock_build_response.call_args - assert ( - call_args[1]["model"] == "claude-3-sonnet-20240229" - ) # Should use request_body model + assert call_args[1]["model"] == "claude-3-sonnet-20240229" # Should use request_body model def test_model_fallback_logic_isolated(self): """Test just the model fallback logic in isolation""" # Test case 1: Model from request body request_body = {"model": "claude-3-sonnet-20240229"} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-haiku-20240307" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-haiku-20240307") # Extract the logic directly from the function model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "claude-3-sonnet-20240229" # Should use request_body model # Test case 2: Fallback to logging obj request_body = {} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-haiku-20240307" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-haiku-20240307") model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "claude-3-haiku-20240307" # Should use fallback model # Test case 3: Empty string in request body, fallback to logging obj request_body = {"model": ""} - logging_obj = self._create_mock_logging_obj( - model_in_details="claude-3-opus-20240229" - ) + logging_obj = self._create_mock_logging_obj(model_in_details="claude-3-opus-20240229") model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "claude-3-opus-20240229" # Should use fallback model @@ -264,11 +234,7 @@ class TestAnthropicLoggingHandlerModelFallback: logging_obj = self._create_mock_logging_obj() model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "" # Should be empty @@ -284,11 +250,7 @@ class TestAnthropicLoggingHandlerModelFallback: # Extract the logic directly from the function model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "" # Should remain empty since no fallback available @@ -298,11 +260,7 @@ class TestAnthropicLoggingHandlerModelFallback: logging_obj = self._create_mock_logging_obj() # Empty dict model = request_body.get("model", "") - if ( - not model - and hasattr(logging_obj, "model_call_details") - and logging_obj.model_call_details.get("model") - ): + if not model and hasattr(logging_obj, "model_call_details") and logging_obj.model_call_details.get("model"): model = logging_obj.model_call_details.get("model") assert model == "" # Should remain empty @@ -311,9 +269,7 @@ class TestAnthropicLoggingHandlerModelFallback: class TestAzureAnthropicCostCalculation: """Test the custom_llm_provider cost calculation logic for Azure AI Anthropic.""" - def _create_mock_logging_obj( - self, model: str = None, custom_llm_provider: str = None - ) -> LiteLLMLoggingObj: + def _create_mock_logging_obj(self, model: str = None, custom_llm_provider: str = None) -> LiteLLMLoggingObj: """Create a mock logging object with optional model and custom_llm_provider""" mock_logging_obj = MagicMock() mock_model_call_details = {} @@ -326,9 +282,7 @@ class TestAzureAnthropicCostCalculation: return mock_logging_obj @patch("litellm.completion_cost") - def test_cost_calculation_with_azure_ai_custom_llm_provider( - self, mock_completion_cost - ): + def test_cost_calculation_with_azure_ai_custom_llm_provider(self, mock_completion_cost): """Test that custom_llm_provider is passed to completion_cost for Azure AI Anthropic""" from litellm.types.utils import ModelResponse from datetime import datetime @@ -397,9 +351,7 @@ class TestAzureAnthropicCostCalculation: assert call_kwargs["custom_llm_provider"] is None @patch("litellm.completion_cost") - def test_cost_calculation_does_not_duplicate_provider_prefix( - self, mock_completion_cost - ): + def test_cost_calculation_does_not_duplicate_provider_prefix(self, mock_completion_cost): """Test that provider prefix is not duplicated if already present in model name""" from litellm.types.utils import ModelResponse from datetime import datetime @@ -436,9 +388,7 @@ class TestAzureAnthropicCostCalculation: assert call_kwargs["custom_llm_provider"] == "azure_ai" @patch("litellm.completion_cost") - def test_cost_calculation_resolves_unknown_model_from_litellm_params( - self, mock_completion_cost - ): + def test_cost_calculation_resolves_unknown_model_from_litellm_params(self, mock_completion_cost): """When the body model is the "unknown" sentinel, the deployment model from litellm_params must be used for costing, not "unknown" (which makes completion_cost raise and the cost silently fall back to $0).""" @@ -451,9 +401,7 @@ class TestAzureAnthropicCostCalculation: logging_obj = self._create_mock_logging_obj(model="unknown") logging_obj.model_call_details["litellm_params"] = { "model": "anthropic/claude-3-5-haiku-20241022", - "metadata": { - "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" - }, + "metadata": {"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"}, } logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -471,17 +419,12 @@ class TestAzureAnthropicCostCalculation: ) mock_completion_cost.assert_called_once() - assert ( - mock_completion_cost.call_args[1]["model"] - == "anthropic/claude-3-5-haiku-20241022" - ) + assert mock_completion_cost.call_args[1]["model"] == "anthropic/claude-3-5-haiku-20241022" assert kwargs["response_cost"] == 0.001 assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" @patch("litellm.completion_cost") - def test_cost_calculation_resolves_unknown_model_from_model_group( - self, mock_completion_cost - ): + def test_cost_calculation_resolves_unknown_model_from_model_group(self, mock_completion_cost): """With only model_group available (no deployment litellm_params.model), the leading passthrough/ prefix must be stripped so the cost map can resolve the model.""" @@ -493,9 +436,7 @@ class TestAzureAnthropicCostCalculation: logging_obj = self._create_mock_logging_obj(model="unknown") logging_obj.model_call_details["litellm_params"] = { - "metadata": { - "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" - } + "metadata": {"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"} } logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -513,16 +454,11 @@ class TestAzureAnthropicCostCalculation: ) mock_completion_cost.assert_called_once() - assert ( - mock_completion_cost.call_args[1]["model"] - == "anthropic/claude-3-5-haiku-20241022" - ) + assert mock_completion_cost.call_args[1]["model"] == "anthropic/claude-3-5-haiku-20241022" assert kwargs["response_cost"] == 0.002 @patch("litellm.completion_cost") - def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group( - self, mock_completion_cost - ): + def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group(self, mock_completion_cost): """When litellm_params.model is itself the "unknown" sentinel, the deployment-model branch must not short-circuit; resolution falls through to model_group so costing still prices the real model instead of "unknown".""" @@ -535,9 +471,7 @@ class TestAzureAnthropicCostCalculation: logging_obj = self._create_mock_logging_obj(model="unknown") logging_obj.model_call_details["litellm_params"] = { "model": "unknown", - "metadata": { - "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" - }, + "metadata": {"model_group": "passthrough/anthropic/claude-3-5-haiku-20241022"}, } logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -555,17 +489,12 @@ class TestAzureAnthropicCostCalculation: ) mock_completion_cost.assert_called_once() - assert ( - mock_completion_cost.call_args[1]["model"] - == "anthropic/claude-3-5-haiku-20241022" - ) + assert mock_completion_cost.call_args[1]["model"] == "anthropic/claude-3-5-haiku-20241022" assert kwargs["response_cost"] == 0.003 assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" @patch("litellm.completion_cost") - def test_streaming_cost_calculation_resolves_model_from_message_start_chunk( - self, mock_completion_cost - ): + def test_streaming_cost_calculation_resolves_model_from_message_start_chunk(self, mock_completion_cost): """On the bare /anthropic passthrough path litellm_params carries no model or model_group and the body model is the "unknown" sentinel; the model must be recovered from the message_start SSE event so completion_cost @@ -628,9 +557,7 @@ class TestAzureAnthropicCostCalculation: ), _sse("message_stop", {"type": "message_stop"}), ] - all_chunks = list( - PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) - ) + all_chunks = list(PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames)) logging_obj = RealLoggingObj( model="unknown", @@ -674,9 +601,7 @@ class TestAzureAnthropicCostCalculation: ] assert ( - AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( - chunks - ) + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(chunks) == "claude-3-5-haiku-20241022" ) @@ -692,14 +617,11 @@ class TestAzureAnthropicCostCalculation: ) assert ( - AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( - [raw_event] - ) + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks([raw_event]) == "claude-3-5-haiku-20241022" ) - class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" @@ -799,9 +721,7 @@ class TestAnthropicBatchPassthroughCostTracking: ) mock_batches_config_instance = MagicMock() - mock_batches_config_instance.transform_retrieve_batch_response.return_value = ( - mock_batch_response - ) + mock_batches_config_instance.transform_retrieve_batch_response.return_value = mock_batch_response mock_batches_config.return_value = mock_batches_config_instance # Test the handler @@ -929,9 +849,7 @@ class TestAnthropicBatchPassthroughCostTracking: "transform_retrieve_batch_response", return_value=mock_batch_response, ): - with patch.object( - AnthropicPassthroughLoggingHandler, "_store_batch_managed_object" - ): + with patch.object(AnthropicPassthroughLoggingHandler, "_store_batch_managed_object"): result = AnthropicPassthroughLoggingHandler.batch_creation_handler( httpx_response=mock_httpx_response, logging_obj=mock_logging_obj, @@ -946,10 +864,7 @@ class TestAnthropicBatchPassthroughCostTracking: # Verify unified_object_id contains anthropic/ prefix unified_object_id = result["kwargs"]["unified_object_id"] decoded = base64.urlsafe_b64decode(unified_object_id + "==").decode() - assert ( - "anthropic/claude-sonnet-4-5-20250929" in decoded - or "claude-sonnet-4-5-20250929" in decoded - ) + assert "anthropic/claude-sonnet-4-5-20250929" in decoded or "claude-sonnet-4-5-20250929" in decoded @pytest.mark.parametrize( "kwargs,expected_user_id,expected_team_id", @@ -1007,9 +922,7 @@ class TestAnthropicBatchPassthroughCostTracking: mock_managed_files_hook.store_unified_object_id = AsyncMock() with ( patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, - patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger" - ), + patch("litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"), ): mock_pl.get_proxy_hook.return_value = mock_managed_files_hook AnthropicPassthroughLoggingHandler._store_batch_managed_object( @@ -1060,9 +973,7 @@ class TestAnthropicBatchPassthroughCostTracking: is never tracked. The failure has to reach the log instead of being reported as a success.""" mock_managed_files_hook = MagicMock() - mock_managed_files_hook.store_unified_object_id = AsyncMock( - side_effect=RuntimeError("db down") - ) + mock_managed_files_hook.store_unified_object_id = AsyncMock(side_effect=RuntimeError("db down")) with ( patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, patch( @@ -1104,9 +1015,7 @@ class TestAnthropicBatchPassthroughCostTracking: this gate an id-scoped route reached the store with a mismatched id, where it could only either claim a row it did not create or fail the model_object_id unique constraint.""" - with patch.object( - AnthropicPassthroughLoggingHandler, "_store_batch_managed_object" - ) as mock_store: + with patch.object(AnthropicPassthroughLoggingHandler, "_store_batch_managed_object") as mock_store: AnthropicPassthroughLoggingHandler.batch_creation_handler( httpx_response=mock_httpx_response, logging_obj=mock_logging_obj, @@ -1120,9 +1029,7 @@ class TestAnthropicBatchPassthroughCostTracking: assert mock_store.call_count == (1 if registers else 0) - def test_batch_creation_handler_failure_status_code( - self, mock_logging_obj, mock_request_body - ): + def test_batch_creation_handler_failure_status_code(self, mock_logging_obj, mock_request_body): """Test batch creation handler with non-200 status code""" mock_response = MagicMock() mock_response.status_code = 400 @@ -1145,9 +1052,7 @@ class TestAnthropicBatchPassthroughCostTracking: assert result["kwargs"]["response_cost"] == 0.0 @patch("litellm.proxy.proxy_server.proxy_logging_obj") - def test_store_batch_managed_object_success( - self, mock_proxy_logging_obj, mock_logging_obj - ): + def test_store_batch_managed_object_success(self, mock_proxy_logging_obj, mock_logging_obj): """Test storing batch managed object""" from litellm.types.utils import LiteLLMBatch @@ -1178,9 +1083,7 @@ class TestAnthropicBatchPassthroughCostTracking: ) # Verify managed files hook was called - mock_proxy_logging_obj.get_proxy_hook.assert_called_once_with( - "managed_files" - ) + mock_proxy_logging_obj.get_proxy_hook.assert_called_once_with("managed_files") class TestBuildCompleteStreamingResponseRobustness: @@ -1343,11 +1246,7 @@ class TestPureTextFastPathParity: }, ) ) - frames.append( - self._sse( - "content_block_stop", {"type": "content_block_stop", "index": b} - ) - ) + frames.append(self._sse("content_block_stop", {"type": "content_block_stop", "index": b})) frames.append( self._sse( "message_delta", @@ -1405,9 +1304,7 @@ class TestPureTextFastPathParity: # And the full logged payload (sans non-deterministic response id). def _scrub(p): d = dict(p) - r = d.get("complete_streaming_response_in_db") or d.get( - "complete_streaming_response" - ) + r = d.get("complete_streaming_response_in_db") or d.get("complete_streaming_response") return d, getattr(r, "usage", None) assert _scrub(k_fast)[1] == _scrub(k_legacy)[1] @@ -1419,11 +1316,7 @@ class TestPureTextFastPathParity: self._assert_parity(self._text_stream(["Just one piece of text."])) def test_parity_cache_tokens(self): - self._assert_parity( - self._text_stream( - ["a", "b", "c"], input_tokens=20, cache_creation=5, cache_read=7 - ) - ) + self._assert_parity(self._text_stream(["a", "b", "c"], input_tokens=20, cache_creation=5, cache_read=7)) def test_parity_max_tokens_stop(self): self._assert_parity(self._text_stream(["tok"] * 8, stop_reason="max_tokens")) @@ -1484,12 +1377,7 @@ class TestPureTextFastPathParity: self._sse("message_stop", {"type": "message_stop"}), ] all_chunks = self._to_all_chunks(frames) - assert ( - AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - list(all_chunks) - ) - is None - ) + assert AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(list(all_chunks)) is None def test_collapse_returns_none_for_thinking(self): frames = [ @@ -1526,19 +1414,12 @@ class TestPureTextFastPathParity: self._sse("message_stop", {"type": "message_stop"}), ] all_chunks = self._to_all_chunks(frames) - assert ( - AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - list(all_chunks) - ) - is None - ) + assert AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(list(all_chunks)) is None def test_collapse_actually_shrinks_chunk_count(self): frames = self._text_stream(["a"] * 50) all_chunks = list(self._to_all_chunks(frames)) - collapsed = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - all_chunks - ) + collapsed = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) assert collapsed is not None # 50 text deltas + 50 event markers + 1 ping collapse to far fewer. assert len(collapsed) < len(all_chunks) / 2 @@ -1604,10 +1485,7 @@ class TestPureTextFastPathParity: self._sse("message_stop", {"type": "message_stop"}), ] all_chunks = list(self._to_all_chunks(frames)) - assert ( - AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) - is None - ) + assert AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) is None class TestInterruptedStreamOutputTokenRecovery: @@ -1722,9 +1600,7 @@ class TestInterruptedStreamOutputTokenRecovery: import litellm placeholder = 2 - result = self._run( - self._interrupted_chunks(placeholder_output_tokens=placeholder) - ) + result = self._run(self._interrupted_chunks(placeholder_output_tokens=placeholder)) usage = result["result"].usage expected = litellm.token_counter( @@ -1897,9 +1773,7 @@ class TestStreamFalseDeduplication: "delta": {"type": "text_delta", "text": "Hello"}, }, ), - TestStreamFalseDeduplication._sse( - "content_block_stop", {"type": "content_block_stop", "index": 0} - ), + TestStreamFalseDeduplication._sse("content_block_stop", {"type": "content_block_stop", "index": 0}), TestStreamFalseDeduplication._sse( "message_delta", { @@ -1949,15 +1823,12 @@ class TestStreamFalseDeduplication: # The assembled response must be stored on model_call_details so callbacks # can identify this as a completed streaming call, not an in-progress one. - assert ( - logging_obj.model_call_details.get("complete_streaming_response") - is not None - ), "complete_streaming_response must be set on model_call_details after assembly" + assert logging_obj.model_call_details.get("complete_streaming_response") is not None, ( + "complete_streaming_response must be set on model_call_details after assembly" + ) # The returned result must match what was stored - assert result["result"] is logging_obj.model_call_details.get( - "complete_streaming_response" - ) + assert result["result"] is logging_obj.model_call_details.get("complete_streaming_response") def test_dedup_guard_fires_when_stream_true_on_logging_obj(self): """ @@ -1986,21 +1857,16 @@ class TestStreamFalseDeduplication: assert logging_obj._is_assembled_stream_success(result=mock_response) is True # First dispatch sets the flag - assert not logging_obj.model_call_details.get( - "has_dispatched_final_stream_success" - ) + assert not logging_obj.model_call_details.get("has_dispatched_final_stream_success") logging_obj.model_call_details["has_dispatched_final_stream_success"] = True # Second dispatch would be blocked — simulate the guard check would_skip = bool( logging_obj._is_assembled_stream_success(result=mock_response) - and logging_obj.model_call_details.get( - "has_dispatched_final_stream_success" - ) + and logging_obj.model_call_details.get("has_dispatched_final_stream_success") ) assert would_skip is True, ( - "Dedup guard must block a second dispatch_success_handlers call for the " - "same assembled streaming response" + "Dedup guard must block a second dispatch_success_handlers call for the same assembled streaming response" ) def test_sse_fallback_path_sets_stream_true_for_dedup(self): @@ -2034,9 +1900,7 @@ class TestStreamFalseDeduplication: would_skip = bool( logging_obj._is_assembled_stream_success(result=mock_response) - and logging_obj.model_call_details.get( - "has_dispatched_final_stream_success" - ) + and logging_obj.model_call_details.get("has_dispatched_final_stream_success") ) assert would_skip is True @@ -2100,9 +1964,9 @@ class TestNonStreamingResponseRedaction: logging_obj=logging_obj, ) - assert ( - "complete_streaming_response" not in logging_obj.model_call_details - ), "non-streaming responses must not populate complete_streaming_response" + assert "complete_streaming_response" not in logging_obj.model_call_details, ( + "non-streaming responses must not populate complete_streaming_response" + ) def test_streaming_sets_complete_streaming_response(self): from litellm.types.utils import ModelResponse @@ -2119,10 +1983,7 @@ class TestNonStreamingResponseRedaction: logging_obj=logging_obj, ) - assert ( - logging_obj.model_call_details.get("complete_streaming_response") - is response - ) + assert logging_obj.model_call_details.get("complete_streaming_response") is response def test_non_streaming_response_is_redacted_when_message_logging_off(self): from litellm.litellm_core_utils.redact_messages import ( @@ -2196,10 +2057,8 @@ class TestAnthropicUsageOnlyFallback: ] def test_build_usage_only_recovers_cache_inclusive_usage(self): - response = ( - AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( - all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022" - ) + response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022" ) assert response is not None usage = response.usage @@ -2259,10 +2118,8 @@ class TestAnthropicUsageOnlyFallback: } ), ] - response = ( - AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( - all_chunks=chunks, model="unknown" - ) + response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="unknown" ) assert response is not None assert response.model == "claude-opus-4-6" @@ -2287,9 +2144,7 @@ class TestAnthropicUsageOnlyFallback: ('data: {"a": 1}', {"a": 1}), ], ) - def test_extract_sse_data_handles_malformed_and_sentinel_lines( - self, event_str, expected - ): + def test_extract_sse_data_handles_malformed_and_sentinel_lines(self, event_str, expected): assert extract_sse_data(event_str) == expected def _real_logging_obj(self): @@ -2309,12 +2164,8 @@ class TestAnthropicUsageOnlyFallback: return logging_obj @patch("litellm.completion_cost") - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) - def test_handler_falls_back_when_assembly_returns_none( - self, mock_assemble, mock_cost - ): + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") + def test_handler_falls_back_when_assembly_returns_none(self, mock_assemble, mock_cost): mock_assemble.return_value = None mock_cost.return_value = 0.0021 logging_obj = self._real_logging_obj() @@ -2335,9 +2186,7 @@ class TestAnthropicUsageOnlyFallback: assert result["kwargs"]["response_cost"] == 0.0021 @patch("litellm.completion_cost") - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") def test_handler_falls_back_when_assembly_raises(self, mock_assemble, mock_cost): import litellm @@ -2367,9 +2216,7 @@ class TestAnthropicUsageOnlyFallback: assert result["result"].usage.completion_tokens == 55 assert result["kwargs"]["response_cost"] == 0.0021 - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") def test_handler_returns_none_when_no_usage_recoverable(self, mock_assemble): # assembly fails AND the chunks carry no usage event, so there is nothing # to price; the handler must return None rather than fabricate a response @@ -2391,15 +2238,9 @@ class TestAnthropicUsageOnlyFallback: assert result["result"] is None assert result["kwargs"] == {} - @patch.object( - AnthropicPassthroughLoggingHandler, "build_usage_only_response_from_chunks" - ) - @patch.object( - AnthropicPassthroughLoggingHandler, "build_complete_streaming_response" - ) - def test_handler_does_not_crash_when_usage_only_fallback_raises( - self, mock_assemble, mock_fallback - ): + @patch.object(AnthropicPassthroughLoggingHandler, "build_usage_only_response_from_chunks") + @patch.object(AnthropicPassthroughLoggingHandler, "build_complete_streaming_response") + def test_handler_does_not_crash_when_usage_only_fallback_raises(self, mock_assemble, mock_fallback): # if the usage-only fallback itself raises, it must be treated as None and # drop gracefully, not propagate out and crash the success handler mock_assemble.return_value = None @@ -2427,7 +2268,6 @@ class TestAnthropicResponseCostRecordedOnModelCallDetails: builder must record it there or streaming pass-through logs $0.""" - class TestAnthropicPassthroughFastMode: """Anthropic charges a provider-specific multiplier for ``speed=fast``, applied off ``usage.speed`` and covering every token type, cache included. The response usage @@ -2436,20 +2276,20 @@ class TestAnthropicPassthroughFastMode: MODEL = "claude-opus-4-8" STREAM_CHUNKS = [ - 'event: message_start', + "event: message_start", 'data: {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant",' ' "model": "claude-opus-4-8", "content": [], "stop_reason": null,' ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 0}}}', - 'event: content_block_start', + "event: content_block_start", 'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}', - 'event: content_block_delta', + "event: content_block_delta", 'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}', - 'event: content_block_stop', + "event: content_block_stop", 'data: {"type": "content_block_stop", "index": 0}', - 'event: message_delta', + "event: message_delta", 'data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"},' ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}}', - 'event: message_stop', + "event: message_stop", 'data: {"type": "message_stop"}', ] @@ -2643,7 +2483,6 @@ class TestRecordPartialUsageForFailure: assert "response_cost" not in logging_obj.model_call_details - @pytest.mark.parametrize( ("all_chunks", "interrupted"), [ @@ -2664,9 +2503,7 @@ def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects( @pytest.mark.asyncio -async def test_anthropic_passthrough_handler( - mock_httpx_response, mock_response, mock_logging_obj -): +async def test_anthropic_passthrough_handler(mock_httpx_response, mock_response, mock_logging_obj): """ Unit test - Assert that the anthropic passthrough handler calls the litellm logging object's async_success_handler """ @@ -3007,3 +2844,102 @@ def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj): assert logged.choices[0].finish_reason == "stop" provider_fields: Final = logged.choices[0].message.provider_specific_fields or {} assert "stop_details" not in provider_fields + + +def _go_sse(event: str, data: dict) -> str: + return f"event: {event}\ndata: {json.dumps(data)}\n\n" + + +_GO_MESSAGE_DELTA_USAGE: Final = { + "input_tokens": 1200, + "output_tokens": 300, + "cost": 0, + "opencode": {"billing_source": "go", "list_cost": 0.00412}, +} + + +def _go_stream_chunks() -> list[str]: + message = { + "id": "msg_go", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 1200, "output_tokens": 0}, + } + return [ + _go_sse("message_start", {"type": "message_start", "message": message}), + _go_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _go_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + ), + _go_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _go_sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": _GO_MESSAGE_DELTA_USAGE, + }, + ), + _go_sse("message_stop", {"type": "message_stop"}), + ] + + +def _go_logging_obj() -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="opencode_go/claude-haiku-5-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="go-stream", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "opencode_go" + logging_obj.model_call_details["stream"] = True + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + return logging_obj + + +def test_streamed_messages_bill_the_list_cost_reported_in_the_final_message_delta(): + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=_go_logging_obj(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "opencode_go/claude-haiku-5-5", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=_go_stream_chunks(), + end_time=datetime.now(), + ) + + assert result["kwargs"]["response_cost"] == 0.00412 + assert result["result"]._hidden_params["response_cost"] == 0.00412 + + +def test_streamed_messages_from_other_providers_ignore_a_zero_usage_cost_field(): + logging_obj = _go_logging_obj() + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude-haiku-5-5", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=_go_stream_chunks(), + end_time=datetime.now(), + ) + + headers = result["result"]._hidden_params.get("additional_headers") or {} + assert "llm_provider-x-litellm-response-cost" not in headers + assert result["kwargs"].get("response_cost") != 0.0 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 92bf2b45fab..82dd3479321 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -12027,3 +12027,24 @@ class TestResponseCostFromLoggingObjForCacheHits: ) assert cost == "" + + +def test_maybe_get_model_id_reads_the_chat_metadata_the_router_updated_in_place(): + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"metadata": {"model_info": {"id": "declared-text"}}}) + + assert processing_obj.maybe_get_model_id(None) == "declared-text" + + +def test_maybe_get_model_id_prefers_litellm_metadata_over_caller_metadata_without_a_deployment(): + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={ + "metadata": {"trace_id": "t-1"}, + "litellm_metadata": {"model_info": {"id": "vision"}}, + } + ) + + assert processing_obj.maybe_get_model_id(None) == "vision" + + +def test_maybe_get_model_id_returns_none_before_the_router_picked_a_deployment(): + assert ProxyBaseLLMRequestProcessing(data={"metadata": {}}).maybe_get_model_id(None) is None diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 0d558f206d8..57b68e6db05 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1,6 +1,8 @@ import json import os +from collections.abc import Iterator from typing import Final +from unittest.mock import patch import pytest from fastapi.testclient import TestClient @@ -181,6 +183,12 @@ def mock_auth(): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.fixture(autouse=True) +def isolated_environment() -> Iterator[None]: + with patch.dict(os.environ): + yield + + @pytest.fixture(autouse=True) def fresh_settings_store(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index e83efa21dcc..6557f0d1446 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -18,7 +18,7 @@ from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, ProxyLogging from litellm.types.utils import CachingDetails @@ -194,7 +194,9 @@ async def test_post_call_failure_hook_attributes_single_router_deployment( @pytest.mark.asyncio -async def test_pre_routing_reject_spend_log_keeps_public_model_group(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_pre_routing_reject_spend_log_keeps_public_model_group( + proxy_logging, make_user_api_key_auth, monkeypatch +): from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -1001,7 +1003,9 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: standard_logging_object = kwargs.get("standard_logging_object") - recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None) + recorded.append( + standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None + ) monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) logging_obj = LiteLLMLoggingObj( @@ -1029,3 +1033,52 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( await asyncio.sleep(0) await asyncio.sleep(0) assert recorded == ["failure"] + + +@pytest.mark.asyncio +async def test_handle_logging_proxy_only_path_flags_requests_that_died_before_request_setup( + proxy_logging, make_user_api_key_auth +): + """A reject before the pipeline set up logging (bad key, bad body) reaches the failure hook with no logging + object. The handler mints a ``litellm_call_id`` so error logs can carry it, and marks the request with a + server-owned sentinel so spend tracking can tell it apart from a guardrail block, which arrives with the + pipeline's logging object. ``request_data`` is the client body, so a client-sent call id or marker value must + not change either outcome.""" + rejected_before_setup: dict[str, object] = { + "model": "m", + "messages": [{"role": "user", "content": "x"}], + "metadata": {}, + "litellm_call_id": "client-supplied", + } + await proxy_logging._handle_logging_proxy_only_error( + request_data=rejected_before_setup, + user_api_key_dict=make_user_api_key_auth(), + route="/chat/completions", + original_exception=HTTPException(status_code=401, detail="bad key"), + ) + + logging_obj = MagicMock() + logging_obj.call_type = "acompletion" + logging_obj.model_call_details = {} + logging_obj.async_failure_handler = AsyncMock() + blocked_by_guardrail: dict[str, object] = { + "litellm_logging_obj": logging_obj, + "litellm_call_id": "call-1", + "model": "m", + "messages": [{"role": "user", "content": "x"}], + "metadata": {}, + REJECTED_BEFORE_REQUEST_SETUP_KEY: True, + } + await proxy_logging._handle_logging_proxy_only_error( + request_data=blocked_by_guardrail, + user_api_key_dict=make_user_api_key_auth(), + route="/chat/completions", + original_exception=HTTPException(status_code=400, detail="blocked"), + ) + + assert ( + rejected_before_setup.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP, + rejected_before_setup["litellm_call_id"], + blocked_by_guardrail.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP, + blocked_by_guardrail["litellm_call_id"], + ) == (True, "client-supplied", False, "call-1") diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index a08d44b9b4d..082a29b7001 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -6,6 +6,7 @@ import pytest from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, + SpendResetWrites, budget_cascade_unit_of_work, enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, @@ -39,6 +40,9 @@ class FakeBatch: self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) + def execute_raw(self, query: str, *args: object) -> None: + self.calls.append(("execute_raw", {"query": query}, {"args": args})) + async def commit(self) -> None: self.commit_count += 1 @@ -166,7 +170,7 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch( async with budget_cascade_unit_of_work(_new_batch) as uow: writes = [getattr(uow, field.name) for field in fields(uow)] for write in writes: - if isinstance(write, LinkedSpendResetWrites): + if isinstance(write, SpendResetWrites): write.queue_spend_zero(where={"budget_id": "budget-1"}) else: write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) @@ -192,3 +196,18 @@ async def test_budget_cascade_raising_inside_block_skips_commit(): await _blow_up_mid_transaction() assert batch.commit_count == 0 + + +async def test_membership_total_spend_repair_is_queued_on_the_cascade_batch_ahead_of_the_reset(): + batch = FakeBatch() + + async with budget_cascade_unit_of_work(lambda: batch) as uow: + uow.team_memberships.queue_total_spend_repair(["budget-1", "budget-2"]) + uow.team_memberships.queue_spend_zero(where={"budget_id": "budget-1"}) + assert batch.commit_count == 0 + + assert batch.commit_count == 1 + repair, reset = batch.calls + assert repair[0] == "execute_raw" + assert repair[2] == {"args": (["budget-1", "budget-2"],)} + assert reset == ("litellm_teammembership.update_many", {"budget_id": "budget-1"}, {"spend": 0}) diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index e88d951ff09..285d90ef8de 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -29,6 +29,7 @@ from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.responses.mcp.mcp_streaming_iterator import MCPEnhancedStreamingIterator from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamingResponse, @@ -2384,8 +2385,11 @@ def _first_index(events: Sequence[object], event_type: str) -> int: async def _scripted_mcp_stream( - model: str, upstream_url: str, rounds: Sequence[bytes] -) -> tuple[tuple[object, ...], respx.Route, AsyncMock, AsyncMock]: + model: str, + upstream_url: str, + rounds: Sequence[bytes], + max_agentic_loops: int | None = None, +) -> tuple[tuple[object, ...], respx.Route, AsyncMock, AsyncMock, MCPEnhancedStreamingIterator]: upstream: Final = respx.post(upstream_url).mock( side_effect=[httpx.Response(200, content=body, headers=SSE_HEADERS) for body in rounds] ) @@ -2402,9 +2406,11 @@ async def _scripted_mcp_stream( tool_choice="required", input=[{"role": "user", "type": "message", "content": "What's the weather in San Francisco?"}], stream=True, + **({"max_agentic_loops": max_agentic_loops} if max_agentic_loops is not None else {}), ) + assert isinstance(response, MCPEnhancedStreamingIterator) events: Final = tuple([event async for event in response]) - return events, upstream, fetch, execute + return events, upstream, fetch, execute, response @pytest.mark.asyncio @@ -2417,7 +2423,7 @@ async def test_streaming_responses_api_with_mcp_tools( model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS[provider] with caplog.at_level(logging.ERROR): - events, upstream, fetch, execute = await _scripted_mcp_stream( + events, upstream, fetch, execute, _ = await _scripted_mcp_stream( model, upstream_url, (tool_round(), answer_round()) ) @@ -2436,6 +2442,22 @@ async def test_streaming_responses_api_with_mcp_tools( assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] +@pytest.mark.asyncio +@respx.mock +async def test_streaming_mcp_iterator_uses_max_agentic_loops_from_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS["openai"] + + _, _, _, _, iterator = await _scripted_mcp_stream( + model, + upstream_url, + (tool_round(), answer_round()), + max_agentic_loops=7, + ) + + assert iterator.max_tool_call_rounds == 7 + + @pytest.mark.asyncio @respx.mock async def test_streaming_mcp_event_order_and_response_id_consistency( @@ -2445,7 +2467,7 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( model, upstream_url, tool_round, answer_round = SCRIPTED_PROVIDERS["openai"] with caplog.at_level(logging.ERROR): - events, _, _, _ = await _scripted_mcp_stream(model, upstream_url, (tool_round(), answer_round())) + events, _, _, _, _ = await _scripted_mcp_stream(model, upstream_url, (tool_round(), answer_round())) created: Final = _first_index(events, "response.created") in_progress: Final = _first_index(events, "response.in_progress") diff --git a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py index 8eb73b7c6e1..ed39da1b7a9 100644 --- a/tests/unit/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/unit/responses/mcp/test_mcp_streaming_iterator.py @@ -2,6 +2,7 @@ import sys import types from contextlib import nullcontext from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -9,7 +10,6 @@ from mcp.types import CallToolResult, TextContent import litellm # noqa: F401 - ensures litellm.responses.main is registered in sys.modules from litellm.responses.mcp.mcp_streaming_iterator import ( - MAX_MCP_TOOL_CALL_ROUNDS, MCPEnhancedStreamingIterator, ) from litellm.types.llms.openai import ( @@ -96,7 +96,7 @@ def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return call_tool -def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: +def _make_iterator(initial_chunks, max_agentic_loops: int | str | None = None) -> MCPEnhancedStreamingIterator: return MCPEnhancedStreamingIterator( base_iterator=_FakeAsyncStream(initial_chunks), mcp_events=[], @@ -107,6 +107,7 @@ def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: "model": "gpt-4", "input": "what is berriai/litellm?", "tools": [{"type": "mcp"}], + **({"max_agentic_loops": max_agentic_loops} if max_agentic_loops is not None else {}), }, ) @@ -157,35 +158,43 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp @pytest.mark.asyncio -async def test_tool_call_rounds_are_capped(monkeypatch): +@pytest.mark.parametrize( + ("max_agentic_loops", "expected_cap"), + ( + pytest.param(None, 5, id="unset"), + pytest.param(7, 7, id="integer"), + pytest.param("7", 7, id="integer-string"), + pytest.param(2, 2, id="custom-cap"), + ), +) +async def test_tool_call_rounds_are_capped(monkeypatch, max_agentic_loops, expected_cap): """ A model that keeps calling tools every round must not loop forever — - auto-execution stops at MAX_MCP_TOOL_CALL_ROUNDS, and the follow-up made + auto-execution stops at the configured round cap, and the follow-up made at the cap drops "tools" from the request so the model is forced to answer in text instead of the stream just hanging. """ call_tool = _mock_mcp_environment(monkeypatch) - # Rounds 2..MAX_MCP_TOOL_CALL_ROUNDS keep calling the tool; the call made - # once the cap is hit returns a text-only response. - tool_call_streams = [ - _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, MAX_MCP_TOOL_CALL_ROUNDS + 1) - ] - aresponses_mock = AsyncMock(side_effect=[*tool_call_streams, _text_only_stream("giving up on tools")]) + tool_call_streams: Final = tuple( + _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, expected_cap + 1) + ) + aresponses_mock: Final = AsyncMock(side_effect=(*tool_call_streams, _text_only_stream("giving up on tools"))) monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) iterator = _make_iterator( [ _output_item_added_chunk(), _completed_chunk([_function_call("call_1", "read_wiki_contents")]), - ] + ], + max_agentic_loops=max_agentic_loops, ) _ = [chunk async for chunk in iterator] - assert iterator.tool_call_round == MAX_MCP_TOOL_CALL_ROUNDS - assert call_tool.call_count == MAX_MCP_TOOL_CALL_ROUNDS - assert aresponses_mock.call_count == MAX_MCP_TOOL_CALL_ROUNDS + assert iterator.tool_call_round == expected_cap + assert call_tool.call_count == expected_cap + assert aresponses_mock.call_count == expected_cap # Every follow-up before the cap still offered tools; only the capped one drops them. for call in aresponses_mock.call_args_list[:-1]: @@ -193,6 +202,12 @@ async def test_tool_call_rounds_are_capped(monkeypatch): assert "tools" not in aresponses_mock.call_args_list[-1].kwargs +@pytest.mark.parametrize(("max_agentic_loops", "expected_error"), ((0, ValueError), ("three", TypeError))) +def test_invalid_max_agentic_loops_raises_when_creating_iterator(max_agentic_loops, expected_error): + with pytest.raises(expected_error, match=r"litellm_params\.max_agentic_loops"): + _make_iterator([], max_agentic_loops=max_agentic_loops) + + @pytest.mark.asyncio async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch): """ diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 63204114501..278bd4c1e52 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -10,6 +10,7 @@ import copy import json from collections.abc import Mapping from importlib import import_module +from pathlib import Path from typing import Final, cast from unittest.mock import MagicMock, patch @@ -852,3 +853,92 @@ def test_bridge_reasoning_effort_reads_only_a_string_or_mapping_effort_nested_in extra_body: JsonValue, expected: JsonValue ) -> None: assert bridge_reasoning_effort(None, extra_body) == expected + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("prefix", ["", "responses/"]) +@pytest.mark.parametrize("catalog_mode", [None, "chat", "responses"]) +@pytest.mark.asyncio +async def test_chatgpt_completion_routes_to_responses_without_a_provider_catalog_entry( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + respx_mock: respx.MockRouter, + asynchronous: bool, + stream: bool, + prefix: str, + catalog_mode: str | None, +) -> None: + model: Final = "test-chatgpt-routing" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("CHATGPT_API_BASE", "https://chatgpt.example") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "fake-access-token", "expires_at": 253402300799, "account_id": "fake-account"}) + ) + monkeypatch.delitem(litellm.model_cost, f"chatgpt/{model}", raising=False) + monkeypatch.delitem(litellm.model_cost, model, raising=False) + if catalog_mode is not None: + monkeypatch.setitem(litellm.model_cost, model, {"mode": catalog_mode, "litellm_provider": "openai"}) + litellm.get_model_info.cache_clear() + payload: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [ + { + "type": "message", + "id": "msg_test", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + events: Final = ( + {"type": "response.created", "response": {**payload, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "delta": "hello", + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + }, + {"type": "response.completed", "response": payload}, + ) + upstream: Final = respx_mock.post("https://chatgpt.example/responses").respond( + 200, + headers={"content-type": "text/event-stream"}, + text="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events), + ) + kwargs: Final = { + "model": f"chatgpt/{prefix}{model}", + "messages": [{"role": "user", "content": "Say hello"}], + "api_base": "https://chatgpt.example", + "stream": stream, + "num_retries": 0, + } + response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs) + result: Final = ( + litellm.stream_chunk_builder([chunk async for chunk in response] if asynchronous else list(response)) + if stream + else response + ) + assert result is not None + assert result.choices[0].message.content == "hello" + assert upstream.call_count == 1 + sent: Final = json.loads(upstream.calls[0].request.content) + assert sent["model"] == model + assert sent["store"] is False + assert sent["input"] == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Say hello"}]} + ] + + +@pytest.mark.parametrize("provider", ["openai", "azure", "xai", "github_copilot"]) +def test_missing_catalog_entry_does_not_force_other_providers_to_responses(provider: str) -> None: + info, model = responses_api_bridge_check(model="test-chatgpt-routing-control", custom_llm_provider=provider) + assert (info.get("mode"), model) == (None, "test-chatgpt-routing-control") diff --git a/tests/unit/test_cerebras_llama_model_metadata.py b/tests/unit/test_cerebras_llama_model_metadata.py new file mode 100644 index 00000000000..c5ec096e553 --- /dev/null +++ b/tests/unit/test_cerebras_llama_model_metadata.py @@ -0,0 +1,115 @@ +# What does this test? +""" +Verifies that Cerebras Llama 3.1 models (standard hyphenated aliases): +- cerebras/llama-3.1-8b +- cerebras/llama-3.1-70b + +have correct rate card metadata, context window limits, and cost calculation parity +with the un-hyphenated variants in model_prices_and_context_window.json. +Fixes: https://github.com/BerriAI/litellm/issues/48157 +""" + +import json +from pathlib import Path + +import pytest + +import litellm + + +@pytest.fixture(scope="module") +def model_cost_maps(): + root = Path(__file__).resolve().parents[2] + main_path = root / "model_prices_and_context_window.json" + backup_path = root / "litellm" / "model_prices_and_context_window_backup.json" + + with open(main_path, "r", encoding="utf-8") as f: + main_map = json.load(f) + with open(backup_path, "r", encoding="utf-8") as f: + backup_map = json.load(f) + + return main_map, backup_map + + +CEREBRAS_MODELS_EXPECTED = { + "cerebras/llama-3.1-8b": { + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "mode": "chat", + "litellm_provider": "cerebras", + }, + "cerebras/llama-3.1-70b": { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "mode": "chat", + "litellm_provider": "cerebras", + }, +} + +PAIR_PARITY = [ + ("cerebras/llama-3.1-8b", "cerebras/llama3.1-8b"), + ("cerebras/llama-3.1-70b", "cerebras/llama3.1-70b"), +] + + +@pytest.mark.parametrize("model_id,expected_data", CEREBRAS_MODELS_EXPECTED.items()) +def test_cerebras_llama_model_metadata(model_cost_maps, model_id, expected_data): + """Test that hyphenated Cerebras Llama 3.1 models exist and have expected metadata in both catalogs.""" + main_map, backup_map = model_cost_maps + + assert model_id in main_map, f"Missing {model_id} in main catalog" + assert model_id in backup_map, f"Missing {model_id} in backup catalog" + + main_entry = main_map[model_id] + backup_entry = backup_map[model_id] + assert main_entry == backup_entry, f"Parity mismatch between main and backup for {model_id}" + + for key, expected_val in expected_data.items(): + assert main_entry.get(key) == expected_val, ( + f"{model_id}[{key}] got {main_entry.get(key)}, expected {expected_val}" + ) + + +@pytest.mark.parametrize("hyphenated,unhyphenated", PAIR_PARITY) +def test_cerebras_llama_parity_with_unhyphenated(model_cost_maps, hyphenated, unhyphenated): + """Test that hyphenated and un-hyphenated variants have identical pricing and context limits.""" + main_map, _ = model_cost_maps + + assert hyphenated in main_map + assert unhyphenated in main_map + + entry_hyphen = main_map[hyphenated] + entry_unhyphen = main_map[unhyphenated] + + assert entry_hyphen["input_cost_per_token"] == entry_unhyphen["input_cost_per_token"] + assert entry_hyphen["output_cost_per_token"] == entry_unhyphen["output_cost_per_token"] + assert entry_hyphen["max_input_tokens"] == entry_unhyphen["max_input_tokens"] + assert entry_hyphen["max_output_tokens"] == entry_unhyphen["max_output_tokens"] + + +def test_cerebras_llama_cost_per_token(): + """Test that litellm.cost_per_token accurately computes costs for hyphenated Cerebras models.""" + prompt_tokens = 1_000_000 + completion_tokens = 1_000_000 + + # llama-3.1-8b: $0.10 input + $0.10 output per 1M tokens + cost_prompt, cost_comp = litellm.cost_per_token( + model="cerebras/llama-3.1-8b", + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + assert pytest.approx(cost_prompt, rel=1e-5) == 0.10 + assert pytest.approx(cost_comp, rel=1e-5) == 0.10 + + # llama-3.1-70b: $0.60 input + $0.60 output per 1M tokens + cost_prompt_70b, cost_comp_70b = litellm.cost_per_token( + model="cerebras/llama-3.1-70b", + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + assert pytest.approx(cost_prompt_70b, rel=1e-5) == 0.60 + assert pytest.approx(cost_comp_70b, rel=1e-5) == 0.60 diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 550fe5f0d1a..27b25c1f98e 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -35,7 +35,7 @@ from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIRealtimeStreamList, ResponseAPIUsage, ResponsesAPIResponse from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse -from litellm.types.utils import CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage +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 supports_prompt_caching import os @@ -5947,6 +5947,110 @@ def test_batch_cost_calculator_bills_cache_write_tokens_at_the_batch_input_rate_ assert prompt_cost == pytest.approx(300_048 * 3e-6) +_TWELVE_HOUR_BATCH_MODEL_INFO: Final = cast( + ModelInfo, + { + "input_cost_per_token_batches": 1e-6, + "input_cost_per_token_above_272k_tokens_batches": 3e-6, + "cache_creation_input_token_cost_batches": 1.25e-6, + "cache_creation_input_token_cost_above_272k_tokens_batches": 3.75e-6, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 7.5e-6, + "output_cost_per_token_batches": 4e-6, + }, +) + + +def _twelve_hour_cache_write_usage(prompt_tokens: int) -> Usage: + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=4, + total_tokens=prompt_tokens + 4, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=0, + cache_creation_tokens=500, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=300, ephemeral_12h_input_tokens=200 + ), + ), + ) + + +@pytest.mark.parametrize( + ("prompt_tokens", "input_rate", "write_rate", "twelve_hour_rate"), + [(1_000, 1e-6, 1.25e-6, 2.5e-6), (272_000, 1e-6, 1.25e-6, 2.5e-6), (300_000, 3e-6, 3.75e-6, 7.5e-6)], +) +def test_batch_cost_calculator_bills_12h_cache_writes_at_the_12h_batch_rate_for_the_crossed_tier( + prompt_tokens: int, input_rate: float, write_rate: float, twelve_hour_rate: float +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + prompt_cost, completion_cost_value = batch_cost_calculator( + usage=_twelve_hour_cache_write_usage(prompt_tokens), + model="lit-12h-batch-model", + custom_llm_provider="openai", + model_info=_TWELVE_HOUR_BATCH_MODEL_INFO, + ) + + assert prompt_cost == pytest.approx((prompt_tokens - 500) * input_rate + 300 * write_rate + 200 * twelve_hour_rate) + assert completion_cost_value == pytest.approx(4 * 4e-6) + + +@pytest.mark.parametrize( + ("prompt_tokens", "input_rate", "twelve_hour_write_rate"), [(1_000, 1e-6, 2.5e-6), (300_000, 3e-6, 7.5e-6)] +) +def test_batch_cost_calculator_bills_12h_cache_writes_once_when_only_the_12h_batch_key_exists( + prompt_tokens: int, input_rate: float, twelve_hour_write_rate: float +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + model_info: Final = cast( + ModelInfo, + { + key: rate + for key, rate in _TWELVE_HOUR_BATCH_MODEL_INFO.items() + if not key.startswith("cache_creation_input_token_cost_") + or key.startswith("cache_creation_input_token_cost_above_12hr") + }, + ) + + prompt_cost, _ = batch_cost_calculator( + usage=_twelve_hour_cache_write_usage(prompt_tokens), + model="lit-12h-batch-model", + custom_llm_provider="openai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx((prompt_tokens - 200) * input_rate + 200 * twelve_hour_write_rate) + + +@pytest.mark.parametrize(("prompt_tokens", "write_rate"), [(1_000, 1.25e-6), (300_000, 3.75e-6)]) +def test_batch_cost_calculator_bills_12h_cache_writes_at_the_batch_cache_write_rate_without_a_12h_batch_key( + prompt_tokens: int, write_rate: float +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + model_info: Final = cast( + ModelInfo, + { + key: rate + for key, rate in _TWELVE_HOUR_BATCH_MODEL_INFO.items() + if not key.startswith("cache_creation_input_token_cost_above_12hr") + }, + ) + + prompt_cost, _ = batch_cost_calculator( + usage=_twelve_hour_cache_write_usage(prompt_tokens), + model="lit-12h-batch-model", + custom_llm_provider="openai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + (prompt_tokens - 500) * (1e-6 if prompt_tokens <= 272_000 else 3e-6) + 500 * write_rate + ) + + def test_batch_cost_calculator_prices_modalities_and_cached_tokens_together_in_the_crossed_tier() -> None: from litellm.cost_calculator import batch_cost_calculator diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 9b9fec4c9a1..c0abb18ca8c 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -455,30 +455,23 @@ def is_active_priced_mistral_chat_row(name: str, entry: Mapping[str, object]) -> ) -def cache_read_is_tenth_of_input(entry: Mapping[str, object]) -> bool: +def cache_read_is_below_input(entry: Mapping[str, object]) -> bool: cache_read: Final = entry.get("cache_read_input_token_cost") input_cost: Final = entry.get("input_cost_per_token") - return ( - isinstance(cache_read, float) - and isinstance(input_cost, (int, float)) - and 0 < cache_read < input_cost - and cache_read == pytest.approx(input_cost / 10) - ) + return isinstance(cache_read, float) and isinstance(input_cost, (int, float)) and 0 < cache_read < input_cost @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) def test_active_mistral_chat_rows_price_cache_reads_below_input(path: Path): """A Mistral chat row without a cache-read rate bills cached prompt tokens at zero, so every - active priced row must carry one, and it must be cheaper than a fresh input token. Mistral - bills cached tokens at 10% of the input price for every model (docs.mistral.ai/studio/ - conversations/advanced/prompt-caching, read 2026-09-18), so the ratio is checked as well.""" + active priced row must carry one, and it must be cheaper than a fresh input token.""" rows: Mapping[str, object] = json.loads(path.read_text()) drifted: Final = [ f"{name}: cache_read={entry.get('cache_read_input_token_cost')} input={entry.get('input_cost_per_token')}" for name, entry in rows.items() if isinstance(entry, dict) and is_active_priced_mistral_chat_row(name, entry) - and not cache_read_is_tenth_of_input(entry) + and not cache_read_is_below_input(entry) ] assert drifted == [] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 2288d4311a9..568a05e8ff4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -30022,64 +30022,109 @@ _DECLARED_TEXT_LEG: Final[dict[str, object]] = { } -def _pre_call_router(*deployments: dict[str, object], routing_strategy: str = "latency-based-routing") -> litellm.Router: +def _mixed_router(*deployments: dict[str, object], routing_strategy: str = "latency-based-routing") -> litellm.Router: return litellm.Router( model_list=[{"model_name": "mixed", **deployment} for deployment in deployments], routing_strategy=routing_strategy, - enable_pre_call_checks=True, ) -def _pre_call_ids(router: litellm.Router, deployments: tuple[dict[str, object], ...], **request: object) -> list[str]: - return [d["model_info"]["id"] for d in router._pre_call_checks(model="mixed", healthy_deployments=list(deployments), **request)] +def _candidate_ids( + router: litellm.Router, + messages: list[dict[str, object]] | None = None, + request_kwargs: dict[str, object] | None = None, +) -> list[str]: + candidates: Final = router._without_declared_text_only_deployments( + router.get_model_list(model_name="mixed"), "mixed", messages, request_kwargs + ) + return [d["model_info"]["id"] for d in candidates] -def test_pre_call_checks_drop_declared_text_only_legs_for_image_requests(): - router = _pre_call_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _pre_call_ids(router, (_VISION_LEG, _DECLARED_TEXT_LEG), messages=_IMAGE_CHAT) == ["vision"] +def test_router_picks_the_vision_leg_for_image_requests_without_enable_pre_call_checks(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert router.enable_pre_call_checks is False + for _ in range(5): + assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "vision" -def test_pre_call_checks_keep_declared_text_only_legs_for_text_requests(): - router = _pre_call_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _pre_call_ids(router, (_VISION_LEG, _DECLARED_TEXT_LEG), messages=[{"role": "user", "content": "Say hi"}]) == [ +@pytest.mark.asyncio +async def test_async_router_picks_the_vision_leg_for_responses_images_without_enable_pre_call_checks(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + for _ in range(5): + deployment: Final = await router.async_get_available_deployment( + model="mixed", request_kwargs={"input": _IMAGE_RESPONSES_INPUT} + ) + assert deployment["model_info"]["id"] == "vision" + + +def test_router_sends_an_image_to_the_text_only_leg_when_pre_call_checks_drop_the_vision_leg(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "mixed", **_VISION_LEG, "model_info": {"id": "vision", "max_input_tokens": 1}}, + {"model_name": "mixed", **_DECLARED_TEXT_LEG}, + ], + routing_strategy="latency-based-routing", + enable_pre_call_checks=True, + ) + assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "declared-text" + + +def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooling_down(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + router.cooldown_cache.add_deployment_to_cooldown( + model_id="vision", original_exception=Exception("bad key"), exception_status=401, cooldown_time=60 + ) + assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "declared-text" + + +def test_without_declared_text_only_deployments_passes_a_chosen_deployment_through(): + assert litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) is _DECLARED_TEXT_LEG + + +def test_image_routing_drop_declared_text_only_legs_for_image_requests(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] + + +def test_image_routing_keep_declared_text_only_legs_for_text_requests(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert _candidate_ids(router, messages=[{"role": "user", "content": "Say hi"}]) == [ "vision", "declared-text", ] -def test_pre_call_checks_drop_declared_text_only_legs_for_responses_input_image(): - router = _pre_call_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _pre_call_ids( - router, (_VISION_LEG, _DECLARED_TEXT_LEG), messages=None, request_kwargs={"input": _IMAGE_RESPONSES_INPUT} - ) == ["vision"] +def test_image_routing_drop_declared_text_only_legs_for_responses_input_image(): + router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) + assert _candidate_ids(router, request_kwargs={"input": _IMAGE_RESPONSES_INPUT}) == ["vision"] -def test_pre_call_checks_send_image_to_the_group_when_every_leg_is_declared_text_only(): +def test_image_routing_send_image_to_the_group_when_every_leg_is_declared_text_only(): other_text_leg = {**_DECLARED_TEXT_LEG, "model_info": {"id": "declared-text-2", "supports_vision": False}} - router = _pre_call_router(_DECLARED_TEXT_LEG, other_text_leg) - assert _pre_call_ids(router, (_DECLARED_TEXT_LEG, other_text_leg), messages=_IMAGE_CHAT) == [ + router: Final = _mixed_router(_DECLARED_TEXT_LEG, other_text_leg) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == [ "declared-text", "declared-text-2", ] -def test_pre_call_checks_keep_undeclared_custom_legs_for_image_requests(): +def test_image_routing_keep_undeclared_custom_legs_for_image_requests(): custom_leg = { "litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct", "api_base": "http://localhost:8000"}, "model_info": {"id": "custom"}, } - router = _pre_call_router(_VISION_LEG, custom_leg) - assert _pre_call_ids(router, (_VISION_LEG, custom_leg), messages=_IMAGE_CHAT) == ["vision", "custom"] + router: Final = _mixed_router(_VISION_LEG, custom_leg) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision", "custom"] -def test_pre_call_checks_resolve_vision_through_the_cost_map_and_base_model(): +def test_image_routing_resolve_vision_through_the_cost_map_and_base_model(): catalog_text_leg = {"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {"id": "catalog-text"}} azure_text_leg = { "litellm_params": {"model": "azure/my-audio-deployment", "api_base": "https://x.openai.azure.com"}, "model_info": {"id": "azure-text", "base_model": "azure/gpt-audio"}, } - router = _pre_call_router(_VISION_LEG, catalog_text_leg, azure_text_leg) - assert _pre_call_ids(router, (_VISION_LEG, catalog_text_leg, azure_text_leg), messages=_IMAGE_CHAT) == ["vision"] + router: Final = _mixed_router(_VISION_LEG, catalog_text_leg, azure_text_leg) + assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] @pytest.mark.parametrize( diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a0b4b784a4d..55b7bc78a11 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -21,7 +21,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest import respx -from jsonschema import validate +from jsonschema import Draft7Validator, validate import litellm from litellm.exceptions import MidStreamFallbackError @@ -923,6 +923,456 @@ def validate_model_cost_values(model_data, exceptions=None): return len(violations) == 0, violations +_COST_MAP_ROW_SCHEMA: Final = { + "type": "object", + "additionalProperties": { + "type": "object", + "properties": { + "supports_computer_use": {"type": "boolean"}, + "cache_creation_input_audio_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost_above_1hr": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": { + "type": "number" + }, + "cache_creation_input_token_cost_above_12hr_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_ultrafast": {"type": "number"}, + "cache_creation_input_token_cost_above_32k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_100k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_100k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_128k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_200k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_256k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_batches": {"type": "number"}, + "cache_creation_input_token_cost_decisions": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "cache_read_input_token_cost_decisions": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, + "input_cost_per_token_decisions": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "output_cost_per_token_decisions": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, + "cache_creation_input_token_cost_flex": {"type": "number"}, + "cache_creation_input_token_cost_priority": {"type": "number"}, + "cache_creation_input_token_cost_ultrafast": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "explicit_cache_read_input_token_cost": {"type": "number"}, + "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_100k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_100k_tokens_batches": {"type": "number"}, + "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens_batches": {"type": "number"}, + "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, + "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, + "cache_read_input_token_cost_batches": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, + "cache_read_input_audio_token_cost": {"type": "number"}, + "cache_read_input_audio_token_cost_batches": {"type": "number"}, + "cache_read_input_audio_token_cost_flex": {"type": "number"}, + "cache_read_input_audio_token_cost_priority": {"type": "number"}, + "cache_read_input_image_token_cost": {"type": "number"}, + "audio_transcription_config": {"type": "string"}, + "deprecation_date": {"type": "string"}, + "input_cost_per_audio_per_second": {"type": "number"}, + "input_cost_per_audio_per_second_above_128k_tokens": {"type": "number"}, + "google_maps_grounding_cost_per_query": {"type": "number"}, + "input_cost_per_audio_token": {"type": "number"}, + "input_cost_per_audio_token_flex": {"type": "number"}, + "input_cost_per_image_token": {"type": "number"}, + "input_cost_per_character": {"type": "number"}, + "input_cost_per_character_above_128k_tokens": {"type": "number"}, + "input_cost_per_image": {"type": "number"}, + "input_cost_per_image_above_128k_tokens": {"type": "number"}, + "input_cost_per_video_token": {"type": "number"}, + "input_cost_per_token_above_32k_tokens": {"type": "number"}, + "input_cost_per_token_above_100k_tokens": {"type": "number"}, + "input_cost_per_token_above_100k_tokens_batches": {"type": "number"}, + "input_cost_per_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_batches": {"type": "number"}, + "input_cost_per_token_above_256k_tokens": {"type": "number"}, + "input_cost_per_token_above_272k_tokens": {"type": "number"}, + "input_cost_per_token_above_512k_tokens": {"type": "number"}, + "cache_read_input_token_cost_flex": {"type": "number"}, + "cache_read_input_token_cost_priority": {"type": "number"}, + "cache_read_input_token_cost_balanced": {"type": "number"}, + "cache_read_input_token_cost_ultrafast": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens_priority": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_priority": {"type": "number"}, + "input_cost_per_token_flex": {"type": "number"}, + "input_cost_per_token_priority": {"type": "number"}, + "input_cost_per_token_balanced": {"type": "number"}, + "input_cost_per_token_ultrafast": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_priority": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_batches": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_flex": {"type": "number"}, + "input_cost_per_audio_token_priority": {"type": "number"}, + "output_cost_per_token_flex": {"type": "number"}, + "output_cost_per_token_priority": {"type": "number"}, + "output_cost_per_token_balanced": {"type": "number"}, + "output_cost_per_token_ultrafast": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_batches": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_flex": {"type": "number"}, + "regional_endpoint_uplift_multiplier": {"type": "number"}, + "regional_processing_uplift_multiplier_eu": {"type": "number"}, + "regional_processing_uplift_multiplier_us": {"type": "number"}, + "input_cost_per_pixel": {"type": "number"}, + "input_cost_per_query": {"type": "number"}, + "input_cost_per_request": {"type": "number"}, + "cost_per_second": {"type": "number"}, + "input_cost_per_second": {"type": "number"}, + "input_cost_per_text_message": {"type": "number"}, + "input_cost_per_token": {"type": "number"}, + "input_cost_per_token_above_128k_tokens": {"type": "number"}, + "input_cost_per_audio_token_batches": {"type": "number"}, + "input_cost_per_image_token_batches": {"type": "number"}, + "input_cost_per_token_batches": {"type": "number"}, + "input_cost_per_video_token_batches": {"type": "number"}, + "input_cost_per_token_cache_hit": {"type": "number"}, + "input_cost_per_video_per_second": {"type": "number"}, + "input_cost_per_video_per_second_above_8s_interval": {"type": "number"}, + "input_cost_per_video_per_second_above_15s_interval": {"type": "number"}, + "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, + "input_dbu_cost_per_token": {"type": "number"}, + "annotation_cost_per_page": {"type": "number"}, + "annotation_cost_per_page_batches": {"type": "number"}, + "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_page_batches": {"type": "number"}, + "ocr_cost_per_credit": {"type": "number"}, + "code_interpreter_cost_per_session": {"type": "number"}, + "inference_geo": {"type": "string"}, + "litellm_provider": {"type": "string"}, + "max_input_tokens": {"type": "number"}, + "max_output_tokens": {"type": "number"}, + "max_tokens": {"type": "number"}, + "metadata": {"type": "object"}, + "provider_specific_entry": {"type": "object"}, + "mode": { + "type": "string", + "enum": [ + "audio_speech", + "audio_transcription", + "chat", + "completion", + "container", + "image_edit", + "embedding", + "evaluation", + "extraction", + "guardrail", + "image_generation", + "video_generation", + "moderation", + "rerank", + "realtime", + "responses", + "ocr", + "search", + "vector_store", + ], + }, + "output_cost_per_audio_token": {"type": "number"}, + "output_cost_per_character": {"type": "number"}, + "output_cost_per_character_above_128k_tokens": {"type": "number"}, + "output_cost_per_image": {"type": "number"}, + "output_cost_per_image_512": {"type": "number"}, + "output_cost_per_image_1024": {"type": "number"}, + "output_cost_per_image_1536": {"type": "number"}, + "output_cost_per_image_0.5K": {"type": "number"}, + "output_cost_per_image_1K": {"type": "number"}, + "output_cost_per_image_2K": {"type": "number"}, + "output_cost_per_image_4K": {"type": "number"}, + "output_cost_per_image_token": {"type": "number"}, + "output_cost_per_image_token_batches": {"type": "number"}, + "output_cost_per_video_token": {"type": "number"}, + "output_cost_per_pixel": {"type": "number"}, + "output_cost_per_second": {"type": "number"}, + "output_cost_per_second_480p": {"type": "number"}, + "output_cost_per_second_720p": {"type": "number"}, + "output_cost_per_second_768p": {"type": "number"}, + "output_cost_per_second_2k": {"type": "number"}, + "output_cost_per_second_1080p": {"type": "number"}, + "output_cost_per_second_4k": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "output_cost_per_token_above_32k_tokens": {"type": "number"}, + "output_cost_per_token_above_100k_tokens": {"type": "number"}, + "output_cost_per_token_above_100k_tokens_batches": {"type": "number"}, + "output_cost_per_token_above_128k_tokens": {"type": "number"}, + "output_cost_per_token_above_200k_tokens": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_batches": {"type": "number"}, + "output_cost_per_token_above_256k_tokens": {"type": "number"}, + "output_cost_per_token_above_272k_tokens": {"type": "number"}, + "output_cost_per_token_above_512k_tokens": {"type": "number"}, + "output_cost_per_token_batches": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "output_cost_per_video_per_second": {"type": "number"}, + "output_db_cost_per_token": {"type": "number"}, + "output_dbu_cost_per_token": {"type": "number"}, + "output_vector_size": {"type": "number"}, + "rpd": {"type": "number"}, + "rpm": {"type": "number"}, + "source": {"type": "string"}, + "comment": {"type": "string"}, + "supports_assistant_prefill": {"type": "boolean"}, + "supports_anthropic_compaction": {"type": "boolean"}, + "supports_audio_input": {"type": "boolean"}, + "supports_audio_output": {"type": "boolean"}, + "gemini_native_audio": {"type": "boolean"}, + "gemini_audio_only_live": {"type": "boolean"}, + "supports_embedding_image_input": {"type": "boolean"}, + "supports_forced_tool_use": {"type": "boolean"}, + "supports_function_calling": {"type": "boolean"}, + "supports_image_input": {"type": "boolean"}, + "supports_nova_canvas_image_edit": {"type": "boolean"}, + "supports_parallel_function_calling": {"type": "boolean"}, + "supports_parallel_tool_use_config": {"type": "boolean"}, + "supports_pdf_input": {"type": "boolean"}, + "prompt_cache_min_tokens": {"type": "number"}, + "supports_prompt_cache_breakpoint": {"type": "boolean"}, + "supports_thinking_cache_preservation": {"type": "boolean"}, + "supports_prompt_caching": {"type": "boolean"}, + "supports_response_schema": {"type": "boolean"}, + "supports_system_messages": {"type": "boolean"}, + "supports_tool_choice": {"type": "boolean"}, + "supports_tool_search": {"type": "boolean"}, + "supports_video_input": {"type": "boolean"}, + "supports_box_questions": {"type": "boolean"}, + "supports_vision": {"type": "boolean"}, + "supports_web_search": {"type": "boolean"}, + "token_count_framing": {"type": "integer"}, + "token_count_model": { + "type": "object", + "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, + "required": ["model"], + }, + "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, + "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, + "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, + "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, + "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, + "supports_url_context": {"type": "boolean"}, + "supports_multimodal": {"type": "boolean"}, + "uses_embed_content": {"type": "boolean"}, + "supports_reasoning": {"type": "boolean"}, + "supports_minimal_reasoning_effort": {"type": "boolean"}, + "supports_low_reasoning_effort": {"type": "boolean"}, + "supports_none_reasoning_effort": {"type": "boolean"}, + "supports_xhigh_reasoning_effort": {"type": "boolean"}, + "supports_max_reasoning_effort": {"type": "boolean"}, + "reasoning_effort_levels": { + "type": "array", + "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, + }, + "default_reasoning_effort": { + "type": "string", + "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], + }, + "supports_adaptive_thinking": {"type": "boolean"}, + "supports_anthropic_thinking_payload": {"type": "boolean"}, + "supports_legacy_thinking": {"type": "boolean"}, + "thinking_always_on": {"type": "boolean"}, + "supports_mid_conversation_system": {"type": "boolean"}, + "supports_sampling_params": {"type": "boolean"}, + "supports_output_config": {"type": "boolean"}, + "supports_speed": {"type": "boolean"}, + "supports_fast_mode": {"type": "boolean"}, + "supported_audio_formats": { + "type": "array", + "items": { + "type": "string", + "enum": ["mp3", "wav"], + }, + }, + "vertex_ai_audio_api": { + "type": "string", + "enum": ["lyria_predict", "lyria_interactions"], + }, + "bedrock_output_config_effort_ceiling": { + "type": "string", + "enum": ["low", "medium", "high", "max", "xhigh"], + }, + "bedrock_converse_supports_strict_tools": {"type": "boolean"}, + "supports_regex_lookaround": {"type": "boolean"}, + "tpm": {"type": "number"}, + "supported_endpoints": { + "type": "array", + "items": { + "type": "string", + "enum": [ + "/v1/responses", + "/v1/embeddings", + "/v1/chat/completions", + "/v1/completions", + "/v1/messages", + "/v1/images/generations", + "/v1/realtime", + "/v1/realtime/transcription_sessions", + "/v1/images/variations", + "/v1/images/edits", + "/v1/batch", + "/v1beta/interactions", + "/v1beta/models", + "/v1/audio/transcriptions", + "/v1/audio/speech", + "/v1/ocr", + "/v1/videos", + "/vertex_ai/live", + "/v1/listen", + "/v1/systemone", + "/v1/decisions", + "/v1beta/interactions", + ], + }, + }, + "supported_regions": { + "type": "array", + "items": { + "type": "string", + }, + }, + "guardrail_cost_per_unit": { + "type": "object", + "additionalProperties": {"type": "number"}, + }, + "search_context_cost_per_query": { + "type": "object", + "properties": { + "search_context_size_low": {"type": "number"}, + "search_context_size_medium": {"type": "number"}, + "search_context_size_high": {"type": "number"}, + }, + "additionalProperties": False, + }, + "web_search_billing_unit": { + "type": "string", + "enum": ["per_prompt", "per_query"], + }, + "search_free_requests": { + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "period": {"type": "string", "enum": ["day", "month"]}, + "shared_group": {"type": "string"}, + }, + "required": ["count", "period", "shared_group"], + "additionalProperties": False, + }, + "citation_cost_per_token": {"type": "number"}, + "supported_modalities": { + "type": "array", + "items": { + "type": "string", + "enum": ["text", "audio", "image", "video"], + }, + }, + "supported_output_modalities": { + "type": "array", + "items": { + "type": "string", + "enum": ["text", "image", "audio", "code", "video"], + }, + }, + "supports_native_streaming": {"type": "boolean"}, + "supports_image_size": {"type": "boolean"}, + "supports_native_structured_output": {"type": "boolean"}, + "supports_bedrock_native_structured_output": {"type": "boolean"}, + "use_openai_responses_path": {"type": "boolean"}, + "off_peak_pricing": { + "type": "object", + "properties": { + "hours_utc": { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + }, + "windows": { + "type": "array", + "items": { + "type": "object", + "properties": { + "hours_utc": { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + }, + "weekdays": { + "type": "array", + "items": {"oneOf": [{"type": "integer"}, {"type": "string"}]}, + }, + }, + "required": ["hours_utc"], + "additionalProperties": False, + }, + }, + "weekday_timezone": {"type": "string"}, + "input_cost_per_token": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + }, + "additionalProperties": False, + }, + "tiered_pricing": { + "type": "array", + "items": { + "type": "object", + "properties": { + "range": { + "type": "array", + "items": {"type": "number"}, + "minItems": 2, + "maxItems": 2, + }, + "input_cost_per_token": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "max_results_range": { + "type": "array", + "items": {"type": "number"}, + "minItems": 2, + "maxItems": 2, + }, + "input_cost_per_query": {"type": "number"}, + }, + "additionalProperties": False, + }, + }, + }, + "additionalProperties": False, + }, +} + + def test_aaamodel_prices_and_context_window_json_is_valid(): """ Validates the `model_prices_and_context_window.json` file. @@ -930,447 +1380,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): If this test fails after you update the json, you need to update the schema or correct the change you made. """ - INTENDED_SCHEMA = { - "type": "object", - "additionalProperties": { - "type": "object", - "properties": { - "supports_computer_use": {"type": "boolean"}, - "cache_creation_input_audio_token_cost": {"type": "number"}, - "cache_creation_input_token_cost": {"type": "number"}, - "cache_creation_input_token_cost_above_1hr": {"type": "number"}, - "cache_creation_input_token_cost_above_12hr": {"type": "number"}, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_flex": {"type": "number"}, - "cache_creation_input_token_cost_above_12hr_priority": {"type": "number"}, - "cache_creation_input_token_cost_above_32k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_100k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_100k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_above_128k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_200k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_256k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_flex": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_priority": {"type": "number"}, - "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_batches": {"type": "number"}, - "cache_creation_input_token_cost_decisions": {"type": "number"}, - "cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, - "cache_read_input_token_cost_decisions": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"}, - "input_cost_per_token_decisions": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, - "output_cost_per_token_decisions": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_decisions": {"type": "number"}, - "cache_creation_input_token_cost_flex": {"type": "number"}, - "cache_creation_input_token_cost_priority": {"type": "number"}, - "cache_creation_input_token_cost_ultrafast": {"type": "number"}, - "cache_read_input_token_cost": {"type": "number"}, - "explicit_cache_read_input_token_cost": {"type": "number"}, - "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_100k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_100k_tokens_batches": {"type": "number"}, - "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_200k_tokens_batches": {"type": "number"}, - "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, - "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, - "cache_read_input_token_cost_batches": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_batches": {"type": "number"}, - "cache_creation_input_token_cost_above_1hr_above_100k_tokens": {"type": "number"}, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, - "cache_read_input_audio_token_cost": {"type": "number"}, - "cache_read_input_audio_token_cost_batches": {"type": "number"}, - "cache_read_input_audio_token_cost_flex": {"type": "number"}, - "cache_read_input_audio_token_cost_priority": {"type": "number"}, - "cache_read_input_image_token_cost": {"type": "number"}, - "audio_transcription_config": {"type": "string"}, - "deprecation_date": {"type": "string"}, - "input_cost_per_audio_per_second": {"type": "number"}, - "input_cost_per_audio_per_second_above_128k_tokens": {"type": "number"}, - "google_maps_grounding_cost_per_query": {"type": "number"}, - "input_cost_per_audio_token": {"type": "number"}, - "input_cost_per_audio_token_flex": {"type": "number"}, - "input_cost_per_image_token": {"type": "number"}, - "input_cost_per_character": {"type": "number"}, - "input_cost_per_character_above_128k_tokens": {"type": "number"}, - "input_cost_per_image": {"type": "number"}, - "input_cost_per_image_above_128k_tokens": {"type": "number"}, - "input_cost_per_video_token": {"type": "number"}, - "input_cost_per_token_above_32k_tokens": {"type": "number"}, - "input_cost_per_token_above_100k_tokens": {"type": "number"}, - "input_cost_per_token_above_100k_tokens_batches": {"type": "number"}, - "input_cost_per_token_above_200k_tokens": {"type": "number"}, - "input_cost_per_token_above_200k_tokens_batches": {"type": "number"}, - "input_cost_per_token_above_256k_tokens": {"type": "number"}, - "input_cost_per_token_above_272k_tokens": {"type": "number"}, - "input_cost_per_token_above_512k_tokens": {"type": "number"}, - "cache_read_input_token_cost_flex": {"type": "number"}, - "cache_read_input_token_cost_priority": {"type": "number"}, - "cache_read_input_token_cost_balanced": {"type": "number"}, - "cache_read_input_token_cost_ultrafast": {"type": "number"}, - "cache_read_input_token_cost_above_200k_tokens_priority": {"type": "number"}, - "cache_read_input_token_cost_above_272k_tokens_priority": {"type": "number"}, - "input_cost_per_token_flex": {"type": "number"}, - "input_cost_per_token_priority": {"type": "number"}, - "input_cost_per_token_balanced": {"type": "number"}, - "input_cost_per_token_ultrafast": {"type": "number"}, - "input_cost_per_token_above_200k_tokens_priority": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_priority": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_batches": {"type": "number"}, - "input_cost_per_token_above_272k_tokens_flex": {"type": "number"}, - "input_cost_per_audio_token_priority": {"type": "number"}, - "output_cost_per_token_flex": {"type": "number"}, - "output_cost_per_token_priority": {"type": "number"}, - "output_cost_per_token_balanced": {"type": "number"}, - "output_cost_per_token_ultrafast": {"type": "number"}, - "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_batches": {"type": "number"}, - "output_cost_per_token_above_272k_tokens_flex": {"type": "number"}, - "regional_endpoint_uplift_multiplier": {"type": "number"}, - "regional_processing_uplift_multiplier_eu": {"type": "number"}, - "regional_processing_uplift_multiplier_us": {"type": "number"}, - "input_cost_per_pixel": {"type": "number"}, - "input_cost_per_query": {"type": "number"}, - "input_cost_per_request": {"type": "number"}, - "cost_per_second": {"type": "number"}, - "input_cost_per_second": {"type": "number"}, - "input_cost_per_text_message": {"type": "number"}, - "input_cost_per_token": {"type": "number"}, - "input_cost_per_token_above_128k_tokens": {"type": "number"}, - "input_cost_per_audio_token_batches": {"type": "number"}, - "input_cost_per_image_token_batches": {"type": "number"}, - "input_cost_per_token_batches": {"type": "number"}, - "input_cost_per_video_token_batches": {"type": "number"}, - "input_cost_per_token_cache_hit": {"type": "number"}, - "input_cost_per_video_per_second": {"type": "number"}, - "input_cost_per_video_per_second_above_8s_interval": {"type": "number"}, - "input_cost_per_video_per_second_above_15s_interval": {"type": "number"}, - "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, - "input_dbu_cost_per_token": {"type": "number"}, - "annotation_cost_per_page": {"type": "number"}, - "annotation_cost_per_page_batches": {"type": "number"}, - "ocr_cost_per_page": {"type": "number"}, - "ocr_cost_per_page_batches": {"type": "number"}, - "ocr_cost_per_credit": {"type": "number"}, - "code_interpreter_cost_per_session": {"type": "number"}, - "inference_geo": {"type": "string"}, - "litellm_provider": {"type": "string"}, - "max_input_tokens": {"type": "number"}, - "max_output_tokens": {"type": "number"}, - "max_tokens": {"type": "number"}, - "metadata": {"type": "object"}, - "provider_specific_entry": {"type": "object"}, - "mode": { - "type": "string", - "enum": [ - "audio_speech", - "audio_transcription", - "chat", - "completion", - "container", - "image_edit", - "embedding", - "evaluation", - "extraction", - "guardrail", - "image_generation", - "video_generation", - "moderation", - "rerank", - "realtime", - "responses", - "ocr", - "search", - "vector_store", - ], - }, - "output_cost_per_audio_token": {"type": "number"}, - "output_cost_per_character": {"type": "number"}, - "output_cost_per_character_above_128k_tokens": {"type": "number"}, - "output_cost_per_image": {"type": "number"}, - "output_cost_per_image_512": {"type": "number"}, - "output_cost_per_image_1024": {"type": "number"}, - "output_cost_per_image_1536": {"type": "number"}, - "output_cost_per_image_0.5K": {"type": "number"}, - "output_cost_per_image_1K": {"type": "number"}, - "output_cost_per_image_2K": {"type": "number"}, - "output_cost_per_image_4K": {"type": "number"}, - "output_cost_per_image_token": {"type": "number"}, - "output_cost_per_image_token_batches": {"type": "number"}, - "output_cost_per_video_token": {"type": "number"}, - "output_cost_per_pixel": {"type": "number"}, - "output_cost_per_second": {"type": "number"}, - "output_cost_per_second_480p": {"type": "number"}, - "output_cost_per_second_720p": {"type": "number"}, - "output_cost_per_second_768p": {"type": "number"}, - "output_cost_per_second_2k": {"type": "number"}, - "output_cost_per_second_1080p": {"type": "number"}, - "output_cost_per_second_4k": {"type": "number"}, - "output_cost_per_token": {"type": "number"}, - "output_cost_per_token_above_32k_tokens": {"type": "number"}, - "output_cost_per_token_above_100k_tokens": {"type": "number"}, - "output_cost_per_token_above_100k_tokens_batches": {"type": "number"}, - "output_cost_per_token_above_128k_tokens": {"type": "number"}, - "output_cost_per_token_above_200k_tokens": {"type": "number"}, - "output_cost_per_token_above_200k_tokens_batches": {"type": "number"}, - "output_cost_per_token_above_256k_tokens": {"type": "number"}, - "output_cost_per_token_above_272k_tokens": {"type": "number"}, - "output_cost_per_token_above_512k_tokens": {"type": "number"}, - "output_cost_per_token_batches": {"type": "number"}, - "output_cost_per_reasoning_token": {"type": "number"}, - "output_cost_per_video_per_second": {"type": "number"}, - "output_db_cost_per_token": {"type": "number"}, - "output_dbu_cost_per_token": {"type": "number"}, - "output_vector_size": {"type": "number"}, - "rpd": {"type": "number"}, - "rpm": {"type": "number"}, - "source": {"type": "string"}, - "comment": {"type": "string"}, - "supports_assistant_prefill": {"type": "boolean"}, - "supports_anthropic_compaction": {"type": "boolean"}, - "supports_audio_input": {"type": "boolean"}, - "supports_audio_output": {"type": "boolean"}, - "gemini_native_audio": {"type": "boolean"}, - "gemini_audio_only_live": {"type": "boolean"}, - "supports_embedding_image_input": {"type": "boolean"}, - "supports_forced_tool_use": {"type": "boolean"}, - "supports_function_calling": {"type": "boolean"}, - "supports_image_input": {"type": "boolean"}, - "supports_nova_canvas_image_edit": {"type": "boolean"}, - "supports_parallel_function_calling": {"type": "boolean"}, - "supports_parallel_tool_use_config": {"type": "boolean"}, - "supports_pdf_input": {"type": "boolean"}, - "prompt_cache_min_tokens": {"type": "number"}, - "supports_prompt_cache_breakpoint": {"type": "boolean"}, - "supports_thinking_cache_preservation": {"type": "boolean"}, - "supports_prompt_caching": {"type": "boolean"}, - "supports_response_schema": {"type": "boolean"}, - "supports_system_messages": {"type": "boolean"}, - "supports_tool_choice": {"type": "boolean"}, - "supports_tool_search": {"type": "boolean"}, - "supports_video_input": {"type": "boolean"}, - "supports_box_questions": {"type": "boolean"}, - "supports_vision": {"type": "boolean"}, - "supports_web_search": {"type": "boolean"}, - "token_count_framing": {"type": "integer"}, - "token_count_model": { - "type": "object", - "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, - "required": ["model"], - }, - "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, - "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, - "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, - "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, - "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, - "supports_url_context": {"type": "boolean"}, - "supports_multimodal": {"type": "boolean"}, - "uses_embed_content": {"type": "boolean"}, - "supports_reasoning": {"type": "boolean"}, - "supports_minimal_reasoning_effort": {"type": "boolean"}, - "supports_low_reasoning_effort": {"type": "boolean"}, - "supports_none_reasoning_effort": {"type": "boolean"}, - "supports_xhigh_reasoning_effort": {"type": "boolean"}, - "supports_max_reasoning_effort": {"type": "boolean"}, - "reasoning_effort_levels": { - "type": "array", - "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, - }, - "default_reasoning_effort": { - "type": "string", - "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], - }, - "supports_adaptive_thinking": {"type": "boolean"}, - "supports_anthropic_thinking_payload": {"type": "boolean"}, - "supports_legacy_thinking": {"type": "boolean"}, - "thinking_always_on": {"type": "boolean"}, - "supports_mid_conversation_system": {"type": "boolean"}, - "supports_sampling_params": {"type": "boolean"}, - "supports_output_config": {"type": "boolean"}, - "supports_speed": {"type": "boolean"}, - "supports_fast_mode": {"type": "boolean"}, - "supported_audio_formats": { - "type": "array", - "items": { - "type": "string", - "enum": ["mp3", "wav"], - }, - }, - "vertex_ai_audio_api": { - "type": "string", - "enum": ["lyria_predict", "lyria_interactions"], - }, - "bedrock_output_config_effort_ceiling": { - "type": "string", - "enum": ["low", "medium", "high", "max", "xhigh"], - }, - "bedrock_converse_supports_strict_tools": {"type": "boolean"}, - "supports_regex_lookaround": {"type": "boolean"}, - "tpm": {"type": "number"}, - "supported_endpoints": { - "type": "array", - "items": { - "type": "string", - "enum": [ - "/v1/responses", - "/v1/embeddings", - "/v1/chat/completions", - "/v1/completions", - "/v1/messages", - "/v1/images/generations", - "/v1/realtime", - "/v1/realtime/transcription_sessions", - "/v1/images/variations", - "/v1/images/edits", - "/v1/batch", - "/v1beta/interactions", - "/v1beta/models", - "/v1/audio/transcriptions", - "/v1/audio/speech", - "/v1/ocr", - "/v1/videos", - "/vertex_ai/live", - "/v1/listen", - "/v1/systemone", - "/v1/decisions", - "/v1beta/interactions", - ], - }, - }, - "supported_regions": { - "type": "array", - "items": { - "type": "string", - }, - }, - "guardrail_cost_per_unit": { - "type": "object", - "additionalProperties": {"type": "number"}, - }, - "search_context_cost_per_query": { - "type": "object", - "properties": { - "search_context_size_low": {"type": "number"}, - "search_context_size_medium": {"type": "number"}, - "search_context_size_high": {"type": "number"}, - }, - "additionalProperties": False, - }, - "web_search_billing_unit": { - "type": "string", - "enum": ["per_prompt", "per_query"], - }, - "search_free_requests": { - "type": "object", - "properties": { - "count": {"type": "integer", "minimum": 1}, - "period": {"type": "string", "enum": ["day", "month"]}, - "shared_group": {"type": "string"}, - }, - "required": ["count", "period", "shared_group"], - "additionalProperties": False, - }, - "citation_cost_per_token": {"type": "number"}, - "supported_modalities": { - "type": "array", - "items": { - "type": "string", - "enum": ["text", "audio", "image", "video"], - }, - }, - "supported_output_modalities": { - "type": "array", - "items": { - "type": "string", - "enum": ["text", "image", "audio", "code", "video"], - }, - }, - "supports_native_streaming": {"type": "boolean"}, - "supports_image_size": {"type": "boolean"}, - "supports_native_structured_output": {"type": "boolean"}, - "supports_bedrock_native_structured_output": {"type": "boolean"}, - "use_openai_responses_path": {"type": "boolean"}, - "off_peak_pricing": { - "type": "object", - "properties": { - "hours_utc": { - "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], - }, - "windows": { - "type": "array", - "items": { - "type": "object", - "properties": { - "hours_utc": { - "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], - }, - "weekdays": { - "type": "array", - "items": {"oneOf": [{"type": "integer"}, {"type": "string"}]}, - }, - }, - "required": ["hours_utc"], - "additionalProperties": False, - }, - }, - "weekday_timezone": {"type": "string"}, - "input_cost_per_token": {"type": "number"}, - "output_cost_per_token": {"type": "number"}, - "output_cost_per_reasoning_token": {"type": "number"}, - "cache_read_input_token_cost": {"type": "number"}, - "cache_creation_input_token_cost": {"type": "number"}, - }, - "additionalProperties": False, - }, - "tiered_pricing": { - "type": "array", - "items": { - "type": "object", - "properties": { - "range": { - "type": "array", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "input_cost_per_token": {"type": "number"}, - "output_cost_per_token": {"type": "number"}, - "cache_read_input_token_cost": {"type": "number"}, - "cache_creation_input_token_cost": {"type": "number"}, - "output_cost_per_reasoning_token": {"type": "number"}, - "max_results_range": { - "type": "array", - "items": {"type": "number"}, - "minItems": 2, - "maxItems": 2, - }, - "input_cost_per_query": {"type": "number"}, - }, - "additionalProperties": False, - }, - }, - }, - "additionalProperties": False, - }, - } - prod_json = os.path.join(os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json") with open(prod_json, "r") as model_prices_file: actual_json = json.load(model_prices_file) @@ -1379,7 +1388,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): actual_json.pop("fallback_generalizations", None) # reserved meta key, not a model entry # Validate schema - validate(actual_json, INTENDED_SCHEMA) + validate(actual_json, _COST_MAP_ROW_SCHEMA) # Validate cost values # Define exceptions for models that are allowed to have costs > 1 @@ -1401,6 +1410,21 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): raise AssertionError(error_message) +def test_cost_map_schema_accepts_12hr_cache_write_rates_for_batch_and_ultrafast_tiers() -> None: + row: Final = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost_above_12hr_batches": 3e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 4e-6, + "cache_creation_input_token_cost_above_12hr_ultrafast": 5e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-6, + } + + assert Draft7Validator(_COST_MAP_ROW_SCHEMA).is_valid({"synthetic-12hr-cache-write-model": row}) + + def test_max_tokens_consistency(): """ Test that max_tokens == max_output_tokens for all models. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx index a66f31a0dc0..54decbfff3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx @@ -102,6 +102,7 @@ const AccessGroupBudgetModal: React.FC = ({ id={id} value={value || null} onChange={(next) => onChange(next ?? undefined)} + placeholder="Keep current" /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.integration.test.tsx index 4d0818987ce..60cd880b95e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.integration.test.tsx @@ -116,6 +116,36 @@ describe("AccessGroupBudgetsPanel", () => { expect(await screen.findByLabelText(/Max Budget/)).toHaveValue(2.5); }); + it("keeps the stored reset period when Keep current is selected", async () => { + GET.mockResolvedValue({ + data: { + access_groups: [ + { + ...BUDGETED_GROUP, + budget: { ...BUDGETED_GROUP.budget, budget_duration: "24h" }, + }, + ], + }, + }); + renderPanel(); + await openActions("premium"); + await userEvent.click(await screen.findByText("Edit budget")); + + await userEvent.click(await screen.findByRole("combobox", { name: "Reset Budget" })); + const keepCurrentOption = await screen.findByRole("option", { name: "Keep current" }); + expect(keepCurrentOption).toBeInTheDocument(); + await userEvent.click(keepCurrentOption); + await userEvent.click(screen.getByRole("button", { name: "Save Budget" })); + + await waitFor(() => + expect(PUT).toHaveBeenCalledWith("/access_group/{access_group}/budget", { + params: { path: { access_group: "premium" } }, + body: { max_budget: 2.5 }, + }), + ); + expect(PUT.mock.calls[0][1].body).not.toHaveProperty("budget_duration"); + }); + it("refuses to save a budget with every field blank", async () => { renderPanel(); await openActions("shared"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx index f0a8873a67c..f797d6e4f36 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx @@ -14,6 +14,7 @@ vi.mock("./user_edit_view", () => ({ UserEditView: ({ onSubmit, onCancel }: { onSubmit: (values: any) => void; onCancel: () => void }) => (
+
), @@ -204,6 +205,18 @@ describe("BulkEditUserModal", () => { }); }); + it("omits an unset reset period from the bulk update payload", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: "Submit with unset duration" })); + + await waitFor(() => { + expect(mockUserBulkUpdateUserCall).toHaveBeenCalledWith("test-token", { max_budget: 100 }, ["user1", "user2"]); + }); + expect(mockUserBulkUpdateUserCall.mock.calls[0][1]).not.toHaveProperty("budget_duration"); + }); + it("should call userBulkUpdateUserCall with allUsers flag when update all users is enabled", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx index 156b27672d5..f71384d80d7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx @@ -322,6 +322,35 @@ describe("UserEditView", () => { expect(callArgs.metadata).toEqual(MOCK_USER_DATA.user_info.metadata); }); + it("uses No change for an unset reset period in bulk edit", async () => { + const onSubmit = vi.fn(); + const userDataWithUnsetResetPeriod = { + ...MOCK_USER_DATA, + user_info: { + ...MOCK_USER_DATA.user_info, + max_budget: null, + budget_duration: null, + }, + }; + renderWithProviders( + , + ); + + const resetPeriod = await screen.findByRole("combobox", { name: "Reset Budget" }); + expect(resetPeriod).toHaveTextContent("No change"); + await userEvent.click(resetPeriod); + await userEvent.click(await screen.findByRole("option", { name: "No change" })); + + const unlimitedBudget = await screen.findByRole("checkbox", { name: "Unlimited Budget" }); + await waitFor(() => expect(unlimitedBudget).toBeChecked()); + await userEvent.click(unlimitedBudget); + fireEvent.change(screen.getByRole("spinbutton", { name: /max budget/i }), { target: { value: "12.5" } }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalled()); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ max_budget: "12.5", budget_duration: null }); + }); + it("should set max_budget to null when unlimited budget is checked on submit", async () => { const onSubmitMock = vi.fn(); const userDataWithNullBudget = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx index 647761bb44b..d4317d477f8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx @@ -315,7 +315,14 @@ export function UserEditView({ - {({ id, value, onChange }) => } + {({ id, value, onChange }) => ( + + )} {!isBulkEdit && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx index 948607bff37..d434987b873 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx @@ -201,13 +201,13 @@ describe("UserInfoView add-to-team form", () => { await openEditor(user); await user.click(screen.getByRole("combobox", { name: "Reset Budget" })); - await user.click(await screen.findByRole("option", { name: "n/a" })); + await user.click(await screen.findByRole("option", { name: "Never resets" })); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ budget_duration: null }); await openEditor(user); - expect(screen.getByRole("combobox", { name: "Reset Budget" })).toHaveTextContent("n/a"); + expect(screen.getByRole("combobox", { name: "Reset Budget" })).toHaveTextContent("Never resets"); }); }); diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 7ab6ec2679e..fcfa1ea5e80 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -730,9 +730,9 @@ describe("Teams - Reset Budget in team create", () => { it("should send an explicit null budget_duration when Never resets is selected", async () => { await openCreateModal(); + await screen.findByText("Default: monthly (30d)"); - await userEvent.click(resetBudgetSelect()); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, resetBudgetSelect(), "Never resets"); const payload = await submitCreateModal(); @@ -752,8 +752,7 @@ describe("Teams - Reset Budget in team create", () => { it("should send the picked duration when one is selected", async () => { await openCreateModal(); - await userEvent.click(resetBudgetSelect()); - await userEvent.click(await screen.findByText("weekly")); + await chooseSelectOption(userEvent, resetBudgetSelect(), "weekly"); const payload = await submitCreateModal(); @@ -768,14 +767,18 @@ describe("Teams - Reset Budget in team create", () => { }); }); - it("should fall back to the n/a placeholder when the default settings fetch fails", async () => { + it("should offer a single Never resets option when no server default is configured", async () => { vi.mocked(getDefaultTeamSettings).mockRejectedValue(new Error("Unauthorized")); await openCreateModal(); - await waitFor(() => { - expect(screen.getByText("n/a")).toBeInTheDocument(); + expect(getDefaultTeamSettings).toHaveBeenCalled(); }); + + await userEvent.click(resetBudgetSelect()); + + expect(await screen.findAllByRole("option", { name: "Never resets" })).toHaveLength(1); + expect(screen.queryByText("Not set")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 9971010c3ef..3a04340d3c4 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -295,7 +295,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, isViewOnly const defaultBudgetDuration: string | undefined = defaultTeamSettings?.values?.budget_duration ?? undefined; const budgetDurationPlaceholder = defaultBudgetDuration ? `Default: ${getBudgetDurationLabel(defaultBudgetDuration)} (${defaultBudgetDuration})` - : "n/a"; + : undefined; // Add this useEffect to fetch guardrails useEffect(() => { @@ -825,7 +825,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, isViewOnly {({ id, value, onChange }) => ( onChange(next ?? undefined)} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx new file mode 100644 index 00000000000..79e3be1750d --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx @@ -0,0 +1,90 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { TeamMemberSpendBudgetCell } from "./TeamMemberSpendBudgetCell"; + +describe("TeamMemberSpendBudgetCell", () => { + it("shows team and member budgets with the member reset date", () => { + render( + , + ); + + expect(screen.getByText("Team")).toBeInTheDocument(); + expect(screen.getByText("$600.00")).toBeInTheDocument(); + expect(screen.getByText("/ $1,000.00")).toBeInTheDocument(); + expect(screen.getByText("My Spend")).toBeInTheDocument(); + expect(screen.getByText("$50.00")).toBeInTheDocument(); + expect(screen.getByText("/ $100.00")).toBeInTheDocument(); + expect(screen.getByText("Resets Oct 20, 2026 UTC")).toBeInTheDocument(); + expect(screen.getAllByRole("meter")[0]).toHaveAttribute("aria-valuetext", "Team $600.00 of $1,000.00"); + expect(screen.getAllByRole("meter")[1]).toHaveAttribute("aria-valuetext", "My Spend $50.00 of $100.00"); + }); + + it("shows the team budget and Unlimited for a member without a budget", () => { + render( + , + ); + + expect(screen.getByText("Team")).toBeInTheDocument(); + expect(screen.getByText("$600.00")).toBeInTheDocument(); + expect(screen.getByText("My Spend")).toBeInTheDocument(); + expect(screen.getByText("/ Unlimited")).toBeInTheDocument(); + expect(screen.getAllByRole("meter")).toHaveLength(1); + }); + + it("omits the reset text when no member reset date exists", () => { + render( + , + ); + + expect(screen.queryByText(/Resets/)).not.toBeInTheDocument(); + }); + + it("renders only the team line when caller membership is absent", () => { + render(); + + expect(screen.getByText("Team")).toBeInTheDocument(); + expect(screen.queryByText("My Spend")).not.toBeInTheDocument(); + expect(screen.getAllByRole("meter")).toHaveLength(1); + }); + + it("exposes member spend above its budget through the meter", () => { + render( + , + ); + + expect(screen.getAllByRole("meter")[1]).toHaveAttribute("aria-valuetext", "My Spend $150.00 of $100.00"); + }); + + it("does not render meters for zero or unlimited budgets", () => { + render( + , + ); + + expect(screen.queryByRole("meter")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx new file mode 100644 index 00000000000..9bea9892899 --- /dev/null +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx @@ -0,0 +1,70 @@ +"use client"; + +import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; +import { getSpendBudgetMeterTone } from "@/components/shared/table_cells/spend_budget_cell"; +import type { components } from "@/lib/http/schema"; +import { formatBudgetResetUtc } from "@/utils/budgetUtils"; +import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; + +type CallerMembership = components["schemas"]["TeamListCallerMembership"]; + +interface TeamMemberSpendBudgetCellProps { + teamSpend: number | null | undefined; + teamMaxBudget: number | null | undefined; + callerMembership?: CallerMembership | null; +} + +interface BudgetLineProps { + label: "Team" | "My Spend"; + spend: number | null | undefined; + maxBudget: number | null | undefined; + resetAt?: string | null; +} + +function BudgetLine({ label, spend, maxBudget, resetAt }: BudgetLineProps) { + const spendValue = typeof spend === "number" && Number.isFinite(spend) ? spend : 0; + const budget = typeof maxBudget === "number" && Number.isFinite(maxBudget) ? maxBudget : null; + const spendText = spendValue === 0 ? "$0.00" : getSpendString(spendValue, 2); + const budgetText = budget === null ? "Unlimited" : `$${formatNumberWithCommas(budget, 2)}`; + const resetText = formatBudgetResetUtc(resetAt); + + return ( +
+
+ {label} + + {spendText}{" "} + / {budgetText} + +
+ {budget !== null && budget > 0 && ( + + + + + + )} + {resetText &&
Resets {resetText} UTC
} +
+ ); +} + +export function TeamMemberSpendBudgetCell({ + teamSpend, + teamMaxBudget, + callerMembership, +}: TeamMemberSpendBudgetCellProps) { + return ( +
+ + {callerMembership && ( + + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx index def6750071f..02870bc1571 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx @@ -112,10 +112,30 @@ it("renders a team row with alias, organization, and spend/budget", async () => expect(screen.getByText("Acme Team")).toBeInTheDocument(); expect(screen.getByText("Test Organization")).toBeInTheDocument(); expect(screen.getByText("$42.50")).toBeInTheDocument(); - expect(screen.getByText("of $100.00")).toBeInTheDocument(); + expect(screen.getByText("/ $100.00")).toBeInTheDocument(); }); }); +it("renders caller membership spend and budget in the team row", async () => { + mockUseTeamsTable.mockReturnValue( + teamsResult([ + { + ...mockTeam, + caller_membership: { + spend: 50, + max_budget: 100, + budget_reset_at: "2026-10-20T12:00:00Z", + }, + }, + ]), + ); + renderTable(); + + expect(await screen.findByText("My Spend")).toBeInTheDocument(); + expect(screen.getByText("$50.00")).toBeInTheDocument(); + expect(screen.getByText("Resets Oct 20, 2026 UTC")).toBeInTheDocument(); +}); + it("renders the Resources cell with member, model, and key counts", () => { renderTable(); @@ -154,7 +174,7 @@ describe("sort contract – only backend-sortable columns are sortable", () => { it("sorts by the backend team_alias field (not the label) when the Team header is clicked", async () => { renderTable(); - fireEvent.click(screen.getByText("Team").closest("button") as HTMLElement); + fireEvent.click(within(screen.getByRole("columnheader", { name: "Team" })).getByRole("button")); await waitFor(() => { expect(mockUseTeamsTable).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "team_alias" })); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx index e56e0d16522..abc66333f11 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx @@ -4,7 +4,7 @@ import { ColumnDef } from "@tanstack/react-table"; import { Copy, KeyRound, Layers, MoreHorizontal, Pencil, Trash2, Users } from "lucide-react"; import { DataTableMultiSortHeader, DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdentityCell, SpendBudgetCell } from "@/components/shared/table_cells"; +import { DateCell, IdentityCell } from "@/components/shared/table_cells"; import { buttonVariants } from "@/components/ui/button"; import { DropdownMenu, @@ -20,6 +20,7 @@ import { copyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils"; import { Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; +import { TeamMemberSpendBudgetCell } from "./TeamMemberSpendBudgetCell"; interface ResourceTone { icon: typeof Users; @@ -233,11 +234,10 @@ export const getTeamTableColumns = ({ size: 200, enableSorting: false, cell: ({ row }) => ( - ), }, diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx new file mode 100644 index 00000000000..184e4297a67 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx @@ -0,0 +1,41 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import BudgetDurationDropdown from "./budget_duration_dropdown"; + +describe("BudgetDurationDropdown", () => { + it("labels an unset period as Never resets in the trigger and option", async () => { + const user = userEvent.setup(); + render(); + + expect(screen.getByRole("combobox")).toHaveTextContent("Never resets"); + + await user.click(screen.getByRole("combobox")); + + expect(await screen.findByRole("option", { name: "Never resets" })).toBeInTheDocument(); + }); + + it("emits null when Never resets is selected", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + render(); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByRole("option", { name: "Never resets" })); + + expect(onChange).toHaveBeenCalledOnce(); + expect(onChange.mock.calls[0][0]).toBeNull(); + }); + + it("uses a supplied placeholder for the unset trigger and option", async () => { + const user = userEvent.setup(); + render(); + + expect(screen.getByRole("combobox")).toHaveTextContent("Not set"); + + await user.click(screen.getByRole("combobox")); + + expect(await screen.findByRole("option", { name: "Not set" })).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 40ee857634c..bae6986ab03 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -27,7 +27,7 @@ const BudgetDurationDropdown: React.FC = ({ onChange, className = "", style = {}, - placeholder = "n/a", + placeholder = "Never resets", showNeverResets = false, }) => { return ( diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 2f0b9c9b00e..11dd4f2c2e1 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -1,4 +1,5 @@ import { Member } from "../networking"; +import type { components } from "@/lib/http/schema"; import type { ObjectPermission } from "../object_permission_types"; import type { ModelBudgetUsage, ModelMaxBudget } from "./ModelMaxBudgetEditor"; @@ -21,6 +22,7 @@ export interface Team { keys_count?: number; members_count?: number; members_with_roles: Member[]; + caller_membership?: components["schemas"]["TeamListCallerMembership"] | null; team_member_permissions?: string[] | null; spend: number; access_group_ids?: string[]; diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx index 914d3a98c81..6090b42e93d 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx @@ -13,7 +13,7 @@ interface SpendBudgetCellProps { budgetDecimals?: number; } -const meterTone = (pct: number): "default" | "warning" | "over" => { +export const getSpendBudgetMeterTone = (pct: number): "default" | "warning" | "over" => { if (pct > 100) return "over"; if (pct >= 80) return "warning"; return "default"; @@ -76,7 +76,7 @@ export function SpendBudgetCell({ {hasBudget && ( - + )} diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx index 2c98d17f961..81106bc094a 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx @@ -106,6 +106,27 @@ describe("EditMembership submit payload", () => { }); }); + it("shows Never resets and submits a null duration for a member without a reset period", async () => { + renderEdit(teamMemberConfig, { user_id: "u1", user_email: "a@b.com", role: "user" }); + + expect(screen.getByRole("combobox", { name: "Budget Reset Period" })).toHaveTextContent("Never resets"); + + save(); + + await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce()); + const expectedPayload = { + user_email: "a@b.com", + user_id: "u1", + role: "user", + max_budget_in_team: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + allowed_models: [], + }; + expect(submitted()).toStrictEqual(expectedPayload); + }); + it("seeds the per-model budget editor from the member's stored budget and submits the edited cap", async () => { const modelBudgetConfig = { ...teamMemberConfig, diff --git a/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx b/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx index d9ba3d82230..24ce626ecb7 100644 --- a/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/MyUserTab.test.tsx @@ -45,5 +45,30 @@ describe("MyUserTab", () => { expect(screen.getByText("member@example.com")).toBeInTheDocument(); expect(screen.getByText("model-one")).toBeInTheDocument(); expect(screen.getByText("TPM: 1,000")).toBeInTheDocument(); + expect(screen.getByText("Never resets")).toBeInTheDocument(); + }); + + it("shows the next reset time when the member budget has a reset window", () => { + vi.mocked(useMyTeamMember).mockReturnValue({ + data: { + user_id: "user-1", + user_email: "member@example.com", + team_id: "team-1", + role: "user", + spend: 12.5, + total_spend: 30, + litellm_budget_table: { + max_budget: 100, + budget_duration: "30d", + budget_reset_at: new Date(2026, 6, 15, 12, 0, 0).toISOString(), + }, + }, + isLoading: false, + error: null, + } as ReturnType); + + renderWithProviders(); + + expect(screen.getByText("Resets Jul 15, 2026 12:00")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/team/MyUserTab.tsx b/ui/litellm-dashboard/src/components/team/MyUserTab.tsx index 9aae6a25dc2..bf1a2c2c242 100644 --- a/ui/litellm-dashboard/src/components/team/MyUserTab.tsx +++ b/ui/litellm-dashboard/src/components/team/MyUserTab.tsx @@ -1,4 +1,4 @@ -import { formatBudgetReset } from "@/utils/budgetUtils"; +import { describeBudgetResetSchedule } from "@/utils/budgetUtils"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { Badge } from "@/components/ui/badge"; @@ -67,7 +67,7 @@ export default function MyUserTab({ teamId }: MyUserTabProps) { const totalSpend = data.total_spend ?? 0; const tpmLimit = budgetTable?.tpm_limit ?? null; const rpmLimit = budgetTable?.rpm_limit ?? null; - const budgetReset = formatBudgetReset(budgetTable?.budget_reset_at); + const budgetReset = describeBudgetResetSchedule(budgetTable?.budget_duration, budgetTable?.budget_reset_at); const allowedModels = budgetTable?.allowed_models ?? null; return ( @@ -93,17 +93,16 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
- {labelWithTooltip( - "Current Cycle Spend (USD)", - "Spend for the current budget cycle. Resets to $0 when the budget window rolls over.", - )} + {labelWithTooltip("Current Cycle Spend (USD)", budgetReset.tooltip)}

${formatNumber(spend, 4)}

of {maxBudget === null ? "Unlimited" : `$${formatNumber(maxBudget, 4)}`}
- {budgetReset &&
Resets {budgetReset}
} +
+ {budgetReset.resetAt ? `Resets ${budgetReset.label}` : budgetReset.label} +
@@ -120,7 +119,10 @@ export default function MyUserTab({ teamId }: MyUserTabProps) { - {labelWithTooltip("Total Spend (USD)", "Cumulative spend across all budget cycles within this team.")} + {labelWithTooltip( + "Total Spend (USD)", + "Cumulative spend across budget cycles within this team. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.", + )}

${formatNumber(totalSpend, 4)}

 diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index cd51bcedb7b..38a8e4f97d5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -313,7 +313,47 @@ describe("TeamMembersComponent", () => { />, ); - expect(screen.getByText("Jul 15, 2026")).toBeInTheDocument(); + expect(screen.getByText("Jul 15, 2026 12:00")).toBeInTheDocument(); + }); + + it("shows Never resets for a member whose budget has no reset window", () => { + const teamData = createMockTeamData(); + teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null; + renderWithProviders( + , + ); + + expect(screen.queryByText("Jul 15, 2026 12:00")).not.toBeInTheDocument(); + expect(screen.getAllByText("Never resets").length).toBeGreaterThan(0); + }); + + it("shows Unknown for a member whose budget has a duration but no scheduled reset yet", () => { + const teamData = createMockTeamData(); + teamData.team_memberships[0].litellm_budget_table!.budget_duration = "30d"; + teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null; + renderWithProviders( + , + ); + + expect(screen.getByText("Unknown")).toBeInTheDocument(); }); it("shows each member the values from its own membership row", () => { diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index ddfd251cf82..319f8c973f6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -10,8 +10,9 @@ import MemberTable from "@/components/common_components/MemberTable"; import type { StoredModelMaxBudget } from "@/components/key_team_helpers/modelMaxBudgetPayload"; import { Member } from "@/components/networking"; import { parseErrorMessage } from "@/components/shared/errorUtils"; -import { DateCell, MoneyCell } from "@/components/shared/table_cells"; +import { MoneyCell } from "@/components/shared/table_cells"; import { toast } from "@/lib/toast"; +import { describeBudgetResetSchedule } from "@/utils/budgetUtils"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles"; import { CircleHelp } from "lucide-react"; @@ -152,6 +153,16 @@ export default function TeamMemberTab({ return findMembership(userId)?.litellm_budget_table?.budget_reset_at ?? null; }; + const renderBudgetReset = (userId: string | null) => { + const budget = findMembership(userId)?.litellm_budget_table; + const schedule = describeBudgetResetSchedule(budget?.budget_duration ?? null, budget?.budget_reset_at ?? null); + return ( + + {schedule.label} + + ); + }; + const extraColumns: NonNullable["extraColumns"]> = [ { title: ( @@ -190,7 +201,7 @@ export default function TeamMemberTab({ title: ( Current Cycle Spend (USD) - + @@ -204,7 +215,7 @@ export default function TeamMemberTab({ title: ( Total Spend (USD) - + @@ -255,7 +266,7 @@ export default function TeamMemberTab({ title: "Budget Reset", key: "budget_reset", sortValue: (record: Member) => getUserBudgetReset(record.user_id), - render: (record: Member) => , + render: (record: Member) => renderBudgetReset(record.user_id), }, { title: ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d5b17abab46..c025f27e67f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35251,14 +35251,22 @@ export interface components { cache_creation_input_token_cost_above_12hr?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens */ cache_creation_input_token_cost_above_12hr_above_272k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Batches */ + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Flex */ cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Priority */ cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Ultrafast */ + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Batches */ + cache_creation_input_token_cost_above_12hr_batches?: number | null; /** Cache Creation Input Token Cost Above 12Hr Flex */ cache_creation_input_token_cost_above_12hr_flex?: number | null; /** Cache Creation Input Token Cost Above 12Hr Priority */ cache_creation_input_token_cost_above_12hr_priority?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Ultrafast */ + cache_creation_input_token_cost_above_12hr_ultrafast?: number | null; /** Cache Creation Input Token Cost Above 1Hr */ cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 1Hr Above 100K Tokens */ @@ -46370,6 +46378,21 @@ export interface components { /** Team Id */ team_id: string; }; + /** + * TeamListCallerMembership + * @description The calling user's own spend and effective member budget in a team. + */ + TeamListCallerMembership: { + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** + * Spend + * @default 0 + */ + spend: number; + }; /** * TeamListItem * @description A team item in the paginated list response, enriched with computed fields. @@ -46406,6 +46429,7 @@ export interface components { budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null; /** Budget Reset At */ budget_reset_at?: string | null; + caller_membership?: components["schemas"]["TeamListCallerMembership"] | null; /** Created At */ created_at?: string | null; /** Default Team Member Models */ @@ -50288,14 +50312,22 @@ export interface components { cache_creation_input_token_cost_above_12hr?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens */ cache_creation_input_token_cost_above_12hr_above_272k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Batches */ + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Flex */ cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Priority */ cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Ultrafast */ + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Batches */ + cache_creation_input_token_cost_above_12hr_batches?: number | null; /** Cache Creation Input Token Cost Above 12Hr Flex */ cache_creation_input_token_cost_above_12hr_flex?: number | null; /** Cache Creation Input Token Cost Above 12Hr Priority */ cache_creation_input_token_cost_above_12hr_priority?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Ultrafast */ + cache_creation_input_token_cost_above_12hr_ultrafast?: number | null; /** Cache Creation Input Token Cost Above 1Hr */ cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Read Input Token Cost */ @@ -50593,14 +50625,22 @@ export interface components { cache_creation_input_token_cost_above_12hr?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens */ cache_creation_input_token_cost_above_12hr_above_272k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Batches */ + cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Flex */ cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex?: number | null; /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Priority */ cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Above 272K Tokens Ultrafast */ + cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Batches */ + cache_creation_input_token_cost_above_12hr_batches?: number | null; /** Cache Creation Input Token Cost Above 12Hr Flex */ cache_creation_input_token_cost_above_12hr_flex?: number | null; /** Cache Creation Input Token Cost Above 12Hr Priority */ cache_creation_input_token_cost_above_12hr_priority?: number | null; + /** Cache Creation Input Token Cost Above 12Hr Ultrafast */ + cache_creation_input_token_cost_above_12hr_ultrafast?: number | null; /** Cache Creation Input Token Cost Above 1Hr */ cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 1Hr Above 100K Tokens */ diff --git a/ui/litellm-dashboard/src/utils/budgetUtils.test.ts b/ui/litellm-dashboard/src/utils/budgetUtils.test.ts new file mode 100644 index 00000000000..a3b5f4da00a --- /dev/null +++ b/ui/litellm-dashboard/src/utils/budgetUtils.test.ts @@ -0,0 +1,46 @@ +import { afterEach, describe, expect, it } from "vitest"; + +import { describeBudgetResetSchedule, formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils"; + +const originalTimezone: string | undefined = process.env.TZ; + +afterEach(() => { + if (originalTimezone === undefined) { + delete process.env.TZ; + return; + } + process.env.TZ = originalTimezone; +}); + +describe("formatBudgetResetUtc", () => { + it("formats the reset date in UTC regardless of the local timezone", () => { + process.env.TZ = "America/Los_Angeles"; + const resetAt: string = "2026-10-20T00:00:00Z"; + + expect(formatBudgetResetUtc(resetAt)).toBe("Oct 20, 2026"); + expect(formatBudgetReset(resetAt)).toBe("Oct 19, 2026"); + }); +}); + +describe("describeBudgetResetSchedule", () => { + it("reports Never resets when the budget has no duration and no scheduled reset", () => { + const schedule = describeBudgetResetSchedule(null, null); + expect(schedule.label).toBe("Never resets"); + expect(schedule.resetAt).toBeNull(); + }); + + it("reports Unknown when a duration is set but no reset has been scheduled yet", () => { + const schedule = describeBudgetResetSchedule("30d", null); + expect(schedule.label).toBe("Unknown"); + expect(schedule.tooltip).toContain("every 30d"); + }); + + it("shows the next reset with its time zone and the rollover carryover note", () => { + const resetAt = new Date(2026, 6, 15, 12, 0, 0); + const schedule = describeBudgetResetSchedule("30d", resetAt.toISOString()); + expect(schedule.label).toBe("Jul 15, 2026 12:00"); + expect(schedule.resetAt).toBe(resetAt.toISOString()); + expect(schedule.tooltip).toContain(Intl.DateTimeFormat().resolvedOptions().timeZone); + expect(schedule.tooltip).toContain("carries into the next cycle"); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/budgetUtils.ts b/ui/litellm-dashboard/src/utils/budgetUtils.ts index 3d3278db88f..171ce5818f7 100644 --- a/ui/litellm-dashboard/src/utils/budgetUtils.ts +++ b/ui/litellm-dashboard/src/utils/budgetUtils.ts @@ -1,4 +1,7 @@ import dayjs from "dayjs"; +import utc from "dayjs/plugin/utc"; + +dayjs.extend(utc); export function formatBudgetReset(iso: string | null | undefined): string | null { if (!iso) return null; @@ -6,3 +9,48 @@ export function formatBudgetReset(iso: string | null | undefined): string | null if (!resetDate.isValid()) return null; return resetDate.format("MMM D, YYYY"); } + +export function formatBudgetResetUtc(iso: string | null | undefined): string | null { + if (!iso) return null; + const resetDate = dayjs.utc(iso); + if (!resetDate.isValid()) return null; + return resetDate.format("MMM D, YYYY"); +} + +export interface BudgetResetSchedule { + label: string; + resetAt: string | null; + tooltip: string; +} + +const ROLLOVER_NOTE = "Spend above a rollover cap at reset carries into the next cycle."; + +export function describeBudgetResetSchedule( + budgetDuration: string | null | undefined, + budgetResetAt: string | null | undefined, +): BudgetResetSchedule { + const resetDate = budgetResetAt ? dayjs(budgetResetAt) : null; + if (resetDate?.isValid()) { + const cadence = budgetDuration ? `every ${budgetDuration}` : "on a schedule"; + const timeZone = Intl.DateTimeFormat().resolvedOptions().timeZone; + const label = resetDate.format("MMM D, YYYY HH:mm"); + return { + label, + resetAt: resetDate.toISOString(), + tooltip: `Current Cycle Spend resets ${cadence}. Next reset ${label} (${timeZone}). ${ROLLOVER_NOTE}`, + }; + } + if (!budgetDuration) { + return { + label: "Never resets", + resetAt: null, + tooltip: + "This budget has no reset window, so Current Cycle Spend accumulates for as long as the membership exists.", + }; + } + return { + label: "Unknown", + resetAt: null, + tooltip: `Resets every ${budgetDuration}, but the next reset time has not been scheduled yet. ${ROLLOVER_NOTE}`, + }; +}