From 455923014a6ef95ac5d98a3f8845969fd44d1450 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 09:07:02 -0700 Subject: [PATCH 01/43] feat(cost): bill openai 12-hour prompt cache writes at their own rate (#48088) * feat: add OpenAI 12-hour cache write pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: omit empty cache-write response fields * style: format cache-write fields * fix(cost): bill openai 30m and 12h cache writes from the usage split, add gpt-5.6 rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost): apply the 12h long-context write rate only above the token threshold Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(cost): key cache write rates by TTL instead of growing the rate tuple _get_token_base_cost returned a positional 6-tuple and every cache lifetime added a slot to it, to the off-peak helper and to every caller. It now returns a TokenBaseCosts dataclass whose cache write rates live in one mapping keyed by TTL. The TTLs come from the ephemeral_*_input_tokens fields on CacheCreationTokenDetails, so supporting a new lifetime only means adding that usage field and its rate key A lifetime with no rate of its own now always falls back to the cache creation rate after off-peak pricing is applied. Before this the non-tiered 12h rate and the tiered 1h rate fell back to the rate from before off-peak, so the same request billed differently depending on the pricing path * refactor(cost): hardcode the cache write TTLs as a Literal instead of parsing usage field names Parsing ephemeral_*_input_tokens names saved nothing, since a new TTL still needs a usage field and provider parsing code. The TTLs are now a CacheWriteTTL Literal with an explicit rate key per TTL * refactor(cost): resolve cache write lifetime rates to floats at lookup Decide the 1h and 12h cost map key in the threshold loop, then resolve each rate once after off-peak with the base write rate as the default, so no None flows through the cost code. Add an integration case for a 12h write on a row without a 12h rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(cost): name every cache write TTL, marking 5m and 30m as base rate Co-Authored-By: Claude Opus 5.5 * refactor(cost): list every cache write TTL in one rate key table, None meaning the base write rate Co-Authored-By: Claude Opus 5.5 * refactor(cost): drop base rate entries from the cache write rate table and rely on fallback Co-Authored-By: Claude Opus 5.5 * refactor(cost): key cache write rates by plain string TTL Co-Authored-By: Claude Opus 5.5 * refactor(cost): fill missing cache write TTL rates with the base write rate after off-peak Co-Authored-By: Claude Opus 5.5 * refactor(cost): default missing cache write TTL rates to the base write rate inline Co-Authored-By: Claude Opus 5.5 * refactor(cost): rename cache write rate key table to make clear it only covers long TTLs Co-Authored-By: Claude Opus 5.5 * fix(cost): list every cache write TTL and bill writes of an unknown TTL at the base write rate Co-Authored-By: Claude Opus 5.5 * refactor(cost): resolve tiered cache write TTL rates after off-peak instead of passing None Co-Authored-By: Claude Opus 5.5 * refactor(cost): carry cache write TTL rates on TokenRates and drop TokenBaseCosts Co-Authored-By: Claude Opus 5.5 * refactor(cost): define TokenRates before its first use to drop the string annotation Co-Authored-By: Claude Opus 5.5 * fix(cost): build the TokenRates TTL default with default_factory and type tiers as Mapping MappingProxyType({}) as a dataclass default is rejected on Python 3.11, where mappingproxy is unhashable, so import litellm failed there. The tier type moves from dict[str, object] to Mapping[str, object] (LIT001), with tier_rate accepting a Mapping to match. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost): type select_tier_for_input tiers as Mapping so tier_rate callers stay typed The tier returned as a bare dict made every tier_rate call site a partially unknown argument once tier_rate took Mapping[str, object]. The range bounds are now validated with a TypeAdapter instead of indexing an unknown list. --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../llm_cost_calc/tiered_pricing.py | 42 ++- .../litellm_core_utils/llm_cost_calc/utils.py | 278 ++++++++++-------- litellm/llms/dashscope/cost_calculator.py | 3 +- ...odel_prices_and_context_window_backup.json | 48 +++ litellm/proxy/_lazy_openapi_snapshot.json | 40 +++ .../transformation.py | 26 +- litellm/responses/utils.py | 22 +- litellm/types/llms/openai.py | 17 +- litellm/types/utils.py | 14 + litellm/utils.py | 18 ++ model_prices_and_context_window.json | 48 +++ .../cost_calculation/cases/openai.py | 102 +++++++ .../integration/cost_calculation/cost_map.py | 2 + .../cost_calculation/cost_tracking_case.py | 6 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 269 ++++++++++++----- .../test_litellm_completion_responses.py | 4 +- tests/unit/responses/test_responses_utils.py | 41 ++- tests/unit/test_utils.py | 12 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 52 +++- 19 files changed, 811 insertions(+), 233 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py index 9bcc2b1743c..3040e36ec5f 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py @@ -5,30 +5,42 @@ Shared by provider cost calculators (e.g. Dashscope) and the proxy budget reservation logic so neither has to depend on the other. """ -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, SupportsFloat + +from pydantic import TypeAdapter, ValidationError + +_TIER_RANGE: Final = TypeAdapter(tuple[float, float]) -def _coerce_cost_per_token(value: float | str | None) -> float: +def _coerce_cost_per_token(value: object) -> float: """ Coerce a per-token cost into a float. Model cost values loaded from YAML config may arrive as strings (e.g. scientific notation like "4e-07"), which would break arithmetic. """ - if value is None: - return 0.0 if isinstance(value, str): try: return float(value) except ValueError: return 0.0 - return float(value) + if isinstance(value, SupportsFloat): + return float(value) + return 0.0 + + +def _tier_range(tier: Mapping[str, object]) -> tuple[float, float] | None: + try: + return _TIER_RANGE.validate_python(tier.get("range")) + except ValidationError: + return None def select_tier_for_input( - tiered_pricing: list[dict], + tiered_pricing: Sequence[Mapping[str, object]], input_tokens: int, -) -> dict | None: +) -> Mapping[str, object] | None: """ Select the pricing tier for a request based on its total input token count. @@ -44,19 +56,21 @@ def select_tier_for_input( if not tiered_pricing or input_tokens <= 0: return None - sorted_tiers: Final = sorted(tiered_pricing, key=lambda t: t.get("range", [0, 0])[0]) - valid_tiers: Final = [tier for tier in sorted_tiers if len(tier.get("range", [])) == 2] - if not valid_tiers: + ranged_tiers: Final = sorted( + ((bounds, tier) for tier in tiered_pricing if (bounds := _tier_range(tier)) is not None), + key=lambda ranged: ranged[0][0], + ) + if not ranged_tiers: return None - matching: Final = [tier for tier in valid_tiers if tier["range"][0] < input_tokens <= tier["range"][1]] + matching: Final = tuple(tier for (start, end), tier in ranged_tiers if start < input_tokens <= end) if matching: return matching[0] - return valid_tiers[-1] + return ranged_tiers[-1][1] def tier_rate( - tier: dict, + tier: Mapping[str, object], cost_key: str, fallback_cost_key: str | None = None, ) -> float: @@ -68,4 +82,6 @@ def tier_rate( primary: Final = tier.get(cost_key) if primary is not None: return _coerce_cost_per_token(primary) + if fallback_cost_key is None: + return 0.0 return _coerce_cost_per_token(tier.get(fallback_cost_key, 0)) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index a874459798d..5d4b1a82759 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -3,7 +3,7 @@ import re from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field, replace from datetime import datetime, timezone, tzinfo from types import MappingProxyType from typing import Final, Literal, TypedDict, cast @@ -52,6 +52,22 @@ _VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency) _DEPLOYMENT_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields) +_BASE_CACHE_WRITE_RATE_KEY: Final = "cache_creation_input_token_cost" + +_CACHE_WRITE_RATE_KEYS: Final[Mapping[str, str]] = MappingProxyType( + { + "5m": _BASE_CACHE_WRITE_RATE_KEY, + "30m": _BASE_CACHE_WRITE_RATE_KEY, + "1hr": "cache_creation_input_token_cost_above_1hr", + "12hr": "cache_creation_input_token_cost_above_12hr", + } +) + +_TTL_PRICED_CACHE_WRITE_RATE_KEYS: Final[Mapping[str, str]] = MappingProxyType( + {ttl: key for ttl, key in _CACHE_WRITE_RATE_KEYS.items() if key != _BASE_CACHE_WRITE_RATE_KEY} +) + + _IMAGE_TOKEN_RATE_KEYS: Final[tuple[str, ...]] = ( "input_cost_per_token", "output_cost_per_token", @@ -356,7 +372,21 @@ def get_batch_cost_rates(model_info: ModelInfo, usage: Usage, custom_llm_provide ) -def _select_priced_tier(model_info: ModelInfo, usage: Usage) -> dict[str, object] | None: +@dataclass(frozen=True, slots=True) +class TokenRates: + input_rate: float + output_rate: float + cache_read_rate: float + cache_creation_rate: float + reasoning_rate: float | None + cache_creation_rates_by_ttl: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + +def _select_priced_tier(model_info: ModelInfo, usage: Usage) -> Mapping[str, object] | None: tiered_pricing: Final = model_info.get("tiered_pricing") if not isinstance(tiered_pricing, list) or not tiered_pricing: return None @@ -376,7 +406,7 @@ def _get_tiered_reasoning_rate(model_info: ModelInfo, usage: Usage) -> float | N return tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") -def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, float, float, float, float] | None: +def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[TokenRates, Mapping[str, object]] | None: """ Resolve the base rates from a model's ``tiered_pricing`` table, if it has one. @@ -391,22 +421,20 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, if tier is None: return None - cache_creation_cost: Final = tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") completion_cost: Final = ( tier_rate(tier, "output_cost_per_token") if "output_cost_per_token" in tier else get_cost_per_unit(model_info, "output_cost_per_token") or 0.0 ) - return ( - tier_rate(tier, "input_cost_per_token"), - completion_cost, - cache_creation_cost, - tier_rate(tier, "cache_creation_input_token_cost_above_1hr") - if "cache_creation_input_token_cost_above_1hr" in tier - and tier["cache_creation_input_token_cost_above_1hr"] is not None - else cache_creation_cost, - tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate: Final = tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") + rates: Final = TokenRates( + input_rate=tier_rate(tier, "input_cost_per_token"), + output_rate=completion_cost, + cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=cache_creation_rate, + reasoning_rate=None, ) + return rates, tier def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_time: datetime | None = None) -> bool: @@ -532,19 +560,6 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = return False -@dataclass(frozen=True, slots=True) -class TokenRates: - input_rate: float - output_rate: float - cache_read_rate: float - cache_creation_rate: float - reasoning_rate: float | None - - @property - def billed_reasoning_rate(self) -> float: - return self.output_rate if self.reasoning_rate is None else self.reasoning_rate - - def resolve_explicit_cache_read_cost_rate( model_info: ModelInfo, usage: Usage, @@ -649,7 +664,8 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, if off_peak is None: return rates off_peak_reasoning_rate: Final = _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token")) - return TokenRates( + return replace( + rates, input_rate=_off_peak_rate(off_peak, "input_cost_per_token", rates.input_rate), output_rate=_off_peak_rate(off_peak, "output_cost_per_token", rates.output_rate), cache_read_rate=_off_peak_rate(off_peak, "cache_read_input_token_cost", rates.cache_read_rate), @@ -661,32 +677,28 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, def _apply_off_peak_to_base_costs( model_info: ModelInfo, current_time: datetime | None, - base_costs: tuple[float, float, float, float | None, float], -) -> tuple[float, float, float, float, float]: - """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path - produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a - present one passes through untouched and an absent one resolves to the applied - cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate. + base_rates: TokenRates, + tier: Mapping[str, object], +) -> TokenRates: + """Apply off-peak rates to the base costs resolved from a pricing tier. off_peak_pricing has no + field for the long-TTL cache-creation rates, so a tier that declares one keeps it, and a tier + that does not bills that TTL at the off-peak-adjusted write rate. Reasoning is left to + _resolve_billed_reasoning_rate. """ - prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs - rates: Final = apply_off_peak_pricing( - model_info, - current_time, - TokenRates( - input_rate=prompt, - output_rate=completion, - cache_read_rate=cache_read, - cache_creation_rate=cache_creation, - reasoning_rate=None, + rates: Final = apply_off_peak_pricing(model_info, current_time, base_rates) + return replace( + rates, + cache_creation_rates_by_ttl=MappingProxyType( + { + ttl: rates.cache_creation_rate if tier.get(key) is None else tier_rate(tier, key) + for ttl, key in _TTL_PRICED_CACHE_WRITE_RATE_KEYS.items() + } ), ) - return ( - rates.input_rate, - rates.output_rate, - rates.cache_creation_rate, - rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr, - rates.cache_read_rate, - ) + + +def _first_present_key(model_info: ModelInfo, preferred_key: str, fallback_key: str) -> str: + return preferred_key if model_info.get(preferred_key) is not None else fallback_key def _get_token_base_cost( @@ -696,7 +708,7 @@ def _get_token_base_cost( current_time: datetime | None = None, *, threshold_is_inclusive: bool = False, -) -> tuple[float, float, float, float, float]: +) -> TokenRates: """ Return prompt cost, completion cost, and cache costs for a given model and usage. @@ -709,15 +721,12 @@ def _get_token_base_cost( An absent cache-creation or cache-read rate always resolves to the resolved input rate, the way the tiered table and custom deployment pricing already do, since a provider that publishes no cache price bills cached tokens as ordinary input. An - absent 1h write rate resolves to the cache-creation rate, off-peak included. An + absent long-TTL write rate resolves to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for all of them. - - Returns: - Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) """ tiered_base_costs: Final = _get_tiered_base_costs(model_info=model_info, usage=usage) if tiered_base_costs is not None: - return _apply_off_peak_to_base_costs(model_info, current_time, tiered_base_costs) + return _apply_off_peak_to_base_costs(model_info, current_time, *tiered_base_costs) # Get service tier aware cost keys input_cost_key: Final = get_service_tier_cost_key("input_cost_per_token", service_tier) @@ -734,13 +743,13 @@ def _get_token_base_cost( # For image generation models that don't have output_cost_per_token, # use output_cost_per_image_token as the base cost (all output tokens are image tokens) - if completion_base_cost == 0.0 or completion_base_cost is None: + if completion_base_cost == 0.0: output_image_cost: Final = get_cost_per_unit(model_info, "output_cost_per_image_token", None) if output_image_cost is not None: completion_base_cost = output_image_cost cache_creation_cost = get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None) - cache_creation_cost_above_1hr = get_cost_per_unit( - model_info, "cache_creation_input_token_cost_above_1hr", default_value=None + cache_creation_keys_by_ttl: Mapping[str, str] = MappingProxyType( + {ttl: get_service_tier_cost_key(key, service_tier) for ttl, key in _TTL_PRICED_CACHE_WRITE_RATE_KEYS.items()} ) cache_read_cost = get_cost_per_unit(model_info, cache_read_cost_key, default_value=None) @@ -811,14 +820,6 @@ def _get_token_base_cost( if service_tier else f"cache_creation_input_token_cost_above_{threshold_str}_tokens" ) - cache_creation_1hr_tiered_key = ( - get_service_tier_cost_key( - f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens", - service_tier, - ) - if service_tier - else f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens" - ) cache_read_tiered_key = ( get_service_tier_cost_key( f"cache_read_input_token_cost_above_{threshold_str}_tokens", @@ -830,10 +831,19 @@ def _get_token_base_cost( cache_creation_cost = get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost) - cache_creation_cost_above_1hr = get_cost_per_unit( - model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr + cache_creation_keys_by_ttl = MappingProxyType( + { + ttl: _first_present_key( + model_info, + get_service_tier_cost_key( + f"{_TTL_PRICED_CACHE_WRITE_RATE_KEYS[ttl]}_above_{threshold_str}_tokens", + service_tier, + ), + key, + ) + for ttl, key in cache_creation_keys_by_ttl.items() + } ) - cache_read_cost = get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) except (IndexError, ValueError): continue @@ -849,16 +859,26 @@ def _get_token_base_cost( resolved_cache_creation_cost: Final = ( input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost ) - - return _apply_off_peak_to_base_costs( + base_rates: Final = apply_off_peak_pricing( model_info, current_time, - ( - prompt_base_cost, - completion_base_cost, - resolved_cache_creation_cost, - cache_creation_cost_above_1hr, - resolved_cache_read_cost, + TokenRates( + input_rate=prompt_base_cost, + output_rate=completion_base_cost, + cache_read_rate=resolved_cache_read_cost, + cache_creation_rate=resolved_cache_creation_cost, + reasoning_rate=None, + ), + ) + return replace( + base_rates, + cache_creation_rates_by_ttl=MappingProxyType( + { + ttl: base_rates.cache_creation_rate + if (rate := get_cost_per_unit(model_info, key, default_value=None)) is None + else rate + for ttl, key in cache_creation_keys_by_ttl.items() + } ), ) @@ -975,25 +995,28 @@ def resolve_image_model_info(model: str, custom_llm_provider: str, model_info: M def calculate_cache_writing_cost( cache_creation_tokens: int, cache_creation_token_details: CacheCreationTokenDetails | None, - cache_creation_cost_above_1hr: float, cache_creation_cost: float, + cache_creation_cost_by_ttl: Mapping[str, float], ) -> float: """ - Adjust cost of cache creation tokens based on the cache creation token details. + Bill cache writes per lifetime: Anthropic reports 5m and 1h counts, OpenAI reports 30m and 12h counts. + Writes no lifetime count covers, such as a lifetime a provider adds later, bill at the base write rate. """ - total_cost: float = 0.0 - if cache_creation_token_details is not None: - # get the number of 5m and 1h cache creation tokens - cache_creation_tokens_5m: Final = cache_creation_token_details.ephemeral_5m_input_tokens - cache_creation_tokens_1h: Final = cache_creation_token_details.ephemeral_1h_input_tokens - # add the number of 5m and 1h cache creation tokens to the cache creation tokens - total_cost += cache_creation_tokens_5m * cache_creation_cost if cache_creation_tokens_5m is not None else 0.0 - total_cost += ( - cache_creation_tokens_1h * cache_creation_cost_above_1hr if cache_creation_tokens_1h is not None else 0.0 - ) - else: - total_cost += cache_creation_tokens * cache_creation_cost - return total_cost + if cache_creation_token_details is None: + return cache_creation_tokens * cache_creation_cost + tokens_by_ttl: Final[Mapping[str, int | None]] = MappingProxyType( + { + "5m": cache_creation_token_details.ephemeral_5m_input_tokens, + "30m": cache_creation_token_details.ephemeral_30m_input_tokens, + "1hr": cache_creation_token_details.ephemeral_1h_input_tokens, + "12hr": cache_creation_token_details.ephemeral_12h_input_tokens, + } + ) + unattributed_tokens: Final = max(cache_creation_tokens - sum(t or 0 for t in tokens_by_ttl.values()), 0) + return unattributed_tokens * cache_creation_cost + sum( + (tokens or 0) * cache_creation_cost_by_ttl.get(ttl, cache_creation_cost) + for ttl, tokens in tokens_by_ttl.items() + ) class PromptTokensDetailsResult(TypedDict): @@ -1178,7 +1201,7 @@ def _calculate_input_cost( prompt_base_cost: float, cache_read_cost: float, cache_creation_cost: float, - cache_creation_cost_above_1hr: float, + cache_creation_cost_by_ttl: Mapping[str, float], service_tier: str | None = None, ) -> float: """ @@ -1248,8 +1271,8 @@ def _calculate_input_cost( prompt_cost += calculate_cache_writing_cost( cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"], cache_creation_token_details=prompt_tokens_details["cache_creation_token_details"], - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost, + cache_creation_cost_by_ttl=cache_creation_cost_by_ttl, ) ### CHARACTER COST @@ -1510,34 +1533,31 @@ def generic_cost_per_token( ) billing_time: Final = current_time if current_time is not None else current_billing_time() - ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - cache_read_cost, - ) = _get_token_base_cost( + base_costs: Final = _get_token_base_cost( model_info=resolved_model_info, usage=usage, service_tier=service_tier, current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) + completion_base_cost: Final = base_costs.output_rate resolved_cache_read_cost: Final = resolve_explicit_cache_read_cost_rate( model_info=resolved_model_info, usage=usage, current_time=billing_time, service_tier=service_tier, ) - effective_cache_read_cost: Final = cache_read_cost if resolved_cache_read_cost is None else resolved_cache_read_cost + effective_cache_read_cost: Final = ( + base_costs.cache_read_rate if resolved_cache_read_cost is None else resolved_cache_read_cost + ) prompt_cost = _calculate_input_cost( prompt_tokens_details=prompt_tokens_details, model_info=resolved_model_info, - prompt_base_cost=prompt_base_cost, + prompt_base_cost=base_costs.input_rate, cache_read_cost=effective_cache_read_cost, - cache_creation_cost=cache_creation_cost, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, + cache_creation_cost=base_costs.cache_creation_rate, + cache_creation_cost_by_ttl=base_costs.cache_creation_rates_by_ttl, service_tier=service_tier, ) @@ -1640,7 +1660,7 @@ class BilledTokenRates: cache_read_input_token_cost: float cache_read_input_audio_token_cost: float cache_creation_input_token_cost: float - cache_creation_input_token_cost_above_1hr: float + cache_creation_input_token_cost_by_ttl: Mapping[str, float] output_cost_per_reasoning_token: float def scaled(self, multiplier: float) -> "BilledTokenRates": @@ -1652,7 +1672,9 @@ class BilledTokenRates: cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, cache_read_input_audio_token_cost=self.cache_read_input_audio_token_cost * multiplier, cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, - cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, + cache_creation_input_token_cost_by_ttl=MappingProxyType( + {ttl: rate * multiplier for ttl, rate in self.cache_creation_input_token_cost_by_ttl.items()} + ), output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, ) @@ -1702,7 +1724,7 @@ def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRat cache_read_input_token_cost=cache_read_rate, cache_read_input_audio_token_cost=cache_read_rate, cache_creation_input_token_cost=cache_creation_rate, - cache_creation_input_token_cost_above_1hr=cache_creation_rate, + cache_creation_input_token_cost_by_ttl=MappingProxyType({}), output_cost_per_reasoning_token=output_rate, ) @@ -1717,13 +1739,7 @@ def _cost_map_billed_rates( current_time: datetime | None, ) -> BilledTokenRates: billing_time: Final = current_time if current_time is not None else current_billing_time() - ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost_rate, - cache_creation_cost_above_1hr_rate, - cache_read_cost_rate, - ) = _get_token_base_cost( + base_costs: Final = _get_token_base_cost( model_info=model_info, usage=usage, service_tier=service_tier, @@ -1737,14 +1753,14 @@ def _cost_map_billed_rates( service_tier=service_tier, ) effective_cache_read_cost_rate: Final = ( - cache_read_cost_rate if provider_cache_read_cost_rate is None else provider_cache_read_cost_rate + base_costs.cache_read_rate if provider_cache_read_cost_rate is None else provider_cache_read_cost_rate ) reasoning_rate: Final = _resolve_billed_reasoning_rate( model_info=model_info, usage=usage, service_tier=service_tier, - completion_base_cost=completion_base_cost, + completion_base_cost=base_costs.output_rate, current_time=billing_time, ) audio_cache_read_rate: Final = get_cost_per_unit( @@ -1759,14 +1775,14 @@ def _cost_map_billed_rates( * get_provider_specific_speed_multiplier(model_info=model_info, usage=usage) ) return BilledTokenRates( - input_cost_per_token=prompt_base_cost, - output_cost_per_token=completion_base_cost, + input_cost_per_token=base_costs.input_rate, + output_cost_per_token=base_costs.output_rate, cache_read_input_token_cost=effective_cache_read_cost_rate, cache_read_input_audio_token_cost=( audio_cache_read_rate if audio_cache_read_rate is not None else effective_cache_read_cost_rate ), - cache_creation_input_token_cost=cache_creation_cost_rate, - cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, + cache_creation_input_token_cost=base_costs.cache_creation_rate, + cache_creation_input_token_cost_by_ttl=base_costs.cache_creation_rates_by_ttl, output_cost_per_reasoning_token=reasoning_rate, ).scaled(multiplier) @@ -1854,8 +1870,8 @@ def get_token_type_cost_breakdown( else calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr, cache_creation_cost=rates.cache_creation_input_token_cost, + cache_creation_cost_by_ttl=rates.cache_creation_input_token_cost_by_ttl, ) ) return TokenTypeCostBreakdown( @@ -1886,7 +1902,7 @@ def calculate_prompt_caching_savings( biller saw them rather than at the later spend write. """ model_info_with_cache_read_default: Final = apply_provider_cache_read_default(model_info, custom_llm_provider) - prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost( + base_costs: Final = _get_token_base_cost( model_info=model_info_with_cache_read_default, usage=usage, service_tier=service_tier, @@ -1899,9 +1915,11 @@ def calculate_prompt_caching_savings( current_time=billed_at, service_tier=service_tier, ) - effective_cache_read_cost: Final = cache_read_cost if provider_cache_read_cost is None else provider_cache_read_cost - write_rate: Final = cache_creation_cost or prompt_base_cost - write_rate_1h: Final = cache_creation_cost_above_1hr or write_rate + prompt_base_cost: Final = base_costs.input_rate + effective_cache_read_cost: Final = ( + base_costs.cache_read_rate if provider_cache_read_cost is None else provider_cache_read_cost + ) + write_rate: Final = base_costs.cache_creation_rate or prompt_base_cost prompt_tokens_details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"], 0) cache_creation_tokens: Final = max(prompt_tokens_details["cache_creation_tokens"], 0) @@ -1909,7 +1927,9 @@ def calculate_prompt_caching_savings( cache_creation_details: Final = ( CacheCreationTokenDetails( ephemeral_5m_input_tokens=max(details.ephemeral_5m_input_tokens or 0, 0), + ephemeral_30m_input_tokens=max(details.ephemeral_30m_input_tokens or 0, 0), ephemeral_1h_input_tokens=max(details.ephemeral_1h_input_tokens or 0, 0), + ephemeral_12h_input_tokens=max(details.ephemeral_12h_input_tokens or 0, 0), ) if details is not None else None @@ -1918,8 +1938,10 @@ def calculate_prompt_caching_savings( write_premium: Final = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_details, - cache_creation_cost_above_1hr=write_rate_1h - prompt_base_cost, cache_creation_cost=write_rate - prompt_base_cost, + cache_creation_cost_by_ttl=MappingProxyType( + {ttl: rate - prompt_base_cost for ttl, rate in base_costs.cache_creation_rates_by_ttl.items()} + ), ) uplift: Final = get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift( model_info, vertex_location diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 98504f54533..3f388637a83 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,6 +7,7 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ +from collections.abc import Mapping from dataclasses import dataclass, replace from datetime import datetime from typing import Final @@ -72,7 +73,7 @@ def _flat_rates(model_info: ModelInfo) -> TokenRates: ) -def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: +def _tier_rates(model_info: ModelInfo, tier: Mapping[str, object]) -> TokenRates: # A tier that declares output rates keeps the request on them, all-or-nothing. A tier table # spelling out only input rates would serve every completion for free, so there the model's # own output rates stand in diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 14fe355a805..8c0c6dffdc7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -36795,6 +36795,12 @@ }, "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, + "cache_creation_input_token_cost_above_12hr_priority": 5e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36880,6 +36886,12 @@ }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -36957,6 +36969,12 @@ }, "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, + "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_decisions": 2.5e-07, @@ -37043,6 +37061,12 @@ }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37117,6 +37141,12 @@ }, "gpt-5.6-sol": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37193,6 +37223,12 @@ }, "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -37268,6 +37304,12 @@ }, "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_12hr": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, @@ -74933,6 +74975,12 @@ }, "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0dd978ee855..4e183bf9d5c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -22858,6 +22858,25 @@ "title": "Body_image_edit_api_openai_deployments__model__images_edits_post", "type": "object" }, + "CacheWriteTokensDetails": { + "additionalProperties": true, + "properties": { + "ttl_12h_tokens": { + "title": "Ttl 12H Tokens", + "type": "integer" + }, + "ttl_30m_tokens": { + "title": "Ttl 30M Tokens", + "type": "integer" + } + }, + "required": [ + "ttl_30m_tokens", + "ttl_12h_tokens" + ], + "title": "CacheWriteTokensDetails", + "type": "object" + }, "CachedTokensDetails": { "properties": { "audio_tokens": { @@ -24352,6 +24371,27 @@ ], "title": "Audio Tokens" }, + "cache_write_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Cache Write Tokens" + }, + "cache_write_tokens_details": { + "anyOf": [ + { + "$ref": "#/components/schemas/CacheWriteTokensDetails" + }, + { + "type": "null" + } + ] + }, "cached_tokens": { "default": 0, "title": "Cached Tokens", diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 02d9566d396..a9c7818b577 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -3663,18 +3663,20 @@ class LiteLLMCompletionResponsesConfig: if count is not None } ) - response_usage.input_tokens_details = InputTokensDetails( - cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, - text_tokens=prompt_details.text_tokens, - audio_tokens=prompt_details.audio_tokens, - image_tokens=prompt_details.image_tokens, - video_tokens=prompt_details.video_tokens, - cached_tokens_details=( - cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None - ), - **cache_write_extra, - **cache_type_extra, - **grounding_request_counts, + response_usage.input_tokens_details = InputTokensDetails.model_validate( + { + "cached_tokens": prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, + "text_tokens": prompt_details.text_tokens, + "audio_tokens": prompt_details.audio_tokens, + "image_tokens": prompt_details.image_tokens, + "video_tokens": prompt_details.video_tokens, + "cached_tokens_details": ( + cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None + ), + **cache_write_extra, + **cache_type_extra, + **grounding_request_counts, + } ) # Translate completion_tokens_details to output_tokens_details diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 503d85aaed9..f256480d952 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel from litellm.types.llms.openai import ( AllMessageValues, + InputTokensDetails, OutputTokensDetails, ResponseAPIUsage, ResponseInputParam, @@ -23,6 +24,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import DecodedResponseId from litellm.types.utils import ( + CacheCreationTokenDetails, CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, SpecialEnums, @@ -41,6 +43,17 @@ def _output_token_detail(details: object, field: str) -> int | None: return value if isinstance(value, int) else None +def _cache_write_lifetimes(cache_write_tokens_details: object) -> CacheCreationTokenDetails | None: + """OpenAI Responses splits `cache_write_tokens` into `ttl_30m_tokens` and `ttl_12h_tokens`.""" + ttl_30m_tokens: Final = _output_token_detail(cache_write_tokens_details, "ttl_30m_tokens") + ttl_12h_tokens: Final = _output_token_detail(cache_write_tokens_details, "ttl_12h_tokens") + if ttl_30m_tokens is None and ttl_12h_tokens is None: + return None + return CacheCreationTokenDetails( + ephemeral_30m_input_tokens=ttl_30m_tokens, ephemeral_12h_input_tokens=ttl_12h_tokens + ) + + def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything return isinstance(value, list) @@ -1247,7 +1260,11 @@ class ResponseAPILoggingUtils: prompt_tokens_details: PromptTokensDetailsWrapper | None = None if response_api_usage.input_tokens_details: if isinstance(response_api_usage.input_tokens_details, dict): - prompt_tokens_details = PromptTokensDetailsWrapper(**response_api_usage.input_tokens_details) + input_details: Final = InputTokensDetails.model_validate(response_api_usage.input_tokens_details) + prompt_tokens_details = PromptTokensDetailsWrapper( + **input_details.model_dump(exclude={"cache_write_tokens_details"}), + cache_creation_token_details=_cache_write_lifetimes(input_details.cache_write_tokens_details), + ) else: prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None), @@ -1259,6 +1276,9 @@ class ResponseAPILoggingUtils: ), video_tokens=getattr(response_api_usage.input_tokens_details, "video_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), + cache_creation_token_details=_cache_write_lifetimes( + getattr(response_api_usage.input_tokens_details, "cache_write_tokens_details", None) + ), cache_type=getattr(response_api_usage.input_tokens_details, "cache_type", None), web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 62dc152ca84..9a4ecf37127 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -683,12 +683,12 @@ class ChatCompletionCachedContent(TypedDict): class PromptCacheBreakpoint(TypedDict): mode: ReadOnly[Literal["explicit"]] - ttl: NotRequired[ReadOnly[Literal["30m"]]] + ttl: NotRequired[ReadOnly[Literal["30m", "12h"]]] class PromptCacheOptions(TypedDict, total=False): mode: ReadOnly[Literal["implicit", "explicit"]] - ttl: ReadOnly[Literal["30m"]] + ttl: ReadOnly[Literal["30m", "12h"]] class ChatCompletionThinkingBlock(TypedDict, total=False): @@ -1382,10 +1382,23 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class CacheWriteTokensDetails(BaseLiteLLMOpenAIResponseObject): + ttl_30m_tokens: int + ttl_12h_tokens: int + + model_config = ConfigDict(extra="allow", frozen=True) + + +def _is_none(value: object) -> bool: + return value is None + + class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 cached_tokens_details: CachedTokensDetails | None = None + cache_write_tokens: int | None = Field(default=None, exclude_if=_is_none) + cache_write_tokens_details: CacheWriteTokensDetails | None = Field(default=None, exclude_if=_is_none) image_tokens: int | None = None text_tokens: int | None = None video_tokens: int | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5939a350970..ed66cbd8cf0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -322,6 +322,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_above_272k_tokens_flex: float | None cache_creation_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_creation_input_token_cost_above_1hr: float | None + cache_creation_input_token_cost_above_12hr: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_flex: ReadOnly[float | None] + cache_creation_input_token_cost_above_12hr_priority: ReadOnly[float | None] + 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_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 @@ -1804,7 +1810,9 @@ class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for ol class CacheCreationTokenDetails(LiteLLMBaseModel): ephemeral_5m_input_tokens: int | None = None + ephemeral_30m_input_tokens: int | None = None ephemeral_1h_input_tokens: int | None = None + ephemeral_12h_input_tokens: int | None = None class PromptTokensDetailsWrapper( @@ -3954,6 +3962,12 @@ class MirroredPricingParams(LiteLLMBaseModel): explicit_cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None + cache_creation_input_token_cost_above_12hr: float | None = None + cache_creation_input_token_cost_above_12hr_flex: float | None = None + cache_creation_input_token_cost_above_12hr_priority: float | None = None + 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 tiered_pricing: list[dict[str, Any]] | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 809a94cd93e..58eaa43e687 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6817,6 +6817,24 @@ def get_model_info_helper( cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), + cache_creation_input_token_cost_above_12hr=_model_info.get( + "cache_creation_input_token_cost_above_12hr", None + ), + cache_creation_input_token_cost_above_12hr_flex=_model_info.get( + "cache_creation_input_token_cost_above_12hr_flex", None + ), + cache_creation_input_token_cost_above_12hr_priority=_model_info.get( + "cache_creation_input_token_cost_above_12hr_priority", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens", None + ), + cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex=_model_info.get( + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex", None + ), + 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 + ), 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 14fe355a805..8c0c6dffdc7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -36795,6 +36795,12 @@ }, "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, + "cache_creation_input_token_cost_above_12hr_priority": 5e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36880,6 +36886,12 @@ }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -36957,6 +36969,12 @@ }, "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, + "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_decisions": 2.5e-07, @@ -37043,6 +37061,12 @@ }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37117,6 +37141,12 @@ }, "gpt-5.6-sol": { "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_12hr": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, @@ -37193,6 +37223,12 @@ }, "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, @@ -37268,6 +37304,12 @@ }, "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_12hr": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, @@ -74933,6 +74975,12 @@ }, "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index acba66223e5..53fd5c1cec2 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -6985,6 +6985,106 @@ GPT_5_6_RESPONSES_CACHE_READ: Final = CostTrackingTestCase( ), ) +GPT_5_6_RESPONSES_CACHE_WRITE_12H: Final = CostTrackingTestCase( + name="gpt-5.6-responses_cache_write_12h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "summarize this text", + "prompt_cache_options": {"mode": "implicit", "ttl": "12h"}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + } + ], + "usage": { + "input_tokens": 12928, + "output_tokens": 380, + "total_tokens": 13308, + "input_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 5000, + "cache_write_tokens_details": {"ttl_30m_tokens": 3000, "ttl_12h_tokens": 2000}, + }, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.0344865, + input_cost=0.0291865, + output_cost=0.00532, + prompt_tokens=12928, + completion_tokens=380, + cache_creation_cost=0.0153125, + ), +) + +GPT_6_LUNA_RESPONSES_CACHE_WRITE_12H_NO_12H_RATE: Final = CostTrackingTestCase( + name="gpt-6-luna-responses_cache_write_12h_no_12h_rate", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-6-luna", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "summarize this text", + "prompt_cache_options": {"mode": "implicit", "ttl": "12h"}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6-luna", + "output": [ + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + } + ], + "usage": { + "input_tokens": 12928, + "output_tokens": 380, + "total_tokens": 13308, + "input_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 5000, + "cache_write_tokens_details": {"ttl_30m_tokens": 3000, "ttl_12h_tokens": 2000}, + }, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.0016078, + input_cost=0.0014178, + output_cost=0.00019, + prompt_tokens=12928, + completion_tokens=380, + cache_creation_cost=0.000625, + ), +) + GPT_5_6_RESPONSES_REASONING: Final = CostTrackingTestCase( name="gpt-5.6-responses_reasoning", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -11913,6 +12013,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( OMNI_MODERATIONS_NEXT_SINGLE, OMNI_MODERATIONS_NEXT_LIST, GPT_5_6_RESPONSES_CACHE_READ, + GPT_5_6_RESPONSES_CACHE_WRITE_12H, + GPT_6_LUNA_RESPONSES_CACHE_WRITE_12H_NO_12H_RATE, GPT_5_6_RESPONSES_REASONING, GPT_5_6_RESPONSES_STREAM, GPT_5_6_RESPONSES_STREAM_CACHE_READ, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 012d847e676..cd979528fe6 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -174,6 +174,8 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { }, }, "gpt-5.6": { + "cache_creation_input_token_cost": 2.1875e-06, + "cache_creation_input_token_cost_above_12hr": 4.375e-06, "cache_read_input_token_cost": 1.75e-07, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 1.75e-06, diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index e98f6a44222..e978a9a316a 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -85,6 +85,12 @@ class CostMapEntry(BaseModel): cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None + cache_creation_input_token_cost_above_12hr: float | None = None + cache_creation_input_token_cost_above_12hr_flex: float | None = None + cache_creation_input_token_cost_above_12hr_priority: float | None = None + 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_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/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 c99147d2e63..1729db94ed8 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 @@ -65,8 +65,8 @@ def test_missing_cache_read_rate_resolves_to_input_rate(prompt_tokens, read_rate prompt_cost, _ = generic_cost_per_token( "policy-fixture", usage, "openai", service_tier=service_tier, model_info=info ) - assert billed[4] == pytest.approx(read_rate if read_rate is not None else billed[0]) - assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4]) + assert billed.cache_read_rate == pytest.approx(read_rate if read_rate is not None else billed.input_rate) + assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed.input_rate + 100 * billed.cache_read_rate) def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_read_rate() -> None: @@ -248,7 +248,7 @@ def test_missing_cache_read_uses_off_peak_input_rate(): } when = datetime(2026, 9, 7, 12, tzinfo=timezone.utc) billed = _get_token_base_cost(info, Usage(prompt_tokens=100), current_time=when) - assert billed[0] == billed[4] == 5e-6 + assert billed.input_rate == billed.cache_read_rate == 5e-6 def test_reasoning_tokens_no_price_set(_local_model_cost_map): @@ -324,8 +324,6 @@ def test_reasoning_tokens_gemini(_local_model_cost_map): ) - - def test_image_tokens_with_custom_pricing(): """Test that image_tokens in completion are properly costed with output_cost_per_image_token.""" from unittest.mock import patch @@ -578,7 +576,7 @@ def test_get_token_base_cost_picks_highest_crossed_tier(): } usage = Usage(prompt_tokens=150_000, completion_tokens=10, total_tokens=150_010) - prompt_base_cost = _get_token_base_cost(model_info, usage)[0] + prompt_base_cost = _get_token_base_cost(model_info, usage).input_rate assert prompt_base_cost == 9e-6 @@ -779,12 +777,12 @@ def test_get_token_base_cost_weekend_only_off_peak_rate(): saturday_peak_hours = _get_token_base_cost( model_info, usage, current_time=datetime(2026, 8, 29, 2, 0, tzinfo=timezone.utc) ) - assert saturday_peak_hours[:2] == (5e-7, 1e-6) + assert (saturday_peak_hours.input_rate, saturday_peak_hours.output_rate) == (5e-7, 1e-6) monday_same_hours = _get_token_base_cost( model_info, usage, current_time=datetime(2026, 8, 24, 2, 0, tzinfo=timezone.utc) ) - assert monday_same_hours[:2] == (1e-6, 2e-6) + assert (monday_same_hours.input_rate, monday_same_hours.output_rate) == (1e-6, 2e-6) def test_get_token_base_cost_applies_off_peak_pricing(): @@ -810,14 +808,14 @@ def test_get_token_base_cost_applies_off_peak_pricing(): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert off_peak[0] == 5e-7 - assert off_peak[1] == 1e-6 - assert off_peak[4] == 5e-8 + assert off_peak.input_rate == 5e-7 + assert off_peak.output_rate == 1e-6 + assert off_peak.cache_read_rate == 5e-8 peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert peak[0] == 1e-6 - assert peak[1] == 2e-6 - assert peak[4] == 1e-7 + assert peak.input_rate == 1e-6 + assert peak.output_rate == 2e-6 + assert peak.cache_read_rate == 1e-7 def test_get_token_base_cost_non_mapping_off_peak_block_bills_standard_rates(): @@ -843,8 +841,8 @@ def test_get_token_base_cost_non_mapping_off_peak_block_bills_standard_rates(): }, ) result = _get_token_base_cost(model_info, usage, current_time=when) - assert result[0] == 1e-6 - assert result[1] == 2e-6 + assert result.input_rate == 1e-6 + assert result.output_rate == 2e-6 def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): @@ -864,8 +862,8 @@ def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) result = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert result[0] == 5e-7 - assert result[1] == 2e-6 + assert result.input_rate == 5e-7 + assert result.output_rate == 2e-6 def test_get_token_base_cost_off_peak_wins_over_threshold(): @@ -891,12 +889,12 @@ def test_get_token_base_cost_off_peak_wins_over_threshold(): usage = Usage(prompt_tokens=250000, completion_tokens=250000, total_tokens=500000) off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert off_peak[0] == 5e-7 - assert off_peak[1] == 1e-6 + assert off_peak.input_rate == 5e-7 + assert off_peak.output_rate == 1e-6 peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert peak[0] == 3e-6 - assert peak[1] == 4e-6 + assert peak.input_rate == 3e-6 + assert peak.output_rate == 4e-6 def test_get_model_info_propagates_off_peak_fields(): @@ -948,10 +946,10 @@ def test_get_token_base_cost_off_peak_wins_over_tiered_pricing(): usage = Usage(prompt_tokens=1_000, completion_tokens=100, total_tokens=1_100) inside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) - assert inside[:2] == (5e-7, 1e-6) + assert (inside.input_rate, inside.output_rate) == (5e-7, 1e-6) outside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert outside[:2] == (3e-6, 6e-6) + assert (outside.input_rate, outside.output_rate) == (3e-6, 6e-6) def _register_off_peak_reasoning_model( @@ -1153,17 +1151,17 @@ def test_get_token_base_cost_off_peak_cache_creation_rate(): inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) inside = _get_token_base_cost(model_info, usage, current_time=inside_window) - assert inside[2] == 5e-7 - assert inside[3] == 2e-6 + assert inside.cache_creation_rate == 5e-7 + assert inside.cache_creation_rates_by_ttl["1hr"] == 2e-6 outside = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) - assert outside[2] == 1.25e-6 + assert outside.cache_creation_rate == 1.25e-6 without_key = cast( ModelInfo, {**model_info, "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}}, ) - assert _get_token_base_cost(without_key, usage, current_time=inside_window)[2] == 1.25e-6 + assert _get_token_base_cost(without_key, usage, current_time=inside_window).cache_creation_rate == 1.25e-6 def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates(): @@ -1583,19 +1581,14 @@ def test_generic_cost_per_token_tier_1hr_cache_rate_preserves_zero_and_fallback( ), ) costs: Final = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openrouter") - breakdown: Final = get_token_type_cost_breakdown( - model=model, usage=usage, custom_llm_provider="openrouter" - ) + breakdown: Final = get_token_type_cost_breakdown(model=model, usage=usage, custom_llm_provider="openrouter") ordinary_rate: Final = cache_rates.get("cache_creation_input_token_cost") - expected_write_cost: Final = ( - (20 - one_hour_tokens) * (ordinary_rate if ordinary_rate is not None else 1.0) - + one_hour_tokens * expected_write_rate - ) + expected_write_cost: Final = (20 - one_hour_tokens) * ( + ordinary_rate if ordinary_rate is not None else 1.0 + ) + one_hour_tokens * expected_write_rate assert costs == pytest.approx((40 * 1.0 + 40 * 0.1 + expected_write_cost, 10 * 2.0)) assert breakdown.cache_creation_cost == pytest.approx(expected_write_cost) - assert breakdown.rates is not None - assert breakdown.rates.cache_creation_input_token_cost_above_1hr == expected_write_rate finally: litellm.model_cost.pop(model, None) @@ -2020,8 +2013,8 @@ def test_calculate_cache_writing_cost(): result = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost, + cache_creation_cost_by_ttl={"1hr": cache_creation_cost_above_1hr, "12hr": cache_creation_cost_above_1hr}, ) # Expected calculation: @@ -2041,8 +2034,8 @@ def test_calculate_cache_writing_cost(): result_no_details = calculate_cache_writing_cost( cache_creation_tokens=cache_creation_tokens_no_details, cache_creation_token_details=cache_creation_token_details_none, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost_fallback, + cache_creation_cost_by_ttl={"1hr": cache_creation_cost_above_1hr, "12hr": cache_creation_cost_above_1hr}, ) # Expected calculation when no details: 1000 * 5e-06 = 0.005 @@ -2059,15 +2052,13 @@ def test_calculate_cache_writing_cost(): result_partial = calculate_cache_writing_cost( cache_creation_tokens=500, cache_creation_token_details=cache_creation_token_details_partial, - cache_creation_cost_above_1hr=6e-06, cache_creation_cost=3e-06, + cache_creation_cost_by_ttl={"1hr": 6e-06, "12hr": 6e-06}, ) - # Expected calculation: 0 (for None 5m tokens) + (100 * 6e-06) = 0.0006 - expected_cost_partial = (0.0) + (100 * 6e-06) + expected_cost_partial = (400 * 3e-06) + (100 * 6e-06) assert round(result_partial, 6) == round(expected_cost_partial, 6) - assert round(result_partial, 6) == 0.0006 # Test case 4: Zero costs result_zero = calculate_cache_writing_cost( @@ -2075,13 +2066,119 @@ def test_calculate_cache_writing_cost(): cache_creation_token_details=CacheCreationTokenDetails( ephemeral_5m_input_tokens=50, ephemeral_1h_input_tokens=950 ), - cache_creation_cost_above_1hr=0.0, cache_creation_cost=0.0, + cache_creation_cost_by_ttl={"1hr": 0.0, "12hr": 0.0}, ) assert result_zero == 0.0 +@pytest.mark.parametrize("field", CacheCreationTokenDetails.model_fields) +def test_calculate_cache_writing_cost_bills_every_reported_lifetime(field: str) -> None: + cost = calculate_cache_writing_cost( + cache_creation_tokens=10, + cache_creation_token_details=CacheCreationTokenDetails(**{field: 10}), + cache_creation_cost=2.0, + cache_creation_cost_by_ttl={"1hr": 2.0, "12hr": 2.0}, + ) + assert cost == 20.0 + + +def test_calculate_cache_writing_cost_bills_writes_of_an_unknown_lifetime_at_the_base_rate() -> None: + cost = calculate_cache_writing_cost( + cache_creation_tokens=10, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_30m_input_tokens=2, ephemeral_12h_input_tokens=3 + ), + cache_creation_cost=1.0, + cache_creation_cost_by_ttl={"1hr": 4.0, "12hr": 4.0}, + ) + assert cost == 2 * 1.0 + 3 * 4.0 + 5 * 1.0 + + +@pytest.mark.parametrize( + ("prompt_tokens", "service_tier", "write_rate", "twelve_hour_rate"), + [ + (10_000, "default", 2e-6, 4e-6), + (10_000, "flex", 1e-6, 2e-6), + (10_000, "priority", 3e-6, 6e-6), + (300_000, "default", 4e-6, 8e-6), + (300_000, "flex", 2e-6, 4e-6), + (300_000, "priority", 6e-6, 12e-6), + ], +) +def test_generic_cost_per_token_prices_12h_cache_writes_at_prompt_threshold( + prompt_tokens: int, service_tier: str, write_rate: float, twelve_hour_rate: float +) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 2e-6, + "cache_creation_input_token_cost_flex": 1e-6, + "cache_creation_input_token_cost_priority": 3e-6, + "cache_creation_input_token_cost_above_272k_tokens": 4e-6, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2e-6, + "cache_creation_input_token_cost_above_272k_tokens_priority": 6e-6, + "cache_creation_input_token_cost_above_12hr": 4e-6, + "cache_creation_input_token_cost_above_12hr_flex": 2e-6, + "cache_creation_input_token_cost_above_12hr_priority": 6e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 8e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 4e-6, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 12e-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 + ), + ), + ) + + result: Final = generic_cost_per_token( + model="test-12h-cache-write", + usage=usage, + custom_llm_provider="openai", + service_tier=None if service_tier == "default" else service_tier, + model_info=model_info, + ) + + assert result[0] == pytest.approx( + (usage.prompt_tokens - usage.prompt_tokens_details.cache_creation_tokens) * model_info["input_cost_per_token"] + + 3_000 * write_rate + + 2_000 * twelve_hour_rate + ) + + +def test_generic_cost_per_token_falls_back_to_default_for_missing_12h_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 3e-6, + } + usage: Final = Usage( + prompt_tokens=100, + completion_tokens=0, + total_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper( + cache_creation_tokens=100, + cache_creation_token_details=CacheCreationTokenDetails(ephemeral_12h_input_tokens=100), + ), + ) + + result: Final = generic_cost_per_token( + model="test-12h-cache-write-fallback", + usage=usage, + custom_llm_provider="openai", + model_info=model_info, + ) + + assert result[0] == pytest.approx(100 * model_info["cache_creation_input_token_cost"]) + + 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 @@ -2119,7 +2216,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): prompt_base_cost=0.0, cache_read_cost=0.0, cache_creation_cost=cache_creation_cost, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, + cache_creation_cost_by_ttl={"1hr": cache_creation_cost_above_1hr, "12hr": cache_creation_cost_above_1hr}, ) # Expected: (100 * 3.75e-06) + (200 * 6e-06) = 0.000375 + 0.0012 = 0.001575 @@ -2330,10 +2427,6 @@ def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(_local_mo assert round(cost, 10) == round(expected_cost, 10) - - - - def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map): usage = Usage( prompt_tokens=0, @@ -2512,8 +2605,6 @@ def test_vertex_global_or_absent_location_no_uplift(vertex_location, _local_mode assert base == located - - def test_vertex_uplift_invalid_multiplier_defaults_to_one(): """A malformed multiplier in the cost map degrades to base pricing, never raises.""" from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -2569,10 +2660,10 @@ def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None: model_info=info, usage=Usage(prompt_tokens=300_001, completion_tokens=3, total_tokens=300_004), ) - assert rates[0] == 0.001 - assert rates[1] == (0.005 if field == "output_cost_per_token" else 0.002) - assert rates[2] == (0.005 if field == "cache_creation_input_token_cost" else 0.004) - assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003) + assert rates.input_rate == 0.001 + assert rates.output_rate == (0.005 if field == "output_cost_per_token" else 0.002) + assert rates.cache_creation_rate == (0.005 if field == "cache_creation_input_token_cost" else 0.004) + assert rates.cache_read_rate == (0.005 if field == "cache_read_input_token_cost" else 0.003) @pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001)) @@ -2599,10 +2690,12 @@ def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, h model_info=info, usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3), ) - assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001) - assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002) - assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004) - assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003) + assert rates.input_rate == (0.005 if prompt_tokens > 200_000 else 0.001) + assert rates.output_rate == ( + higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002 + ) + assert rates.cache_creation_rate == (0.007 if prompt_tokens > 200_000 else 0.004) + assert rates.cache_read_rate == (higher_rate if prompt_tokens > 512_000 else 0.003) def test_threshold_keys_exclude_service_tier_variants(): @@ -2626,8 +2719,7 @@ def test_threshold_keys_exclude_service_tier_variants(): }, ) usage = Usage(prompt_tokens=350_000, completion_tokens=1_000, total_tokens=351_000) - prompt_base, *_ = _get_token_base_cost(model_info=model_info, usage=usage) - assert prompt_base == 3e-6 + assert _get_token_base_cost(model_info=model_info, usage=usage).input_rate == 3e-6 @pytest.mark.parametrize( @@ -2897,7 +2989,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp cache_read_input_token_cost=6e-7, cache_read_input_audio_token_cost=6e-7, cache_creation_input_token_cost=7.5e-6, - cache_creation_input_token_cost_above_1hr=7.5e-6, + cache_creation_input_token_cost_by_ttl={"1hr": 7.5e-6, "12hr": 7.5e-6}, output_cost_per_reasoning_token=3e-5, ) assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost) @@ -3713,6 +3805,23 @@ def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a 3e-7, id="no-1h-price-uses-the-off-peak-write-price", ), + pytest.param( + { + "tiered_pricing": [ + { + "range": [0, 128000], + "input_cost_per_token": 2e-7, + "output_cost_per_token": 1.25e-6, + "cache_creation_input_token_cost": 2.5e-7, + } + ], + "off_peak_pricing": {"hours_utc": "00:00-23:59", "cache_creation_input_token_cost": 3e-7}, + }, + datetime(2026, 9, 14, 12, tzinfo=timezone.utc), + 3e-7, + 3e-7, + id="tiered-no-1h-price-uses-the-off-peak-write-price", + ), ), ) def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path( @@ -3724,10 +3833,17 @@ def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_ model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates} usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) - _, _, creation, creation_1h, _ = _get_token_base_cost(model_info, usage, current_time=current_time) + rates = _get_token_base_cost(model_info, usage, current_time=current_time) - assert creation == pytest.approx(expected_creation) - assert creation_1h == pytest.approx(expected_creation_1h) + assert rates.cache_creation_rate == pytest.approx(expected_creation) + assert calculate_cache_writing_cost( + cache_creation_tokens=3, + cache_creation_token_details=CacheCreationTokenDetails( + ephemeral_1h_input_tokens=1, ephemeral_12h_input_tokens=2 + ), + cache_creation_cost=rates.cache_creation_rate, + cache_creation_cost_by_ttl=rates.cache_creation_rates_by_ttl, + ) == pytest.approx(3 * expected_creation_1h) def _image_response(num_images: int = 1, usage: ImageUsage | None = None) -> ImageResponse: @@ -3787,8 +3903,6 @@ def test_route_image_generation_cost_openai_honors_deployment_input_cost_per_ima assert cost == pytest.approx(0.07) - - @pytest.mark.parametrize( ("custom_llm_provider", "model"), [ @@ -4206,17 +4320,28 @@ def test_token_breakdown_uses_supplied_pricing_without_catalog_lookup(rate: floa "cache_creation_input_token_cost": rate, } usage: Final = Usage( - prompt_tokens=20, completion_tokens=3, total_tokens=23, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), cache_creation_input_tokens=5, + prompt_tokens=20, + completion_tokens=3, + total_tokens=23, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), + cache_creation_input_tokens=5, ) prompt_cost, output_cost = generic_cost_per_token( - model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info, + model="supplied-pricing-test", + usage=usage, + custom_llm_provider="bedrock", + model_info=model_info, ) breakdown: Final = get_token_type_cost_breakdown( - model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info, + model="supplied-pricing-test", + usage=usage, + custom_llm_provider="bedrock", + model_info=model_info, ) assert breakdown.cache_read_cost == pytest.approx(5 * rate) assert breakdown.cache_creation_cost == pytest.approx(5 * rate) assert breakdown.rates is not None - assert prompt_cost == pytest.approx(10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost) + assert prompt_cost == pytest.approx( + 10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost + ) assert output_cost == pytest.approx(3 * breakdown.rates.output_cost_per_token) diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 30e13442c92..023d0ea26fb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -3398,7 +3398,9 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is not None assert response_usage.input_tokens_details.cached_tokens == 5 assert response_usage.input_tokens_details.text_tokens == 8 - assert "cache_write_tokens" not in response_usage.input_tokens_details.model_dump() + assert {"cache_write_tokens", "cache_write_tokens_details"}.isdisjoint( + response_usage.input_tokens_details.model_dump() + ) def test_transform_usage_with_cached_tokens_gemini(self): """Test that cached_tokens from Gemini are properly transformed to input_tokens_details""" diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 6ee2c495705..97a39f64001 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -1,5 +1,6 @@ from importlib import import_module import base64 +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -7,10 +8,11 @@ import pytest import litellm from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse -from litellm.types.utils import Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import _invalidate_model_cost_lowercase_map @@ -413,6 +415,43 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens_details.cache_creation_tokens == 10059 assert result.prompt_tokens_details.cached_tokens == 0 + @pytest.mark.parametrize("as_model", [False, True]) + def test_transform_response_api_usage_prices_12h_cache_writes(self, as_model: bool) -> None: + usage_data: Final = { + "input_tokens": 5_000, + "output_tokens": 0, + "total_tokens": 5_000, + "input_tokens_details": { + "cache_write_tokens": 5_000, + "cache_write_tokens_details": {"ttl_30m_tokens": 3_000, "ttl_12h_tokens": 2_000}, + }, + } + usage_input: Final = ResponseAPIUsage(**usage_data) if as_model else usage_data + usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_input) + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 3e-6, + "cache_creation_input_token_cost_above_12hr": 5e-6, + } + + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cache_write_tokens == 5_000 + assert usage.prompt_tokens_details.cache_creation_token_details is not None + assert usage.prompt_tokens_details.cache_creation_token_details.ephemeral_12h_input_tokens == 2_000 + assert usage.prompt_tokens_details.cache_creation_token_details.ephemeral_30m_input_tokens == 3_000 + cost: Final = generic_cost_per_token( + model="test-12h-cache-write", + usage=usage, + custom_llm_provider="openai", + model_info=model_info, + )[0] + + assert cost == pytest.approx( + 3_000 * model_info["cache_creation_input_token_cost"] + + 2_000 * model_info["cache_creation_input_token_cost_above_12hr"] + ) + def test_transform_response_api_usage_mixed_details(self): """Test transformation handles mixed token details (cached + image + audio).""" # Setup - hypothetical usage with mixed token types diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index ad069aa987b..fcd6b8aa402 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -939,6 +939,18 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "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"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1d36f5fd9a2..da9e77e9b58 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28439,6 +28439,15 @@ export interface components { */ uncached_input_tokens: number; }; + /** CacheWriteTokensDetails */ + CacheWriteTokensDetails: { + /** Ttl 12H Tokens */ + ttl_12h_tokens: number; + /** Ttl 30M Tokens */ + ttl_30m_tokens: number; + } & { + [key: string]: unknown; + }; /** CachedTokensDetails */ CachedTokensDetails: { /** Audio Tokens */ @@ -33589,6 +33598,9 @@ export interface components { InputTokensDetails: { /** Audio Tokens */ audio_tokens?: number | null; + /** Cache Write Tokens */ + cache_write_tokens?: number | null; + cache_write_tokens_details?: components["schemas"]["CacheWriteTokensDetails"] | null; /** * Cached Tokens * @default 0 @@ -35207,6 +35219,18 @@ export interface components { cache_creation_input_token_cost_above_100k_tokens?: number | null; /** Cache Creation Input Token Cost Above 100K Tokens Batches */ cache_creation_input_token_cost_above_100k_tokens_batches?: number | null; + /** Cache Creation Input Token Cost Above 12Hr */ + 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 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 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 1Hr */ cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 1Hr Above 100K Tokens */ @@ -41506,9 +41530,9 @@ export interface components { mode?: "implicit" | "explicit"; /** * Ttl - * @constant + * @enum {string} */ - ttl?: "30m"; + ttl?: "30m" | "12h"; }; /** PromptCachingRequest */ PromptCachingRequest: { @@ -50199,6 +50223,18 @@ export interface components { blocked?: boolean | null; /** Cache Creation Input Token Cost */ cache_creation_input_token_cost?: number | null; + /** Cache Creation Input Token Cost Above 12Hr */ + 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 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 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 1Hr */ cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Read Input Token Cost */ @@ -50492,6 +50528,18 @@ export interface components { cache_creation_input_token_cost_above_100k_tokens?: number | null; /** Cache Creation Input Token Cost Above 100K Tokens Batches */ cache_creation_input_token_cost_above_100k_tokens_batches?: number | null; + /** Cache Creation Input Token Cost Above 12Hr */ + 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 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 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 1Hr */ cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 1Hr Above 100K Tokens */ From dc0237ccd722c7e4d010b9a29ff4ba22af217f5b Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 09:12:31 -0700 Subject: [PATCH 02/43] fix(cost): bill gemini audio input and cached audio at flex, priority and batch rates (#48086) * fix(gemini): price tiered audio token rates * fix: correct Gemini audio pricing catalog * fix(gemini): preserve cached modality pricing --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 6 +- .../llm_cost_calc/zero_cost_diagnostic.py | 29 +++- .../vertex_and_google_ai_studio_gemini.py | 31 ++-- ...odel_prices_and_context_window_backup.json | 16 ++ litellm/types/utils.py | 10 ++ litellm/utils.py | 9 ++ model_prices_and_context_window.json | 16 ++ .../llm_cost_calc/test_utils.py | 62 ++++++++ .../test_zero_cost_diagnostic.py | 31 ++++ ...test_vertex_and_google_ai_studio_gemini.py | 147 +++++++++++++++++- tests/unit/test_cost_calculator.py | 77 +++++---- tests/unit/test_utils.py | 4 + tests/unit/test_utils_get_model_info.py | 24 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 20 +++ 14 files changed, 429 insertions(+), 53 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 99944662bf8..98f6f1409f2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2879,6 +2879,7 @@ def _batch_rate( model_info: ModelInfo, key: Literal[ "input_cost_per_audio_token_batches", + "cache_read_input_audio_token_cost_batches", "input_cost_per_image_token_batches", "input_cost_per_video_token_batches", ], @@ -3012,6 +3013,7 @@ def _batch_prompt_cost( ) -> 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 audio_tokens, image_tokens, video_tokens = ( details["audio_tokens"], @@ -3026,7 +3028,9 @@ def _batch_prompt_cost( + audio_tokens * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate) + image_tokens * _batch_rate(model_info, "input_cost_per_image_token_batches", input_rate) + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate) - + cached_tokens * (cache_read_rate or 0.0) + + (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) ) diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py index 9688b511ea8..3ed598b5e97 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -5,6 +5,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.utils import ( CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, @@ -16,6 +17,8 @@ ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" _TEXT_INPUT_RATE: Final = "input_cost_per_token" _AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token" +_CACHE_READ_INPUT_RATE: Final = "cache_read_input_token_cost" +_CACHE_READ_AUDIO_INPUT_RATE: Final = "cache_read_input_audio_token_cost" _TEXT_OUTPUT_RATE: Final = "output_cost_per_token" _AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token" _RATE_KEY_MARKERS: Final = ("cost", "pricing") @@ -33,19 +36,38 @@ def _tokens(value: object) -> int: def used_pricing_keys(usage: Usage) -> tuple[str, ...]: - prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details) + prompt_details: Final = parse_prompt_tokens_details(usage) + prompt_audio: Final = prompt_details["audio_tokens"] completion_audio: Final = _audio_tokens(usage.completion_tokens_details) - prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio + prompt_text: Final = max( + _tokens(usage.prompt_tokens) + - prompt_details["cache_hit_tokens"] + - prompt_audio + - prompt_details["image_tokens"] + - prompt_details["video_tokens"], + 0, + ) + cached_text: Final = max(prompt_details["cache_hit_tokens"] - prompt_details["cache_hit_audio_tokens"], 0) completion_text: Final = _tokens(usage.completion_tokens) - completion_audio components: Final = ( (_TEXT_INPUT_RATE, prompt_text), (_AUDIO_INPUT_RATE, prompt_audio), + (_CACHE_READ_INPUT_RATE, cached_text), + (_CACHE_READ_AUDIO_INPUT_RATE, prompt_details["cache_hit_audio_tokens"]), (_TEXT_OUTPUT_RATE, completion_text), (_AUDIO_OUTPUT_RATE, completion_audio), ) return tuple(key for key, count in components if count > 0) +def _effective_pricing_key(key: str, pricing_entry: Mapping[str, object]) -> str: + if key == _CACHE_READ_AUDIO_INPUT_RATE and pricing_entry.get(key) is None: + return _CACHE_READ_INPUT_RATE if pricing_entry.get(_CACHE_READ_INPUT_RATE) is not None else _TEXT_INPUT_RATE + if key == _CACHE_READ_INPUT_RATE and pricing_entry.get(key) is None: + return _TEXT_INPUT_RATE + return key + + def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None: try: return _NESTED_PRICING.validate_python(value) @@ -89,7 +111,8 @@ def diagnose_zero_cost( pricing_entry: Mapping[str, object], calculation_failed: bool, ) -> StandardLoggingZeroCostDiagnostic | None: - used_keys: Final = used_pricing_keys(usage) + effective_keys: Final = tuple(_effective_pricing_key(key, pricing_entry) for key in used_pricing_keys(usage)) + used_keys: Final = tuple(key for index, key in enumerate(effective_keys) if key not in effective_keys[:index]) if not used_keys: return None missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ce632bc0eef..d788ccb0f1d 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -42,6 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.types.llms.anthropic import AnthropicThinkingParam +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.gemini import BidiGenerateContentServerMessage from litellm.types.llms.openai import ( AllMessageValues, @@ -2364,26 +2365,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif modality == "DOCUMENT": cached_text_tokens = (cached_text_tokens or 0) + token_count - ## Calculate non-cached tokens by subtracting cached from total (per modality) - ## This is necessary because promptTokensDetails includes both cached and non-cached tokens - ## See: https://github.com/BerriAI/litellm/issues/18750 - if cached_text_tokens is not None and prompt_text_tokens is not None: - # Explicit caching: subtract cached tokens per modality from cacheTokensDetails - prompt_text_tokens = prompt_text_tokens - cached_text_tokens - elif ( + if ( cached_tokens is not None and prompt_text_tokens is not None and cached_text_tokens is None and "cacheTokensDetails" not in usage_metadata ): - # Implicit caching: only cachedContentTokenCount is provided (no cacheTokensDetails) - # Subtract from text tokens since implicit caching is primarily for text content - # See: https://github.com/BerriAI/litellm/issues/16341 prompt_text_tokens = prompt_text_tokens - cached_tokens - if cached_audio_tokens is not None and prompt_audio_tokens is not None: - prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens - if cached_image_tokens is not None and prompt_image_tokens is not None: - prompt_image_tokens = prompt_image_tokens - cached_image_tokens if cached_video_tokens is not None and prompt_video_tokens is not None: prompt_video_tokens = prompt_video_tokens - cached_video_tokens @@ -2404,9 +2392,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): billable_text_tokens: Final = ( prompt_text_tokens if prompt_text_tokens is None else prompt_text_tokens + billable_tool_use_prompt_tokens ) + cached_tokens_details_fields: Final = MappingProxyType( + { + key: token_count + for key, token_count in ( + ("text_tokens", cached_text_tokens), + ("audio_tokens", cached_audio_tokens), + ("image_tokens", cached_image_tokens), + ) + if token_count is not None + } + ) + cached_tokens_details: Final = ( + CachedTokensDetails.model_validate(cached_tokens_details_fields) if cached_tokens_details_fields else None + ) prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cached_tokens, + cached_tokens_details=cached_tokens_details, audio_tokens=prompt_audio_tokens, text_tokens=billable_text_tokens, image_tokens=prompt_image_tokens, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8c0c6dffdc7..15dfd196c6f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -31137,11 +31137,15 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31466,11 +31470,15 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_audio_token_cost": 3e-08, + "cache_read_input_audio_token_cost_batches": 3e-08, + "cache_read_input_audio_token_cost_flex": 3e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, + "input_cost_per_audio_token_flex": 1.5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31760,12 +31768,16 @@ }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_audio_token_cost": 5e-08, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "cache_read_input_audio_token_cost_priority": 9e-08, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, + "input_cost_per_audio_token_flex": 2.5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -31886,10 +31898,14 @@ }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ed66cbd8cf0..270d9b43dea 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -333,6 +333,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None cache_read_input_audio_token_cost: ReadOnly[float | None] + cache_read_input_audio_token_cost_flex: ReadOnly[float | None] + cache_read_input_audio_token_cost_priority: ReadOnly[float | None] cache_read_input_image_token_cost: ReadOnly[float | None] explicit_cache_read_input_token_cost: ReadOnly[float | None] cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing @@ -348,6 +350,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_read_input_token_cost_above_512k_tokens: float | None cache_read_input_token_cost_batches: ReadOnly[float | None] + cache_read_input_audio_token_cost_batches: ReadOnly[float | None] cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] @@ -361,6 +364,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): off_peak_pricing: ReadOnly[OffPeakPricing | None] # time-windowed off-peak rates input_cost_per_character: float | None # only for vertex ai models input_cost_per_audio_token: float | None + input_cost_per_audio_token_flex: ReadOnly[float | None] + input_cost_per_audio_token_priority: ReadOnly[float | None] input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models input_cost_per_token_above_200k_tokens: float | None # only for vertex ai gemini-2.5-pro models input_cost_per_token_above_100k_tokens: ReadOnly[float | None] @@ -4020,14 +4025,19 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_100k_tokens_batches: float | None = None cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None + cache_read_input_audio_token_cost_batches: float | None = None cache_creation_input_token_cost_batches: float | None = None cache_creation_input_token_cost_above_100k_tokens_batches: float | None = None cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None cache_read_input_audio_token_cost: float | None = None + cache_read_input_audio_token_cost_flex: float | None = None + cache_read_input_audio_token_cost_priority: float | None = None cache_read_input_image_token_cost: float | None = None input_cost_per_character_above_128k_tokens: float | None = None input_cost_per_audio_token: float | None = None + input_cost_per_audio_token_flex: float | None = None + input_cost_per_audio_token_priority: float | None = None input_cost_per_token_cache_hit: float | None = None input_cost_per_token_above_128k_tokens: float | None = None input_cost_per_token_above_100k_tokens: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 58eaa43e687..bb94ca6d40a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6773,6 +6773,13 @@ def get_model_info_helper( ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None), + cache_read_input_audio_token_cost_flex=_model_info.get("cache_read_input_audio_token_cost_flex", None), + cache_read_input_audio_token_cost_priority=_model_info.get( + "cache_read_input_audio_token_cost_priority", None + ), + cache_read_input_audio_token_cost_batches=_model_info.get( + "cache_read_input_audio_token_cost_batches", None + ), explicit_cache_read_input_token_cost=_model_info.get("explicit_cache_read_input_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( @@ -6858,6 +6865,8 @@ def get_model_info_helper( input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_text_message=_model_info.get("input_cost_per_text_message", None), input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), + input_cost_per_audio_token_flex=_model_info.get("input_cost_per_audio_token_flex", None), + input_cost_per_audio_token_priority=_model_info.get("input_cost_per_audio_token_priority", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8c0c6dffdc7..15dfd196c6f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -31137,11 +31137,15 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31466,11 +31470,15 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_audio_token_cost": 3e-08, + "cache_read_input_audio_token_cost_batches": 3e-08, + "cache_read_input_audio_token_cost_flex": 3e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, + "input_cost_per_audio_token_flex": 1.5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31760,12 +31768,16 @@ }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_audio_token_cost": 5e-08, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "cache_read_input_audio_token_cost_priority": 9e-08, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, + "input_cost_per_audio_token_flex": 2.5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -31886,10 +31898,14 @@ }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index 142d49bfff1..b6c18e0da86 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -11,6 +11,7 @@ import respx import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage TIER_MODEL: Final = "tier-priced-test-model" @@ -31,6 +32,20 @@ PROMPT_TOKENS: Final = 1000 CACHED_TOKENS: Final = 200 COMPLETION_TOKENS: Final = 500 TIER_API_BASE: Final = "https://tier-pricing.invalid/v1" +AUDIO_TIER_MODEL: Final = "tier-priced-gemini-audio-test-model" +AUDIO_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( + { + "input_cost_per_token": 4e-06, + "output_cost_per_token": 8e-06, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_audio_token": 3e-06, + "cache_read_input_audio_token_cost": 2e-06, + "input_cost_per_audio_token_flex": 1e-06, + "input_cost_per_audio_token_priority": 2e-06, + "cache_read_input_audio_token_cost_flex": 4e-07, + "cache_read_input_audio_token_cost_priority": 8e-07, + } +) def _cost_at(prices: Mapping[str, float], column_suffix: str) -> float: @@ -78,6 +93,53 @@ def test_completion_cost_bills_the_price_columns_of_the_service_tier( assert cost == pytest.approx(_cost_at(TIER_ROW, column_suffix)) +@pytest.mark.parametrize( + ("service_tier", "audio_rate_key", "cache_audio_rate_key"), + ( + pytest.param("flex", "input_cost_per_audio_token_flex", "cache_read_input_audio_token_cost_flex"), + pytest.param("priority", "input_cost_per_audio_token_priority", "cache_read_input_audio_token_cost_priority"), + ), +) +def test_gemini_audio_and_cached_audio_use_service_tier_rates( + local_model_cost_map: None, service_tier: str, audio_rate_key: str, cache_audio_rate_key: str +) -> None: + litellm.register_model( + { + AUDIO_TIER_MODEL: { + "litellm_provider": "gemini", + "mode": "chat", + **AUDIO_TIER_ROW, + } + } + ) + usage: Final = Usage( + prompt_tokens=40, + completion_tokens=0, + total_tokens=40, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, + audio_tokens=20, + cached_tokens=10, + cached_tokens_details=CachedTokensDetails(audio_tokens=10), + ), + ) + response: Final = ModelResponse(model=AUDIO_TIER_MODEL, usage=usage) + + cost: Final = litellm.completion_cost( + completion_response=response, + model=AUDIO_TIER_MODEL, + custom_llm_provider="gemini", + service_tier=service_tier, + ) + + expected: Final = ( + 20 * AUDIO_TIER_ROW["input_cost_per_token"] + + 10 * AUDIO_TIER_ROW[audio_rate_key] + + 10 * AUDIO_TIER_ROW[cache_audio_rate_key] + ) + assert cost == pytest.approx(expected) + + LONG_CONTEXT_TIER_MODEL: Final = "long-context-tier-priced-test-model" LONG_CONTEXT_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( { diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py index 921dde8bf7b..bdc14c4934b 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( used_pricing_keys, zero_cost_warning, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage PER_SECOND_ENTRY: Final = {"cost_per_second": 0.00042} @@ -135,6 +136,36 @@ def test_audio_tokens_need_the_audio_rates() -> None: assert diagnostic["missing_pricing_keys"] == ("input_cost_per_audio_token", "output_cost_per_audio_token") +def test_cached_audio_uses_its_cache_rate_in_zero_cost_diagnostic() -> None: + usage: Final = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, + audio_tokens=20, + cached_tokens=20, + cached_tokens_details=CachedTokensDetails(audio_tokens=20), + ), + ) + pricing_entry: Final = { + "input_cost_per_token": 1e-6, + "input_cost_per_audio_token": 2e-6, + "cache_read_input_audio_token_cost": 0, + } + + assert used_pricing_keys(usage) == ("cache_read_input_audio_token_cost",) + assert ( + diagnose_zero_cost( + usage=usage, + pricing_model="gemini-audio", + pricing_entry=pricing_entry, + calculation_failed=False, + ) + is None + ) + + def test_warning_names_the_request_the_entry_the_missing_keys_and_the_counter() -> None: diagnostic = diagnose_zero_cost( usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 4c3ee9e7ca5..3bc864cc047 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -16,8 +16,10 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse, completion, embedding, image_generation from litellm._logging import verbose_logger +from litellm.cost_calculator import batch_cost_calculator from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token, parse_prompt_tokens_details from litellm.llms.anthropic.pass_through.messages import handler as anthropic_messages_handler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig @@ -30,8 +32,9 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( normalize_gemini_speech_config, reject_unknown_gemini_tts_voice_keys, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, Usage +from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger @@ -1515,6 +1518,7 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): v = VertexGeminiConfig() usage_metadata = { "promptTokenCount": 210, + "cachedContentTokenCount": 54, "candidatesTokenCount": 50, "totalTokenCount": 260, "promptTokensDetails": [ @@ -1536,9 +1540,10 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): usage_metadata = UsageMetadata(**usage_metadata) result = v.calculate_usage(completion_response={"usageMetadata": usage_metadata}) - # prompt details are total - cached per modality - assert result.prompt_tokens_details.text_tokens == 16 # 20 - 4 - assert result.prompt_tokens_details.image_tokens == 140 # (90 + 100) - (40 + 10) + assert result.prompt_tokens_details.text_tokens == 20 + assert result.prompt_tokens_details.image_tokens == 190 + assert result.prompt_tokens_details.cached_tokens == 54 + assert result.prompt_tokens_details.cached_tokens_details == CachedTokensDetails(text_tokens=4, image_tokens=50) # candidates details accumulate duplicate modalities assert result.completion_tokens_details.text_tokens == 20 # 15 + 5 @@ -1546,6 +1551,118 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): assert result.completion_tokens == 50 +def _gemini_usage_with_cached_audio(cached_audio_tokens: int) -> Usage: + cache_tokens_details: Final = ( + ( + {"modality": "TEXT", "tokenCount": 5}, + {"modality": "AUDIO", "tokenCount": cached_audio_tokens}, + ) + if cached_audio_tokens > 0 + else ({"modality": "TEXT", "tokenCount": 5},) + ) + completion_response: Final = { + "usageMetadata": { + "promptTokenCount": 40, + "cachedContentTokenCount": 5 + cached_audio_tokens, + "promptTokensDetails": ( + {"modality": "TEXT", "tokenCount": 20}, + {"modality": "AUDIO", "tokenCount": 20}, + ), + "cacheTokensDetails": cache_tokens_details, + "candidatesTokenCount": 7, + "totalTokenCount": 47, + } + } + return VertexGeminiConfig.calculate_usage(completion_response=completion_response) + + +def test_gemini_raw_usage_metadata_prices_cached_audio_at_its_batch_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_batches": 1e-6, + "input_cost_per_audio_token": 6e-6, + "input_cost_per_audio_token_batches": 3e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_token_cost_batches": 4e-7, + "cache_read_input_audio_token_cost": 5e-7, + "cache_read_input_audio_token_cost_batches": 2e-7, + } + usage: Final = _gemini_usage_with_cached_audio(5) + expected_prompt_cost: Final = ( + 15 * model_info["input_cost_per_token_batches"] + + 15 * model_info["input_cost_per_audio_token_batches"] + + 5 * model_info["cache_read_input_token_cost_batches"] + + 5 * model_info["cache_read_input_audio_token_cost_batches"] + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-audio-batch-test-model", + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_gemini_raw_usage_metadata_prices_cached_audio_at_its_standard_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_audio_token": 6e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_audio_token_cost": 5e-7, + } + usage: Final = _gemini_usage_with_cached_audio(5) + expected_prompt_cost: Final = ( + 15 * model_info["input_cost_per_token"] + + 15 * model_info["input_cost_per_audio_token"] + + 5 * model_info["cache_read_input_token_cost"] + + 5 * model_info["cache_read_input_audio_token_cost"] + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_gemini_usage_totals_are_subtracted_once_and_audio_cache_rate_matches_audio_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_audio_token": 6e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_audio_token_cost": 6e-6, + } + cached_usage: Final = _gemini_usage_with_cached_audio(5) + uncached_usage: Final = _gemini_usage_with_cached_audio(0) + cached_details: Final = parse_prompt_tokens_details(cached_usage) + cached_prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=cached_usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + uncached_prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=uncached_usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert cached_usage.prompt_tokens_details.text_tokens == 20 + assert cached_usage.prompt_tokens_details.audio_tokens == 20 + assert cached_usage.prompt_tokens_details.cached_tokens_details == CachedTokensDetails( + text_tokens=5, audio_tokens=5 + ) + assert cached_details["text_tokens"] == 15 + assert cached_details["audio_tokens"] == 15 + assert cached_prompt_cost == pytest.approx(uncached_prompt_cost) + + def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): """ If budget_tokens is 0, do not set includeThoughts to True @@ -6150,7 +6267,7 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): def test_vertex_ai_usage_metadata_with_document_tokens_cached(): - """Test that cached DOCUMENT tokens are correctly subtracted from prompt text tokens.""" + """Test cached DOCUMENT token details and unchanged prompt cost.""" v = VertexGeminiConfig() usage_metadata_dict = { @@ -6174,10 +6291,24 @@ def test_vertex_ai_usage_metadata_with_document_tokens_cached(): completion_response = {"usageMetadata": usage_metadata_dict} result = v.calculate_usage(completion_response=completion_response) - # DOCUMENT cached tokens map to cached_text_tokens, so: - # text_tokens = (8 TEXT + 774 DOCUMENT) - 400 cached = 382 - assert result.prompt_tokens_details.text_tokens == 382, "text_tokens should be (8 + 774) - 400 cached = 382" + assert result.prompt_tokens_details.text_tokens == 782 assert result.prompt_tokens_details.cached_tokens == 400 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 400 + + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "cache_read_input_token_cost": 5e-7, + } + prompt_cost, _ = generic_cost_per_token( + model="gemini-document-cache-test-model", + usage=result, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + 382 * model_info["input_cost_per_token"] + 400 * model_info["cache_read_input_token_cost"] + ) def test_async_streaming_uses_custom_client(): diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 5d4eb5fae12..550fe5f0d1a 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -16,6 +16,7 @@ import litellm from litellm import TranscriptionResponse, model_cost from litellm.constants import REALTIME_METERED_USAGE_KEY from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator +from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, @@ -3337,14 +3338,7 @@ def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models(_loca def test_gemini_cache_tokens_details_no_negative_values(): - """ - Test for Issue #18750: Negative text_tokens with Gemini caching - - When using Gemini with explicit caching, the response includes cacheTokensDetails - which breaks down cached tokens by modality. This test ensures that: - 1. text_tokens is never negative - 2. We correctly subtract cached tokens per modality (not total) - """ + """Ensure issue #18750 cache details remain non-negative after cost parsing.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -3372,27 +3366,17 @@ def test_gemini_cache_tokens_details_no_negative_values(): usage = VertexGeminiConfig.calculate_usage(completion_response) - # Text tokens should be non-cached text only: 9402 - 9393 = 9 - assert usage.prompt_tokens_details.text_tokens == 9, ( - f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" - ) + prompt_details: Final = parse_prompt_tokens_details(usage) - # Image tokens should be non-cached image only: 258 - 258 = 0 - assert usage.prompt_tokens_details.image_tokens == 0, ( - f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" + assert usage.prompt_tokens_details.text_tokens == 9402 + assert usage.prompt_tokens_details.image_tokens == 258 + assert usage.prompt_tokens_details.cached_tokens == 9651 + assert usage.prompt_tokens_details.cached_tokens_details == CachedTokensDetails( + text_tokens=9393, image_tokens=258 ) - - # Total cached should match - assert usage.prompt_tokens_details.cached_tokens == 9651, ( - f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" - ) - - # MOST IMPORTANT: text_tokens should NEVER be negative - assert usage.prompt_tokens_details.text_tokens >= 0, ( - f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - ) - - print("✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative") + assert prompt_details["text_tokens"] == 9 + assert prompt_details["image_tokens"] == 0 + assert prompt_details["text_tokens"] >= 0 def test_gemini_without_cache_tokens_details(): @@ -5990,6 +5974,45 @@ def test_batch_cost_calculator_prices_modalities_and_cached_tokens_together_in_t assert prompt_cost == pytest.approx(298_926 * 3e-6 + 64 * 5e-6 + 10 * 3e-6 + 1_000 * 3e-7) +def test_batch_cost_calculator_prices_cached_audio_at_its_batch_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_batches": 1e-6, + "input_cost_per_audio_token": 6e-6, + "input_cost_per_audio_token_batches": 3e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_token_cost_batches": 4e-7, + "cache_read_input_audio_token_cost": 5e-7, + "cache_read_input_audio_token_cost_batches": 2e-7, + } + usage: Final = Usage( + prompt_tokens=40, + completion_tokens=0, + total_tokens=40, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, + audio_tokens=20, + cached_tokens=10, + cached_tokens_details=CachedTokensDetails(text_tokens=5, audio_tokens=5), + ), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-audio-batch-test-model", + custom_llm_provider="gemini", + model_info=model_info, + ) + + expected: Final = ( + 15 * model_info["input_cost_per_token_batches"] + + 15 * model_info["input_cost_per_audio_token_batches"] + + 5 * model_info["cache_read_input_token_cost_batches"] + + 5 * model_info["cache_read_input_audio_token_cost_batches"] + ) + assert prompt_cost == pytest.approx(expected) + + QWEN3_NEXT_REGIONS: Final = ("ap-northeast-1", "ap-south-1", "ap-southeast-2", "eu-west-1", "eu-west-2", "sa-east-1") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index fcd6b8aa402..338faae188f 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -994,6 +994,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "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"}, @@ -1001,6 +1004,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "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"}, diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index bfaac44a483..4bd16a2997f 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -61,6 +61,30 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): assert model_info["input_cost_per_token"] == 0.0 +def test_get_model_info_passes_through_gemini_tiered_audio_rates() -> None: + model: Final = "gemini-tiered-audio-rates-test" + prices: Final = { + "input_cost_per_audio_token_flex": 1e-6, + "input_cost_per_audio_token_priority": 2e-6, + "cache_read_input_audio_token_cost_flex": 3e-7, + "cache_read_input_audio_token_cost_priority": 4e-7, + "cache_read_input_audio_token_cost_batches": 5e-7, + } + litellm.register_model( + { + model: { + "litellm_provider": "gemini", + "mode": "chat", + **prices, + } + } + ) + + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + assert {key: model_info[key] for key in prices} == prices + + def test_get_model_info_ollama_chat(): from litellm.llms.ollama.completion.transformation import OllamaConfig diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index da9e77e9b58..b5a86dd2187 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35259,6 +35259,12 @@ export interface components { cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; + /** Cache Read Input Audio Token Cost Batches */ + cache_read_input_audio_token_cost_batches?: number | null; + /** Cache Read Input Audio Token Cost Flex */ + cache_read_input_audio_token_cost_flex?: number | null; + /** Cache Read Input Audio Token Cost Priority */ + cache_read_input_audio_token_cost_priority?: number | null; /** Cache Read Input Image Token Cost */ cache_read_input_image_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -35338,6 +35344,10 @@ export interface components { input_cost_per_audio_token?: number | null; /** Input Cost Per Audio Token Batches */ input_cost_per_audio_token_batches?: number | null; + /** Input Cost Per Audio Token Flex */ + input_cost_per_audio_token_flex?: number | null; + /** Input Cost Per Audio Token Priority */ + input_cost_per_audio_token_priority?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -50568,6 +50578,12 @@ export interface components { cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; + /** Cache Read Input Audio Token Cost Batches */ + cache_read_input_audio_token_cost_batches?: number | null; + /** Cache Read Input Audio Token Cost Flex */ + cache_read_input_audio_token_cost_flex?: number | null; + /** Cache Read Input Audio Token Cost Priority */ + cache_read_input_audio_token_cost_priority?: number | null; /** Cache Read Input Image Token Cost */ cache_read_input_image_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -50647,6 +50663,10 @@ export interface components { input_cost_per_audio_token?: number | null; /** Input Cost Per Audio Token Batches */ input_cost_per_audio_token_batches?: number | null; + /** Input Cost Per Audio Token Flex */ + input_cost_per_audio_token_flex?: number | null; + /** Input Cost Per Audio Token Priority */ + input_cost_per_audio_token_priority?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ From 082064569082761bbb8ce1dc1b1be2fb2a145b94 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 24 Oct 2026 10:17:14 -0700 Subject: [PATCH 03/43] fix(mcp): retain toolset scope when refreshing catalog authority (#48095) --- .../mcp_server/auth/user_api_key_auth_mcp.py | 9 +- .../proxy/_experimental/mcp_server/server.py | 2 +- .../auth/test_user_api_key_auth_mcp.py | 19 +++- .../_experimental/mcp_server/test_catalog.py | 101 +++++++++++++++++- 4 files changed, 125 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 53b56c7b31d..7f0d392ab1b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1368,7 +1368,7 @@ class MCPRequestHandler: if auth is None: return None - refreshed = auth.model_copy() + refreshed: Final = auth.model_copy(update={"requires_fresh_policy": True}) if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id: # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is # sealed in the token: refresh the live user policy and keep the token's team, routes and models. @@ -1386,7 +1386,12 @@ class MCPRequestHandler: refreshed.project_id = current.project_id refreshed.user_id = current.user_id refreshed.access_group_ids = current.access_group_ids - refreshed.requires_fresh_policy = True + else: + return refreshed + if refreshed.mcp_toolset_id is not None: + from litellm.proxy._experimental.mcp_server.server import apply_toolset_scope + + return await apply_toolset_scope(refreshed, refreshed.mcp_toolset_id) return refreshed @staticmethod diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b3426089375..bb7c81f9d60 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1797,7 +1797,7 @@ if MCP_AVAILABLE: return acting.model_copy(update={"mcp_toolset_id": toolset_id}) tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=[toolset_id] + toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy ) server_ids: Final = list(tool_permissions.keys()) existing_op: Final = acting.object_permission diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index f73a8547158..e8689b2ba89 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -10906,7 +10906,19 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc @pytest.mark.asyncio @pytest.mark.parametrize("current_groups", [[], ["replacement-group"]]) async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session_scope(monkeypatch, current_groups): - permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_servers=["current-server"]) + from litellm.proxy import proxy_server + from litellm.types.mcp_server.mcp_toolset import MCPToolset + + permission = LiteLLM_ObjectPermissionTable( + object_permission_id="current-policy", mcp_toolsets=["session-toolset"] + ) + toolset: Final = MCPToolset( + toolset_id="session-toolset", toolset_name="session-toolset", + tools=[{"server_id": "current-server", "tool_name": "read"}], + ) + tables: Final = SimpleNamespace(litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset]))) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) current = UserAPIKeyAuth( object_permission=permission, object_permission_id="current-policy", @@ -10931,7 +10943,10 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session caller.mcp_toolset_id = "session-toolset" refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True) - assert refreshed.object_permission == permission + assert refreshed.object_permission == permission.model_copy(update={ + "mcp_servers": ["current-server"], "mcp_toolsets": [], + "mcp_tool_permissions": {"current-server": ["read"]}, + }) assert refreshed.object_permission_id == "current-policy" assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ( "new-team", diff --git a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py index b383078c0ed..4f6530e5de1 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock @@ -13,6 +13,7 @@ from mcp.types import ( ListResourcesResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, + ListToolsRequest, ListToolsResult, PaginatedRequestParams, Tool, @@ -34,6 +35,104 @@ CatalogKind = Literal["tools", "prompts", "resources", "templates"] OptionalCatalogResult = ListPromptsResult | ListResourcesResult | ListResourceTemplatesResult +@pytest.mark.asyncio +@pytest.mark.parametrize( + "selected,revoked,expected", + [ + ("narrow", False, ("alpha-read_item",)), + ("full", False, ("alpha-read_item", "alpha-write_item", "beta-read_item", "beta-write_item")), + (None, False, ("alpha-read_item", "alpha-write_item", "beta-read_item", "beta-write_item")), + ("narrow", True, ()), + ], +) +async def test_catalog_refresh_keeps_selected_toolset_with_current_grants_and_contents( + monkeypatch: pytest.MonkeyPatch, + selected: str | None, + revoked: bool, + expected: tuple[str, ...], +) -> None: + from fastapi import HTTPException + + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.server import apply_toolset_scope + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.types.mcp_server.mcp_toolset import MCPToolset, MCPToolsetTool + + grants: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="catalog-grants", mcp_toolsets=["narrow", "full"] + ) + live_grants: Final = grants.model_copy(update={"mcp_toolsets": ["full"]}) if revoked else grants + tools: Final[tuple[MCPToolsetTool, ...]] = ( + {"server_id": "alpha", "tool_name": "read_item"}, + {"server_id": "alpha", "tool_name": "write_item"}, + {"server_id": "beta", "tool_name": "read_item"}, + {"server_id": "beta", "tool_name": "write_item"}, + ) + toolsets: Final = ( + MCPToolset(toolset_id="narrow", toolset_name="narrow", tools=[tools[0]]), + MCPToolset(toolset_id="full", toolset_name="full", tools=list(tools)), + ) + + async def read_toolsets(where: Mapping[str, Mapping[str, Sequence[str]]]) -> tuple[MCPToolset, ...]: + return tuple(toolset for toolset in toolsets if toolset.toolset_id in where["toolset_id"]["in"]) + + tables: Final = SimpleNamespace( + litellm_objectpermissiontable=SimpleNamespace(find_unique=AsyncMock(return_value=live_grants)), + litellm_mcptoolsettable=SimpleNamespace(find_many=read_toolsets), + ) + database: Final = SimpleNamespace( + db=tables, + writer_db=tables, + get_data=AsyncMock(return_value=UserAPIKeyAuth(object_permission_id=grants.object_permission_id)), + ) + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) + await cache.async_set_cache(key="toolset_perms:narrow", value={"alpha": ["read_item", "write_item"]}) + manager: Final = operations.global_mcp_server_manager + for server_id in ("alpha", "beta"): + monkeypatch.setitem( + manager.registry, + server_id, + MCPServer(server_id=server_id, name=server_id, transport=MCPTransport.http, spec_path="/catalog.yaml"), + ) + + async def tool_handler() -> str: + return "ok" + + for tool in tools: + global_mcp_tool_registry.register_tool( + name=f"{tool['server_id']}-{tool['tool_name']}", + description=tool["tool_name"], + input_schema={"type": "object"}, + handler=tool_handler, + ) + original: Final = UserAPIKeyAuth( + api_key="catalog-key-hash", object_permission=grants, + object_permission_id=grants.object_permission_id, + ) + original.via_virtual_key = True + caller: Final = await apply_toolset_scope(original, selected) if selected is not None else original + before: Final = caller.model_dump() + context: Final = operations.prepare_context(user_api_key_auth=caller) + if revoked: + with pytest.raises(HTTPException, match="does not have access to toolset") as denied: + await catalog.list_gateway_catalog(context, ListToolsRequest(), log_list_tools_to_spendlogs=False) + assert denied.value.status_code == 403 + else: + result: Final = await catalog.list_gateway_catalog( + context, ListToolsRequest(), log_list_tools_to_spendlogs=False + ) + assert isinstance(result, ListToolsResult) + assert tuple(sorted(tool.name for tool in result.tools)) == expected + assert caller.model_dump() == before + assert original.object_permission == grants + + def page(name: str, cursor: str | None = None, revision: str = "stable") -> ListToolsResult: return ListToolsResult( tools=[Tool(name=name, input_schema={"type": "object"})], From 464932b1d19d7ed8d377ed79fc2a9c0110dff39b Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:31:09 -0700 Subject: [PATCH 04/43] refactor(typing): clear 1181 basedpyright Any errors across 431 files (#47789) * refactor(typing): narrow Any sources across providers, proxy and guardrails (wave 1) * refactor(typing): narrow Any sources across hubs, providers and proxy (wave 2) * refactor(typing): validate RunwayML task JSON and let dotprompt loaders accept exported prompts * fix: restore pass-through behavior on paths the typing pass tightened MCP client IP extraction trusts X-Forwarded-For only when use_x_forwarded_for is True, dotprompt skips a malformed JSON prompt entry instead of raising, RunwayML task payloads pass through unvalidated, and a malformed provider body raises the same AttributeError as before at the CompactifAI, Gemini, and mock transport parse sites * fix(runwayml): keep a malformed 2xx task body on the 500 path * fix(edenai,cli): keep a non-object body in the error and fail on a malformed migration report EdenAI's embedding and image generation parse dropped a non-object body from the error they raised, so the message no longer named what the provider sent. A helper now raises with the body in the message, and both paths have a regression test. `lite encryption migrate` exited 0 with no message when the proxy's report was not an object or its residual_legacy was not a number, where base crashed. It now fails with a message naming the value. The CompactifAI and Gemini image-edit non-object body tests accept TypeError or AttributeError, which both base and this branch raise, so they guard against a ValidationError instead of pinning one class. RunwayML's async parse gets the same non-object body test as the sync one. * fix: keep base's error classes, duck-typed checks and fail-closed guardrails the typing pass changed * test: annotate the new regression tests' locals with Final * test: annotate the remaining new test locals with Final --- enterprise/enterprise_hooks/aporia_ai.py | 13 +- .../enterprise_hooks/blocked_user_list.py | 17 +-- .../google_text_moderation.py | 13 +- .../example_logging_api.py | 7 +- .../enterprise_callbacks/llm_guard.py | 18 ++- .../proxy/common_utils/check_batch_cost.py | 4 +- .../proxy/hooks/managed_files.py | 27 ++-- .../proxy/hooks/managed_vector_stores.py | 6 +- .../management_endpoints/project_endpoints.py | 13 +- .../proxy/vector_stores/endpoints.py | 8 +- litellm/__init__.py | 16 ++- litellm/_lazy_imports.py | 4 +- litellm/a2a_protocol/cost_calculator.py | 20 ++- .../bedrock_agentcore/transformation.py | 27 ++-- .../watsonx_orchestrate/transformation.py | 40 ++++-- litellm/a2a_protocol/utils.py | 21 ++-- litellm/anthropic_beta_headers_manager.py | 13 +- .../exceptions/exception_mapping_utils.py | 11 +- litellm/assistants/main.py | 24 ++-- litellm/caching/_embedding_router.py | 6 +- litellm/caching/azure_blob_cache.py | 8 +- litellm/caching/base_cache.py | 2 +- litellm/caching/disk_cache.py | 8 +- litellm/cost_calculator.py | 44 ++++++- litellm/exceptions.py | 2 +- litellm/experimental_mcp_client/tools.py | 20 +-- litellm/files/main.py | 14 ++- litellm/files/streaming.py | 4 +- litellm/harness/handlers/base.py | 3 +- litellm/harness/handlers/cli_handler.py | 4 +- .../harness/handlers/deepagents_handler.py | 2 +- litellm/harness/runtime.py | 11 +- litellm/images/main.py | 11 +- .../SlackAlerting/hanging_request_check.py | 14 ++- .../SlackAlerting/user_spend_alerts.py | 5 +- litellm/integrations/agentops/agentops.py | 23 ++-- litellm/integrations/arize/arize.py | 2 +- litellm/integrations/braintrust_logging.py | 29 +++-- .../integrations/cloudzero/cz_stream_api.py | 7 +- litellm/integrations/custom_guardrail.py | 7 +- litellm/integrations/custom_secret_manager.py | 3 +- litellm/integrations/dotprompt/__init__.py | 4 +- .../dotprompt/dotprompt_manager.py | 15 ++- .../integrations/dotprompt/prompt_manager.py | 29 +++-- .../focus/destinations/factory.py | 5 +- litellm/integrations/focus/export_engine.py | 8 +- litellm/integrations/focus/focus_logger.py | 8 +- litellm/integrations/gcs_bucket/gcs_bucket.py | 8 +- .../gcs_bucket/gcs_bucket_mock_client.py | 3 +- litellm/integrations/gcs_pubsub/pub_sub.py | 8 +- litellm/integrations/langfuse/langfuse.py | 5 +- .../integrations/langfuse/langfuse_otel.py | 5 +- .../langfuse/langfuse_prompt_management.py | 12 +- litellm/integrations/langsmith.py | 11 +- litellm/integrations/langtrace.py | 2 +- litellm/integrations/levo/levo.py | 2 +- litellm/integrations/logfire_logger.py | 5 +- litellm/integrations/newrelic/newrelic.py | 19 +-- .../integrations/newrelic/newrelic_metrics.py | 2 +- .../opik_payload_builder/payload_builders.py | 4 +- litellm/integrations/otel/model/config.py | 3 +- litellm/integrations/otel/model/metadata.py | 6 +- litellm/integrations/otel/model/payloads.py | 15 ++- litellm/integrations/otel/plumbing/metrics.py | 8 +- litellm/integrations/sqs.py | 5 +- .../handler.py | 5 +- litellm/litellm_core_utils/app_crypto.py | 9 +- .../cloud_storage_security.py | 6 +- litellm/litellm_core_utils/core_helpers.py | 47 +++++-- litellm/litellm_core_utils/cost_discount.py | 2 +- .../litellm_core_utils/default_encoding.py | 2 +- .../dot_notation_indexing.py | 2 +- .../exception_mapping_utils.py | 4 +- .../litellm_core_utils/get_model_cost_map.py | 6 +- litellm/litellm_core_utils/hidden_params.py | 3 +- .../initialize_dynamic_callback_params.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 63 ++++++---- .../llm_cost_calc/tool_call_cost_tracking.py | 2 +- litellm/litellm_core_utils/llm_judge.py | 5 +- .../logging_callback_manager.py | 8 +- .../prompt_templates/common_utils.py | 20 +-- .../prompt_templates/factory.py | 32 +++-- .../mid_conversation_system.py | 2 +- .../sensitive_data_masker.py | 8 +- .../streaming_chunk_builder_utils.py | 3 +- .../litellm_core_utils/streaming_handler.py | 86 ++++++++----- litellm/litellm_core_utils/token_counter.py | 4 +- litellm/litellm_core_utils/url_utils.py | 59 +++++---- litellm/llms/__init__.py | 3 +- litellm/llms/a2a/common_utils.py | 6 +- litellm/llms/ai21/chat/transformation.py | 4 +- .../llms/amazon_nova/chat/transformation.py | 4 +- litellm/llms/anthropic/batches/handler.py | 2 +- .../chat/guardrail_translation/handler.py | 9 +- litellm/llms/anthropic/chat/transformation.py | 7 +- .../anthropic/completion/transformation.py | 3 +- litellm/llms/anthropic/files/handler.py | 5 +- .../adapters/streaming_iterator.py | 8 +- .../editors/clear_tool_uses.py | 15 ++- .../context_management/placeholders.py | 4 +- .../pass_through/context_management/result.py | 6 +- .../pass_through/convertible_blocks.py | 17 ++- .../messages/fake_stream_iterator.py | 12 +- .../pass_through/messages/transformation.py | 2 +- .../responses_adapters/transformation.py | 4 +- litellm/llms/azure/audio_transcriptions.py | 4 +- litellm/llms/azure/azure.py | 10 +- litellm/llms/azure/exception_mapping.py | 8 +- .../llms/azure/responses/transformation.py | 2 +- litellm/llms/azure_ai/chat/transformation.py | 2 +- litellm/llms/base_llm/chat/transformation.py | 5 +- .../guardrail_translation/base_translation.py | 2 +- .../base_llm/guardrail_translation/utils.py | 4 +- .../base_llm/image_edit/transformation.py | 5 +- .../base_llm/interactions/transformation.py | 6 +- .../base_managed_resource.py | 4 +- .../base_llm/passthrough/transformation.py | 3 +- .../llms/base_llm/rerank/transformation.py | 2 +- .../base_llm/text_to_speech/transformation.py | 6 +- .../base_llm/vector_store/transformation.py | 4 +- litellm/llms/baseten/chat.py | 3 +- litellm/llms/bedrock/batches/handler.py | 3 +- .../bedrock/chat/agentcore/transformation.py | 15 ++- .../bedrock/chat/converse_transformation.py | 26 +++- .../amazon_ai21_transformation.py | 3 +- .../amazon_cohere_transformation.py | 3 +- .../amazon_llama_transformation.py | 3 +- .../amazon_mistral_transformation.py | 3 +- .../amazon_titan_transformation.py | 3 +- .../anthropic_claude2_transformation.py | 3 +- .../bedrock/chat/mantle/transformation.py | 5 +- litellm/llms/bedrock/common_utils.py | 7 +- .../embed/amazon_titan_g1_transformation.py | 3 +- .../embed/amazon_titan_v2_transformation.py | 3 +- litellm/llms/bedrock/files/handler.py | 4 +- litellm/llms/bedrock/image_edit/handler.py | 5 +- .../amazon_stability3_transformation.py | 4 +- .../amazon_titan_transformation.py | 6 +- .../anthropic_claude3_transformation.py | 4 +- .../bedrock/messages/mantle_transformation.py | 4 +- litellm/llms/bedrock/model_listing.py | 5 +- .../guardrail_translation/handler.py | 6 +- .../image_edit/transformation.py | 3 +- .../image_generation/transformation.py | 3 +- litellm/llms/cerebras/chat.py | 3 +- litellm/llms/chatgpt/chat/streaming_utils.py | 7 +- litellm/llms/clarifai/chat/transformation.py | 3 +- .../claude_code/harness/transformation.py | 5 +- litellm/llms/codestral/completion/handler.py | 4 +- litellm/llms/codex/harness/transformation.py | 7 +- litellm/llms/cohere/chat/transformation.py | 5 +- litellm/llms/cohere/rerank/transformation.py | 3 +- .../llms/cohere/rerank_v2/transformation.py | 4 +- litellm/llms/cometapi/chat/transformation.py | 4 +- .../llms/compactifai/chat/transformation.py | 44 +++++-- .../llms/custom_httpx/async_client_cleanup.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 50 +++++--- litellm/llms/custom_httpx/mock_transport.py | 13 +- litellm/llms/dashscope/cost_calculator.py | 13 +- .../image_generation/transformation.py | 3 +- .../llms/databricks/embed/transformation.py | 6 +- .../databricks/messages/transformation.py | 4 +- .../deepagents/harness/sandbox_backend.py | 4 +- .../llms/deepagents/harness/transformation.py | 7 +- .../audio_transcription/transformation.py | 41 +++++- .../llms/deepinfra/rerank/transformation.py | 5 +- .../llms/deepseek/messages/transformation.py | 36 ++++-- litellm/llms/deprecated_providers/palm.py | 6 +- .../llms/edenai/embedding/transformation.py | 18 ++- .../edenai/image_generation/transformation.py | 18 ++- .../audio_transcription/transformation.py | 35 +++--- .../flux_pro_v11_ultra_transformation.py | 3 +- .../recraft_v3_transformation.py | 3 +- .../stable_diffusion_transformation.py | 3 +- .../featherless_ai/chat/transformation.py | 4 +- .../fireworks_ai/rerank/transformation.py | 3 +- litellm/llms/gdc/chat/transformation.py | 8 +- litellm/llms/gemini/agents/transformation.py | 52 +++++--- .../audio_transcription/transformation.py | 4 +- litellm/llms/gemini/common_utils.py | 12 +- .../guardrail_translation/handler.py | 10 +- .../gemini/google_genai/transformation.py | 16 ++- .../llms/gemini/image_edit/transformation.py | 28 +++-- .../gemini/image_generation/transformation.py | 7 +- .../llms/gemini/realtime/transformation.py | 12 +- litellm/llms/gemini/videos/transformation.py | 5 +- litellm/llms/gigachat/chat/transformation.py | 6 +- litellm/llms/gigachat/file_handler.py | 18 +-- .../gigachat/passthrough/transformation.py | 8 +- .../github_copilot/chat/transformation.py | 11 +- .../github_copilot/messages/transformation.py | 5 +- .../responses/transformation.py | 28 +++-- .../llms/gradient_ai/chat/transformation.py | 4 +- litellm/llms/groq/chat/transformation.py | 3 +- litellm/llms/groq/cost_calculator.py | 5 +- litellm/llms/groq/stt/transformation.py | 6 +- .../llms/hosted_vllm/rerank/transformation.py | 9 +- .../llms/hosted_vllm/videos/transformation.py | 4 +- litellm/llms/huggingface/embedding/handler.py | 18 +-- .../huggingface/embedding/transformation.py | 7 +- .../llms/huggingface/rerank/transformation.py | 2 +- litellm/llms/langflow/chat/transformation.py | 11 +- litellm/llms/langgraph/chat/sse_iterator.py | 5 +- litellm/llms/langgraph/chat/transformation.py | 10 +- litellm/llms/lemonade/chat/transformation.py | 7 +- .../litellm_proxy/skills/prompt_injection.py | 34 ++--- .../litellm_proxy/skills/sandbox_executor.py | 12 +- .../llms/lm_studio/embed/transformation.py | 6 +- litellm/llms/maritalk.py | 3 +- .../milvus/vector_stores/transformation.py | 14 ++- .../llms/minimax/messages/transformation.py | 4 +- .../image_generation/transformation.py | 3 +- litellm/llms/nadir/chat/transformation.py | 9 +- litellm/llms/nvidia_nim/embed.py | 6 +- .../audio_transcription/handler.py | 31 +++-- .../audio_transcription/transformation.py | 17 ++- litellm/llms/oci/chat/transformation.py | 3 +- litellm/llms/oci/embed/transformation.py | 3 +- litellm/llms/ollama/common_utils.py | 3 +- .../llms/ollama/completion/transformation.py | 2 +- litellm/llms/openai/common_utils.py | 17 ++- litellm/llms/openai/evals/transformation.py | 8 +- .../responses/count_tokens/token_counter.py | 4 +- .../responses/count_tokens/transformation.py | 42 +++++-- .../guardrail_translation/handler.py | 11 +- .../openai_like/messages/transformation.py | 6 +- .../llms/opencode/harness/transformation.py | 26 +++- .../openrouter/image_edit/transformation.py | 6 +- .../image_generation/transformation.py | 3 +- .../opensandbox/sandbox/transformation.py | 12 +- .../guardrail_translation/handler.py | 48 +++++-- .../perplexity/responses/transformation.py | 5 +- .../llms/petals/completion/transformation.py | 4 +- .../pg_vector/vector_stores/transformation.py | 2 +- .../image_generation/transformation.py | 3 +- .../image_generation/transformation.py | 40 ++++-- .../runwayml/text_to_speech/transformation.py | 17 ++- .../llms/runwayml/videos/transformation.py | 2 +- litellm/llms/sagemaker/common_utils.py | 11 +- .../embedding/cohere_transformation.py | 6 +- .../embedding/openai_transformation.py | 5 +- .../sagemaker/embedding/transformation.py | 6 +- litellm/llms/sap/chat/transformation.py | 8 +- litellm/llms/sap/credentials.py | 33 +++-- litellm/llms/sap/embed/transformation.py | 8 +- .../stability/image_edit/transformations.py | 11 +- .../llms/tencent/messages/transformation.py | 4 +- .../llms/triton/completion/transformation.py | 7 +- .../llms/triton/embedding/transformation.py | 16 ++- .../vercel_ai_gateway/chat/transformation.py | 14 ++- .../gemini_transcribe_transformation.py | 4 +- .../llms/vertex_ai/batches/transformation.py | 2 +- litellm/llms/vertex_ai/files/handler.py | 4 +- .../llms/vertex_ai/files/transformation.py | 15 ++- litellm/llms/vertex_ai/fine_tuning/handler.py | 13 +- litellm/llms/vertex_ai/gemini/long_decode.py | 3 +- .../batch_embed_content_handler.py | 10 +- .../vertex_gemini_transformation.py | 2 +- .../vertex_imagen_transformation.py | 7 +- .../image_generation_handler.py | 4 +- .../vertex_gemini_transformation.py | 3 +- .../multimodal_embeddings/transformation.py | 2 +- .../vertex_ai/rag_engine/transformation.py | 18 ++- .../llms/vertex_ai/rerank/transformation.py | 16 ++- .../ai21/transformation.py | 6 +- .../transformation.py | 5 +- .../vertex_embeddings/transformation.py | 6 +- .../llms/volcengine/chat/transformation.py | 3 +- .../volcengine/responses/transformation.py | 11 +- litellm/llms/voyage/rerank/transformation.py | 3 +- .../audio_transcription/transformation.py | 3 +- litellm/llms/watsonx/rerank/transformation.py | 3 +- litellm/llms/xai/chat/transformation.py | 13 +- litellm/llms/xai/cost_calculator.py | 9 +- litellm/main.py | 21 ++-- litellm/models/organization_membership.py | 10 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 3 +- .../proxy/_experimental/mcp_server/catalog.py | 7 +- litellm/proxy/_experimental/mcp_server/db.py | 18 +-- .../mcp_server/discoverable_endpoints.py | 2 +- .../mcp_server/faults/classify.py | 5 +- .../mcp_server/faults/list_outcomes.py | 2 +- .../mcp_server/mcp_server_manager.py | 8 +- .../mcp_server/oauth2_token_cache.py | 5 +- .../mcp_server/oauth_identity_binding.py | 6 +- .../token_exchange_provider.py | 4 +- .../mcp_server/result_conversion.py | 13 +- .../mcp_server/semantic_tool_filter.py | 14 ++- .../proxy/_experimental/mcp_server/server.py | 20 ++- .../_experimental/mcp_server/toolset_db.py | 6 +- .../proxy/_experimental/mcp_server/utils.py | 5 +- litellm/proxy/_lazy_features.py | 20 +-- litellm/proxy/_lazy_openapi_snapshot.py | 45 ++++--- litellm/proxy/a2a/discovery.py | 14 ++- .../proxy/agent_endpoints/a2a_endpoints.py | 15 ++- litellm/proxy/agent_endpoints/a2a_routing.py | 6 +- .../proxy/agent_endpoints/databricks_oauth.py | 3 +- .../proxy/agent_endpoints/managed_identity.py | 6 +- .../claude_code_marketplace.py | 3 +- litellm/proxy/auth/auth_checks.py | 4 +- litellm/proxy/auth/auth_utils.py | 16 ++- litellm/proxy/auth/ip_address_utils.py | 4 +- litellm/proxy/auth/rds_iam_token.py | 8 +- litellm/proxy/auth/route_checks.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 8 +- litellm/proxy/caching_routes.py | 7 +- litellm/proxy/client/chat.py | 5 +- .../proxy/client/cli/commands/encryption.py | 30 +++-- .../client/cli/commands/statusline_script.py | 26 +++- litellm/proxy/client/cli/interface.py | 2 +- litellm/proxy/client/cli/main.py | 6 +- litellm/proxy/client/credentials.py | 6 +- litellm/proxy/client/http_client.py | 7 +- litellm/proxy/client/model_groups.py | 4 +- litellm/proxy/client/models.py | 2 +- .../auth_cache_invalidation_pubsub.py | 15 ++- litellm/proxy/common_utils/callback_utils.py | 18 ++- litellm/proxy/common_utils/debug_utils.py | 7 +- .../proxy/common_utils/http_parsing_utils.py | 2 +- .../proxy/common_utils/load_config_utils.py | 3 +- .../common_utils/proxy_rate_limit_error.py | 4 +- .../container_endpoints/handler_factory.py | 5 +- .../proxy/container_endpoints/ownership.py | 25 ++-- litellm/proxy/db/autorouter_session_rollup.py | 2 +- litellm/proxy/db/db_spend_update_writer.py | 69 +++++++++-- .../db_transaction_queue/pod_lock_manager.py | 6 +- .../redis_update_buffer.py | 4 +- litellm/proxy/db/prisma_query_span.py | 5 +- litellm/proxy/db/proxy_worker_heartbeat.py | 5 +- litellm/proxy/db/routing_prisma_wrapper.py | 2 +- litellm/proxy/db/spend_log_tool_index.py | 4 +- .../agent_skills_archive.py | 4 +- .../agent_skills_endpoints.py | 7 +- .../proxy/example_config_yaml/custom_auth.py | 2 +- .../example_config_yaml/custom_callbacks.py | 12 +- litellm/proxy/guardrails/_content_utils.py | 15 +-- .../proxy/guardrails/guardrail_endpoints.py | 3 +- .../guardrail_hooks/agent_365/agent_365.py | 5 +- .../guardrails/guardrail_hooks/airia/airia.py | 13 +- .../block_code_execution/__init__.py | 3 +- .../cato_networks/cato_networks.py | 7 +- .../cisco_ai_defense/cisco_ai_defense.py | 6 +- .../guardrail_hooks/compresr/__init__.py | 3 +- .../guardrail_hooks/compresr/compresr.py | 4 +- .../crowdstrike_aidr/crowdstrike_aidr.py | 6 +- .../guardrail_hooks/custom_code/__init__.py | 2 +- .../generic_guardrail_api.py | 12 +- .../guardrail_hooks/grayswan/grayswan.py | 5 +- .../guardrail_hooks/javelin/javelin.py | 7 +- .../guardrail_hooks/lakera_ai_v2.py | 10 +- .../guardrails/guardrail_hooks/lasso/lasso.py | 43 ++++--- .../competitor_intent/airline.py | 8 +- .../litellm_content_filter/content_filter.py | 8 +- .../llm_shield_proxy/llm_shield_proxy.py | 53 ++++---- .../llm_shield_proxy/stream_restorers.py | 18 ++- .../mcp_end_user_permission/__init__.py | 15 ++- .../mcp_jwt_signer/mcp_jwt_signer.py | 19 +-- .../mcp_security/mcp_security_guardrail.py | 7 +- .../microsoft_purview/__init__.py | 5 +- .../guardrail_hooks/noma/noma_v2.py | 24 ++-- .../guardrails/guardrail_hooks/onyx/onyx.py | 5 +- .../guardrail_hooks/openai/__init__.py | 7 +- .../guardrail_hooks/openai/moderations.py | 7 +- .../panw_prisma_airs/panw_prisma_airs.py | 45 ++++--- .../guardrails/guardrail_hooks/presidio.py | 6 +- .../guardrail_hooks/qualifire/qualifire.py | 4 +- .../semantic_guard/route_loader.py | 4 +- .../semantic_guard/semantic_guard.py | 27 ++-- .../tool_policy/tool_policy_guardrail.py | 2 +- .../unified_guardrail/unified_guardrail.py | 20 +-- .../proxy/guardrails/guardrail_registry.py | 18 +-- .../proxy/guardrails/tool_name_extraction.py | 8 +- litellm/proxy/guardrails/usage_endpoints.py | 50 +++++--- .../proxy/hooks/key_management_event_hooks.py | 2 +- .../proxy/hooks/parallel_request_limiter.py | 2 +- .../hooks/parallel_request_limiter_v3.py | 2 +- litellm/proxy/hooks/rate_limiter_utils.py | 5 +- .../access_group_endpoints.py | 29 +++-- .../auto_router_endpoints.py | 31 +++-- .../credential_migration.py | 3 +- .../key_management_endpoints.py | 46 ++++--- .../management_v1/spend_logs.py | 22 +++- .../mcp_management_endpoints.py | 12 +- .../model_management_endpoints.py | 11 +- .../router_settings_endpoints.py | 5 +- .../management_endpoints/team_endpoints.py | 12 +- litellm/proxy/management_endpoints/ui_sso.py | 30 +++-- .../object_permission_utils.py | 2 +- litellm/proxy/memory/memory_endpoints.py | 4 +- litellm/proxy/moyai_endpoints.py | 23 ++-- .../batch_file_validation.py | 7 +- .../storage_backend_service.py | 12 +- .../llm_passthrough_endpoints.py | 2 +- ...end_medical_passthrough_logging_handler.py | 8 +- .../tinyfish_passthrough_logging_handler.py | 5 +- .../transcribe_passthrough_logging_handler.py | 9 +- .../upstream_usage_headers.py | 3 +- .../policy_engine/attachment_registry.py | 26 ++-- .../proxy/policy_engine/pipeline_executor.py | 14 ++- .../proxy/policy_engine/policy_endpoints.py | 2 +- litellm/proxy/proxy_server.py | 116 ++++++++++++----- .../public_endpoints/public_v1/model_hub.py | 6 +- .../spend_tracking/budget_reservation.py | 4 +- .../spend_tracking/ptu_flat_cost_rollup.py | 7 +- .../spend_management_endpoints.py | 4 +- litellm/proxy/utils.py | 117 ++++++++++++------ litellm/proxy/video_endpoints/endpoints.py | 6 +- litellm/proxy/video_endpoints/utils.py | 2 +- litellm/rag/ingestion/bedrock_ingestion.py | 4 +- litellm/realtime_api/main.py | 14 ++- litellm/repositories/base_repository.py | 3 +- litellm/repositories/team_repository.py | 3 +- litellm/rerank_api/main.py | 4 +- litellm/rerank_api/rerank_utils.py | 4 +- .../responses/file_search/emulated_handler.py | 4 +- .../custom_tools.py | 6 +- .../streaming_iterator.py | 11 +- litellm/responses/main.py | 9 +- .../responses/mcp/chat_completions_handler.py | 5 +- litellm/responses/sse_output_recovery.py | 5 +- litellm/responses/streaming_iterator.py | 5 +- litellm/responses/utils.py | 104 +++++++++------- litellm/router.py | 4 +- .../quality_router/quality_router.py | 6 +- litellm/sandbox/__init__.py | 5 +- .../secret_managers/base_secret_manager.py | 4 +- .../cyberark_secret_manager.py | 4 +- litellm/secret_managers/main.py | 3 +- litellm/utils.py | 23 ++-- litellm/vector_stores/main.py | 14 ++- .../code_coverage_tests/recursive_detector.py | 2 +- tests/unit/a2a_protocol/test_utils.py | 14 +++ .../test_container_proxy_ownership.py | 19 +++ .../proxy/hooks/test_managed_files.py | 15 +++ .../dotprompt/test_prompt_manager.py | 38 ++++++ .../integrations/newrelic/test_newrelic.py | 13 ++ .../litellm_core_utils/attribute_forwarder.py | 6 + .../unit/llms/compactifai/test_compactifai.py | 37 +++++- ...gram_audio_transcription_transformation.py | 22 ++++ .../test_edenai_embedding_transformation.py | 6 + ..._edenai_image_generation_transformation.py | 6 + .../test_google_genai_transformation.py | 20 +++ .../test_gemini_image_edit_transformation.py | 20 +++ .../llms/gemini/test_gemini_common_utils.py | 19 +++ ...test_openai_count_tokens_transformation.py | 18 +++ ...test_openai_responses_guardrail_handler.py | 17 +++ .../runwayml/image_generation/__init__.py | 0 ..._runway_image_generation_transformation.py | 87 +++++++++++++ .../unit/proxy/auth/test_mcp_ip_filtering.py | 14 +++ .../client/cli/test_encryption_commands.py | 19 +++ .../common_utils/test_check_batch_cost.py | 11 ++ .../semantic_guard/test_semantic_guard.py | 14 +++ .../guardrail_hooks/test_llm_shield_proxy.py | 17 +++ .../test_storage_backend_service.py | 37 +++++- .../test_openapi_customization.py | 12 ++ 455 files changed, 3670 insertions(+), 1546 deletions(-) create mode 100644 tests/unit/litellm_core_utils/attribute_forwarder.py create mode 100644 tests/unit/llms/runwayml/image_generation/__init__.py create mode 100644 tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py diff --git a/enterprise/enterprise_hooks/aporia_ai.py b/enterprise/enterprise_hooks/aporia_ai.py index 28b49bfce21..1e840f8dcc5 100644 --- a/enterprise/enterprise_hooks/aporia_ai.py +++ b/enterprise/enterprise_hooks/aporia_ai.py @@ -15,9 +15,10 @@ sys.path.insert( ) # Adds the parent directory to the system path import json import sys -from typing import Any, List, Literal, Optional +from typing import Final, List, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -29,11 +30,13 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.llms.openai.common_utils import decoded_json_object from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata from litellm.types.guardrails import GuardrailEventHooks GUARDRAIL_NAME = "aporia" +_DECODED_JSON: Final = TypeAdapter(object) class AporiaGuardrail(CustomGuardrail): @@ -68,7 +71,7 @@ class AporiaGuardrail(CustomGuardrail): async def prepare_aporia_request( self, new_messages: List[dict], response_string: Optional[str] = None ) -> dict: - data: dict[str, Any] = {} + data: dict[str, object] = {} if new_messages is not None: data["messages"] = new_messages if response_string is not None: @@ -122,8 +125,10 @@ class AporiaGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Aporia AI response: %s", response.text) if response.status_code == 200: # check if the response was flagged - _json_response = response.json() - action: str = _json_response.get( + _json_response: Final = decoded_json_object( + _DECODED_JSON.validate_python(response.json()) + ) + action: Final = _json_response.get( "action" ) # possible values are modify, passthrough, block, rephrase if action == "block": diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py index f6998d58cd5..7e4d2514880 100644 --- a/enterprise/enterprise_hooks/blocked_user_list.py +++ b/enterprise/enterprise_hooks/blocked_user_list.py @@ -7,9 +7,10 @@ ## This accepts a list of user id's for whom calls will be rejected -from typing import Literal, Optional +from typing import Final, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._internal_context import with_service_target @@ -19,6 +20,9 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import EndUserRepository + +_REQUEST_BODY_VALUE: Final = TypeAdapter(object) class ENTERPRISE_BlockedUserList(CustomLogger): @@ -94,15 +98,14 @@ class ENTERPRISE_BlockedUserList(CustomLogger): ) if end_user_cache_obj is None and self.prisma_client is not None: # check db - end_user_obj = ( - await self.prisma_client.db.litellm_endusertable.find_unique( - where={"user_id": user} - ) - ) + end_user_id: Final = _REQUEST_BODY_VALUE.validate_python(user) + end_user_obj = await EndUserRepository( + self.prisma_client + ).table.find_unique(where={"user_id": end_user_id}) if end_user_obj is None: # user not in db - assume not blocked end_user_obj = LiteLLM_EndUserTable(user_id=user, blocked=False) cache.set_cache(key=cache_key, value=end_user_obj, ttl=60) - if end_user_obj is not None and end_user_obj.blocked is True: + if end_user_obj.blocked is True: raise HTTPException( status_code=400, detail={ diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index b3ef5a3d496..7dbbc9ed6ce 100644 --- a/enterprise/enterprise_hooks/google_text_moderation.py +++ b/enterprise/enterprise_hooks/google_text_moderation.py @@ -6,7 +6,10 @@ # +-----------------------------------------------+ # Thank you users! We ❤️ you! - Krrish & Ishaan +from typing import Final + from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -15,6 +18,8 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.utils import CallTypesLiteral +_THRESHOLD_VALUE: Final = TypeAdapter(object) + class ENTERPRISE_GoogleTextModeration(CustomLogger): user_api_key_cache = None @@ -69,9 +74,11 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger): f"{category}_confidence_threshold", default_confidence_threshold, ) - set_confidence_value = getattr( - self, - f"{category}_confidence_threshold", + set_confidence_value = _THRESHOLD_VALUE.validate_python( + getattr( + self, + f"{category}_confidence_threshold", + ) ) verbose_proxy_logger.info( f"Google Text Moderation: {category}_confidence_threshold: {set_confidence_value}" diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py b/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py index 14d34f5d1e8..e0190ea8da0 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py @@ -1,15 +1,20 @@ # this is an example endpoint to receive data from litellm +from typing import Final + from fastapi import FastAPI, HTTPException, Request +from pydantic import TypeAdapter app = FastAPI() +_LOGGED_EVENT: Final = TypeAdapter(object) + @app.post("/log-event") async def log_event(request: Request): try: print("Received /log-event request") # noqa # Assuming the incoming request has JSON data - data = await request.json() + data: Final = _LOGGED_EVENT.validate_python(await request.json()) print("Received request data:") # noqa print(data) # noqa diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index 3422e8969b0..e12539eac53 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -8,25 +8,30 @@ ## This provides an LLM Guard Integration for content moderation on the proxy import asyncio -from typing import Final, Optional +from collections.abc import Mapping +from typing import Final import aiohttp from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.openai.common_utils import decoded_json_object from litellm.proxy._types import UserAPIKeyAuth from litellm.secret_managers.main import get_secret_str from litellm.types.utils import CallTypesLiteral +_DECODED_JSON: Final = TypeAdapter(object) + class _ENTERPRISE_LLMGuard(CustomLogger): # Class variables or attributes def __init__( self, mock_testing: bool = False, - mock_redacted_text: Optional[dict] = None, + mock_redacted_text: Mapping[str, object] | None = None, ): self.mock_redacted_text = mock_redacted_text self.llm_guard_mode = litellm.llm_guard_mode @@ -66,7 +71,9 @@ class _ENTERPRISE_LLMGuard(CustomLogger): async with session.post( analyze_url, json={"prompt": text} ) as response: - redacted_text = await response.json() + redacted_text = _DECODED_JSON.validate_python( + await response.json() + ) verbose_proxy_logger.debug( f"LLM Guard: Received response - {redacted_text}" ) @@ -77,12 +84,13 @@ class _ENTERPRISE_LLMGuard(CustomLogger): "error": f"Invalid content moderation response: {redacted_text}" }, ) - if redacted_text.get("is_valid", None) is False: + redacted_response: Final = decoded_json_object(redacted_text) + if redacted_response.get("is_valid", None) is False: raise HTTPException( status_code=400, detail={"error": "Violated content safety policy"}, ) - sanitized_prompt = redacted_text.get("sanitized_prompt") + sanitized_prompt: Final = redacted_response.get("sanitized_prompt") return sanitized_prompt if isinstance(sanitized_prompt, str) else text except Exception as e: verbose_proxy_logger.exception( diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index df1f263a947..f3b1240f733 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1040,8 +1040,8 @@ class CheckBatchCost: from litellm.proxy.proxy_server import proxy_config from litellm.secret_managers.main import str_to_bool - api_key: Final = getattr(job, "api_key", None) - team_id: Final = getattr(job, "team_id", None) + api_key: Final = job.api_key if hasattr(job, "api_key") else None + team_id: Final = job.team_id if hasattr(job, "team_id") else None has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key) try: key_row: Final = ( diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bc35c4a0aba..48cbc0ab8ff 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -44,6 +44,7 @@ from litellm.constants import ( ) from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.internal_call_metadata import requests_background from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -133,7 +134,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient - Span = Union[_Span, Any] + Span = _Span InternalUsageCache = _InternalUsageCache PrismaClient = _PrismaClient else: @@ -490,6 +491,14 @@ _PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDE _PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0) +class _SupportsKeyedGet(Protocol): + def get(self, key: str, default: object, /) -> object: ... + + +def _keyed_get(source: object, key: str, default: object = None) -> object: + return require_attributes(source, _SupportsKeyedGet, "get").get(key, default) + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -1611,7 +1620,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return holders if holders or "input_file_id" in request else healthy_deployments async def async_pre_call_deployment_hook( - self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] + self, kwargs: dict[str, object], call_type: Optional[CallTypes] ) -> Optional[dict]: """ Allow modifying the request just before it's sent to the deployment. @@ -1623,11 +1632,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) - model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) + model_id = cast(Optional[str], _keyed_get(kwargs.get("model_info", {}), "id")) if model_id is None: model_id = cast( Optional[str], - kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), + _keyed_get(_keyed_get(kwargs.get("litellm_metadata", {}), "model_info", {}), "id"), ) for accessor_key in accessor_keys: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) @@ -1824,7 +1833,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span) -> dict: + async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span | None) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id @@ -1866,7 +1875,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): llm_router: Optional[Router], _create_file_request: CreateFileRequest, target_model_names_list: List[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, ) -> List[OpenAIFileObject]: if llm_router is None: raise Exception("LLM Router not initialized. Ensure models added to proxy.") @@ -1882,7 +1891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): create_file_request: CreateFileRequest, llm_router: Router, target_model_names_list: List[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, user_api_key_dict: UserAPIKeyAuth, ) -> OpenAIFileObject: responses = await self.create_file_for_each_model( @@ -1949,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_objects: List[OpenAIFileObject], create_file_request: CreateFileRequest, internal_usage_cache: InternalUsageCache, - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, target_model_names_list: List[str], ) -> OpenAIFileObject: ## GET THE FILE TYPE FROM THE CREATE FILE REQUEST @@ -2652,7 +2661,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Continue with other files even if one fails continue - def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str: + def _get_content_type_from_file_object(self, file_object: object) -> str: """ Determine content type from file object. diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index 0e9657aba96..fcaf58099b9 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient - Span = Union[_Span, Any] + Span = _Span InternalUsageCache = _InternalUsageCache PrismaClient = _PrismaClient else: @@ -107,7 +107,7 @@ class _PROXY_LiteLLMManagedVectorStores( llm_router: Router, model: str, request_data: dict[str, object] | VectorStoreCreateOptionalRequestParams, - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, ) -> VectorStoreCreateResponse: """ Create a vector store for a specific model. @@ -134,7 +134,7 @@ class _PROXY_LiteLLMManagedVectorStores( create_request: VectorStoreCreateOptionalRequestParams, llm_router: Router, target_model_names_list: List[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, user_api_key_dict: UserAPIKeyAuth, ) -> VectorStoreCreateResponse: """ diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index db419fde1e6..b91b4a970e8 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -36,6 +36,7 @@ from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository @@ -73,6 +74,16 @@ def _verification_token_table( return VerificationTokenRepository(prisma_client).table +class _WriterVerificationTokenRepository(PrismaTableRepository["prisma_models.LiteLLM_VerificationToken"]): + table_name = "litellm_verificationtoken" + + +def _writer_verification_token_table( + prisma_client: PrismaClient, +) -> TableActions["prisma_models.LiteLLM_VerificationToken"]: + return _WriterVerificationTokenRepository(prisma_client, use_writer=True).table + + def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: return BudgetRepository(prisma_client).table @@ -849,7 +860,7 @@ async def update_project( else None ) if current_project is not None and data.team_id != current_project.team_id: - mismatched_key_count: Final = await prisma_client.writer_db.litellm_verificationtoken.count( + mismatched_key_count: Final = await _writer_verification_token_table(prisma_client).count( where={ "project_id": data.project_id, "OR": [{"team_id": {"not": data.team_id}}, {"team_id": None}], diff --git a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py index 9b87f4f1a7f..19d4988452b 100644 --- a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py +++ b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py @@ -14,6 +14,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final, List, Optional, Protocol from fastapi import APIRouter, Depends, HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -37,6 +38,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient router = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) class ManagedVectorStoreRow(Protocol): @@ -348,11 +350,11 @@ async def update_vector_store( raise HTTPException(status_code=500, detail="Database not connected") try: - update_data = data.model_dump(exclude_unset=True) - vector_store_id: Final[str] = update_data.pop("vector_store_id") + update_data = data.model_dump(exclude_unset=True, exclude={"vector_store_id"}) + vector_store_id: Final = data.vector_store_id if update_data.get("vector_store_metadata") is not None: update_data["vector_store_metadata"] = safe_dumps( - update_data["vector_store_metadata"] + _DECODED_JSON.validate_python(update_data["vector_store_metadata"]) ) updated = await managed_vector_store_table(prisma_client).update( diff --git a/litellm/__init__.py b/litellm/__init__.py index 1784deeb797..be57c8f3f16 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -175,7 +175,7 @@ argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload. gcs_pub_sub_use_v1: Optional[bool] = False # if you want to use v1 gcs pubsub logged payload generic_api_use_v1: Optional[bool] = False # if you want to use v1 generic api logged payload -argilla_transformation_object: Optional[Dict[str, Any]] = None +argilla_transformation_object: Mapping[str, str] | None = None _async_input_callback: List[Union[str, Callable[..., object], "CustomLogger"]] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. @@ -333,7 +333,7 @@ in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False enable_azure_ad_token_refresh: Optional[bool] = False # Proxy Authentication - auto-obtain/refresh OAuth2/JWT tokens for LiteLLM Proxy -proxy_auth: Optional[Any] = None +proxy_auth: _ProxyAuthHandlerType | None = None ### DEFAULT AZURE API VERSION ### AZURE_DEFAULT_API_VERSION = "2025-02-01-preview" # this is updated to the latest ### DEFAULT WATSONX API VERSION ### @@ -503,7 +503,7 @@ skill_search_embedding_model: Optional[str] = None # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) -public_model_groups_links: Dict[str, Union[str, Dict[str, Any]]] = {} +public_model_groups_links: dict[str, str | dict[str, object]] = {} #### REQUEST PRIORITIZATION ####### priority_reservation: Optional[Dict[str, Union[float, "PriorityReservationDict"]]] = None # priority_reservation_settings is lazy-loaded via __getattr__ @@ -2289,6 +2289,10 @@ if TYPE_CHECKING: LiteLLM_UpperboundKeyGenerateParams, ) from litellm.utils import ModelResponseListIterator as _ModelResponseListIterator + from collections.abc import Awaitable + from litellm.types.utils import ModelResponse as _ModelResponseType + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper as _CustomStreamWrapperType + from litellm.proxy_auth import ProxyAuthHandler as _ProxyAuthHandlerType # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -2312,7 +2316,7 @@ if TYPE_CHECKING: supports_audio_output: Callable[..., bool] supports_system_messages: Callable[..., bool] supports_reasoning: Callable[..., bool] - acreate: Callable[..., Any] + acreate: Callable[..., Awaitable[_ModelResponseType | _CustomStreamWrapperType]] get_max_tokens: Callable[..., int] get_model_info: Callable[..., _ModelInfoType] register_prompt_template: Callable[..., None] @@ -2610,8 +2614,10 @@ def __getattr__(name: str) -> Any: cancel_run, delete_run, ) + from types import MappingProxyType - return locals()[name] + evals_exports: Final[Mapping[str, object]] = MappingProxyType(locals()) + return evals_exports[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 46126c444a6..1a7dc343745 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -19,7 +19,7 @@ import importlib import sys from collections.abc import Callable, Mapping from types import ModuleType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final from typing_extensions import ReadOnly, TypedDict @@ -456,7 +456,7 @@ def _lazy_import_http_handlers(name: str) -> object: params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"} # Create the client instance - provider_id: Final = cast(Any, "litellm_module_level_client") + provider_id: Final = "litellm_module_level_client" async_client: Final = get_async_httpx_client( llm_provider=provider_id, params=params, diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py index 31ca81c44dd..9a1431d6457 100644 --- a/litellm/a2a_protocol/cost_calculator.py +++ b/litellm/a2a_protocol/cost_calculator.py @@ -5,7 +5,10 @@ Supports dynamic cost parameters that allow platform owners to define custom costs per agent query or per token. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, SupportsFloat, SupportsIndex + +from typing_extensions import Buffer if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( @@ -15,6 +18,15 @@ else: LitellmLoggingObject = Any +def _custom_response_cost(model_call_details: Mapping[str, object]) -> float | None: + response_cost: Final = model_call_details.get("response_cost", None) + if response_cost is None: + return None + if isinstance(response_cost, (str, Buffer, SupportsFloat, SupportsIndex)): + return float(response_cost) + raise TypeError(f"float() argument must be a string or a real number, not {type(response_cost).__name__!r}") + + class A2ACostCalculator: @staticmethod def calculate_a2a_cost( @@ -45,9 +57,9 @@ class A2ACostCalculator: model_call_details: Final = litellm_logging_obj.model_call_details # Check if user set a custom response cost (backward compatibility) - response_cost: Final = model_call_details.get("response_cost", None) - if response_cost is not None: - return float(response_cost) + custom_response_cost: Final = _custom_response_cost(model_call_details) + if custom_response_cost is not None: + return custom_response_cost # Get litellm_params for cost parameters litellm_params: Final = model_call_details.get("litellm_params", {}) or {} diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 404bcca712f..1cea0194715 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -84,14 +84,23 @@ def _filter_reserved_headers( return filtered or None +def _calling_key_hash(litellm_params: Mapping[str, object]) -> str | None: + principal: Final = litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) + if principal is None or isinstance(principal, str): + return principal + if not principal: + return None + raise AttributeError(f"{type(principal).__name__!r} object has no attribute 'encode'", name="encode", obj=principal) + + def _request_scoped_runtime_session_id( params: Mapping[str, object], - litellm_params: Mapping[str, Any], + litellm_params: Mapping[str, object], ) -> str | None: context_id: Final = get_session_id_from_a2a_params(params) if not isinstance(context_id, str) or not context_id: return None - return scope_session_to_principal(context_id, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)) + return scope_session_to_principal(context_id, _calling_key_hash(litellm_params)) def _validate_runtime_session_id(session_id: str, model: str) -> str: @@ -109,6 +118,13 @@ def _validate_runtime_session_id(session_id: str, model: str) -> str: ) +def _agentcore_model_name(litellm_params: Mapping[str, object]) -> str: + model: Final = litellm_params.get("model", "") + if isinstance(model, str): + return model + raise AttributeError(f"{type(model).__name__!r} object has no attribute 'startswith'", name="startswith", obj=model) + + class BedrockAgentCoreA2ATransformation: """ Request/response transformation for Bedrock AgentCore A2A agents. @@ -152,11 +168,8 @@ class BedrockAgentCoreA2ATransformation: """ # Extract model and strip the "bedrock/" prefix # "bedrock/agentcore/arn:aws:..." → "agentcore/arn:aws:..." - model: Final = litellm_params.get("model", "") - if model.startswith("bedrock/"): - agentcore_model = model[len("bedrock/") :] - else: - agentcore_model = model + model: Final = _agentcore_model_name(litellm_params) + agentcore_model: Final = model.removeprefix("bedrock/") # Build optional_params from litellm_params (everything except model and custom_llm_provider) optional_params: Final = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")} diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index 57c4a4677d0..c2d581321cd 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -8,11 +8,31 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model: """ import asyncio -from collections.abc import AsyncIterator, Mapping -from typing import Any, Final +from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from itertools import chain +from typing import Final from uuid import uuid4 from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items +from litellm.llms.openai.common_utils import decoded_json_object + + +def _a2a_text_part(part: object) -> Sequence[object]: + part_fields: Final = as_str_mapping(part) + if not isinstance(part_fields, dict): + return () + if part_fields.get("kind") in (None, "", "text") and part_fields.get("text"): + return (part_fields["text"],) + return () + + +def _join_with_spaces(texts: Sequence[object]) -> str: + strings: Final = [text for text in texts if isinstance(text, str)] + if len(strings) == len(texts): + return " ".join(strings) + first_bad: Final = next(position for position, text in enumerate(texts) if not isinstance(text, str)) + raise TypeError(f"sequence item {first_bad}: expected str instance, {type(texts[first_bad]).__name__} found") class WatsonxOrchestrateTransformation: @@ -29,22 +49,16 @@ class WatsonxOrchestrateTransformation: return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}" @staticmethod - def extract_text_from_a2a_params(params: dict[str, Any]) -> str: + def extract_text_from_a2a_params(params: Mapping[str, object]) -> str: """ Extract user message text from A2A MessageSendParams. A2A format: params.message.parts[*] where part.kind == "text" """ - message: Final = params.get("message", {}) - parts: Final = message.get("parts", []) - texts: Final = [] - for part in parts: - if not isinstance(part, dict): - continue - kind = part.get("kind") - if kind in (None, "", "text") and part.get("text"): - texts.append(part["text"]) - return " ".join(texts) or "" + message: Final = decoded_json_object(params.get("message", {})) + parts: Final[Iterable[object]] = iter_items(message.get("parts", [])) + texts: Final = list(chain.from_iterable(map(_a2a_text_part, parts))) + return _join_with_spaces(texts) or "" @staticmethod def build_wxo_run_body( diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 7400844bf28..17d1b86a609 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -4,15 +4,20 @@ Utility functions for A2A protocol. import hashlib from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import has_attributes if TYPE_CHECKING: from a2a.types import SendMessageRequest, SendStreamingMessageRequest +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str) -> object: ... + + class A2ARequestUtils: """Utility class for A2A request/response processing.""" @@ -74,7 +79,7 @@ class A2ARequestUtils: @staticmethod def get_input_message_from_request( request: "SendMessageRequest | SendStreamingMessageRequest", - ) -> Any: + ) -> object: """ Extract the input message from an A2A request. @@ -84,7 +89,7 @@ class A2ARequestUtils: Returns: The message object/dict or None """ - params: Final = getattr(request, "params", None) + params: Final[object] = getattr(request, "params", None) if params is None: return None return getattr(params, "message", None) @@ -127,10 +132,12 @@ class A2ARequestUtils: # the dict branch — request-side parts are a2a-sdk Part RootModels whose # kind/text live on part.root, which the object branch cannot read. This # mirrors how the response side already works (it operates on model_dump). - input_message = A2ARequestUtils.get_input_message_from_request(request) - if input_message is not None and hasattr(input_message, "model_dump"): - input_message = input_message.model_dump(mode="json") - input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) + input_message: Final = A2ARequestUtils.get_input_message_from_request(request) + input_text: Final = A2ARequestUtils.extract_text_from_message( + input_message.model_dump(mode="json") + if has_attributes(input_message, _SupportsModelDump, "model_dump") + else input_message + ) prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) # Count output tokens diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 61e36bd1e96..fec1d8e4465 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -45,6 +45,7 @@ _BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_na BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] _BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) +_DECODED_JSON: Final = TypeAdapter(object) class GetAnthropicBetaHeadersConfig: @@ -76,7 +77,7 @@ class GetAnthropicBetaHeadersConfig: } @staticmethod - def _check_is_valid_dict(fetched_config: dict) -> bool: + def _check_is_valid_dict(fetched_config: object) -> bool: """Check if fetched config is a non-empty dict with expected structure.""" if not isinstance(fetched_config, dict): verbose_logger.warning( @@ -110,7 +111,7 @@ class GetAnthropicBetaHeadersConfig: return True @classmethod - def validate_beta_headers_config(cls, fetched_config: dict) -> bool: + def validate_beta_headers_config(cls, fetched_config: object) -> bool: """ Validate the integrity of a fetched beta headers config. @@ -119,7 +120,7 @@ class GetAnthropicBetaHeadersConfig: return cls._check_is_valid_dict(fetched_config) @staticmethod - def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> dict: + def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> object: """ Fetch the beta headers config from a remote URL. @@ -128,7 +129,7 @@ class GetAnthropicBetaHeadersConfig: """ response: Final = httpx.get(url, timeout=timeout) response.raise_for_status() - return response.json() + return _DECODED_JSON.validate_python(response.json()) def get_beta_headers_config(url: str) -> dict: @@ -161,7 +162,9 @@ def get_beta_headers_config(url: str) -> dict: return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() # Validate the fetched config - if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content): + if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content) or not isinstance( + content, dict + ): verbose_logger.warning( "LiteLLM: Fetched beta headers config failed integrity check. Using local backup instead. url=%s", url, diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index eb3ec8aaee2..d8198e41239 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -7,10 +7,14 @@ Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthrop import json from typing import Final +from pydantic import TypeAdapter + from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from .exceptions import AnthropicErrorDetail, AnthropicErrorResponse, AnthropicErrorType +_PARSED_ERROR_BODY: Final = TypeAdapter(object) + # HTTP status code -> Anthropic error type # Source: https://docs.anthropic.com/en/api/errors ANTHROPIC_ERROR_TYPE_MAP: Final[dict[int, AnthropicErrorType]] = { @@ -79,7 +83,7 @@ class AnthropicExceptionMapping: - Generic: {"message": "..."} - Plain strings """ - parsed: Final = safe_json_loads(raw_message) + parsed: Final = _PARSED_ERROR_BODY.validate_python(safe_json_loads(raw_message)) if isinstance(parsed, dict): # Bedrock format if "detail" in parsed and isinstance(parsed["detail"], dict): @@ -145,9 +149,8 @@ class AnthropicExceptionMapping: AnthropicErrorResponse dict """ # Try to parse as JSON once - parsed: dict | None = safe_json_loads(raw_message) - if not isinstance(parsed, dict): - parsed = None + decoded: Final = _PARSED_ERROR_BODY.validate_python(safe_json_loads(raw_message)) + parsed: Final[dict[str, object] | None] = decoded if isinstance(decoded, dict) else None # If parsed and already in Anthropic format - passthrough if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(parsed): diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index a3428261636..a5721c6f047 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -77,7 +77,7 @@ def get_assistants( api_base: str | None = None, api_version: str | None = None, **kwargs, -) -> SyncCursorPage[Assistant]: +) -> SyncCursorPage[Assistant] | Coroutine[None, None, AsyncCursorPage[Assistant]]: aget_assistants: Final[bool | None] = kwargs.pop("aget_assistants", None) if aget_assistants is not None and not isinstance(aget_assistants, bool): raise Exception("Invalid value passed in for aget_assistants. Only bool or None allowed") @@ -100,7 +100,7 @@ def get_assistants( elif timeout is None: timeout = 600.0 - response: SyncCursorPage[Assistant] | None = None + response: SyncCursorPage[Assistant] | Coroutine[None, None, AsyncCursorPage[Assistant]] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -564,7 +564,7 @@ def create_thread( tool_resources: OpenAICreateThreadParamsToolResources | None = None, client: OpenAI | None = None, **kwargs, -) -> Thread: +) -> Thread | Coroutine[None, None, Thread]: """ - get the llm provider - if openai - route it there @@ -610,7 +610,7 @@ def create_thread( api_base: str | None = None api_key: str | None = None - response: Thread | None = None + response: Thread | Coroutine[None, None, Thread] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -735,7 +735,7 @@ def get_thread( thread_id: str, client=None, **kwargs, -) -> Thread: +) -> Thread | Coroutine[None, None, Thread]: """Get the thread object, given a thread_id""" aget_thread: Final = kwargs.pop("aget_thread", None) optional_params: Final = GenericLiteLLMParams(**kwargs) @@ -757,7 +757,7 @@ def get_thread( timeout = 600.0 api_base: str | None = None api_key: str | None = None - response: Thread | None = None + response: Thread | Coroutine[None, None, Thread] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -903,7 +903,7 @@ def add_message( metadata: dict | None = None, client=None, **kwargs, -) -> OpenAIMessage: +) -> OpenAIMessage | Coroutine[None, None, OpenAIMessage]: ### COMMON OBJECTS ### a_add_message: Final = kwargs.pop("a_add_message", None) _message_data: Final = MessageData(role=role, content=content, attachments=attachments, metadata=metadata) @@ -935,7 +935,7 @@ def add_message( timeout = 600.0 api_key: str | None = None api_base: str | None = None - response: OpenAIMessage | None = None + response: OpenAIMessage | Coroutine[None, None, OpenAIMessage] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -1065,7 +1065,7 @@ def get_messages( thread_id: str, client: Any | None = None, **kwargs, -) -> SyncCursorPage[OpenAIMessage]: +) -> SyncCursorPage[OpenAIMessage] | Coroutine[None, None, AsyncCursorPage[OpenAIMessage]]: aget_messages: Final = kwargs.pop("aget_messages", None) optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) @@ -1086,7 +1086,7 @@ def get_messages( elif timeout is None: timeout = 600.0 - response: SyncCursorPage[OpenAIMessage] | None = None + response: SyncCursorPage[OpenAIMessage] | Coroutine[None, None, AsyncCursorPage[OpenAIMessage]] | None = None api_key: str | None = None api_base: str | None = None if custom_llm_provider == "openai": @@ -1256,7 +1256,7 @@ def run_thread( client: object | None = None, event_handler: AssistantEventHandler | None = None, # for stream=True calls **kwargs, -) -> Run: +) -> Run | Coroutine[None, None, Run]: """Run a given thread + assistant.""" arun_thread: Final = kwargs.pop("arun_thread", None) optional_params: Final = GenericLiteLLMParams(**kwargs) @@ -1278,7 +1278,7 @@ def run_thread( elif timeout is None: timeout = 600.0 - response: Run | None = None + response: Run | Coroutine[None, None, Run] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index d82e58ecc91..0e7589d88ab 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -13,7 +13,7 @@ This module is dependency-injected: callers pass the proxy ``llm_router`` and from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import litellm from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS @@ -25,12 +25,12 @@ if TYPE_CHECKING: def resolve_embedding_router( embedding_model: str, llm_router: Router | None, - llm_model_list: list[dict[str, Any]] | None, + llm_model_list: Sequence[Mapping[str, object]] | None, ) -> Router | None: """Return ``llm_router`` iff it serves ``embedding_model`` as a deployment.""" if llm_router is None: return None - router_model_names: Final[list[str]] = ( + router_model_names: Final[Sequence[object]] = ( [m["model_name"] for m in llm_model_list if "model_name" in m] if llm_model_list is not None else [] ) if embedding_model in router_model_names: diff --git a/litellm/caching/azure_blob_cache.py b/litellm/caching/azure_blob_cache.py index bf0e15421bc..5b9071dc920 100644 --- a/litellm/caching/azure_blob_cache.py +++ b/litellm/caching/azure_blob_cache.py @@ -13,10 +13,14 @@ import json from contextlib import suppress from typing import Final +from pydantic import TypeAdapter + from litellm._logging import print_verbose, verbose_logger from .base_cache import BaseCache +_DECODED_BLOB: Final = TypeAdapter(object) + class AzureBlobCache(BaseCache): def __init__(self, account_url, container) -> None: @@ -65,7 +69,7 @@ class AzureBlobCache(BaseCache): print_verbose(f"Get Azure Blob Cache: key: {key}") as_bytes: Final = self.container_client.download_blob(key).readall() as_str: Final = as_bytes.decode("utf-8") - cached_response: Final[object] = json.loads(as_str) + cached_response: Final = _DECODED_BLOB.validate_python(json.loads(as_str)) verbose_logger.debug( "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s", @@ -86,7 +90,7 @@ class AzureBlobCache(BaseCache): blob: Final = await self.async_container_client.download_blob(key) as_bytes: Final = await blob.readall() as_str: Final = as_bytes.decode("utf-8") - cached_response: Final[object] = json.loads(as_str) + cached_response: Final = _DECODED_BLOB.validate_python(json.loads(as_str)) verbose_logger.debug( "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s", key, diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py index 51c169ba796..fb91d13a8c9 100644 --- a/litellm/caching/base_cache.py +++ b/litellm/caching/base_cache.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index 50939ad51ca..9470f23570a 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -1,15 +1,19 @@ import json from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter + from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any +_DECODED_DISK_VALUE: Final = TypeAdapter(object) + class DiskCache(BaseCache): def __init__(self, disk_cache_dir: str | None = None): @@ -44,7 +48,7 @@ class DiskCache(BaseCache): original_cached_response: Final = self.disk_cache.get(key) if original_cached_response: try: - cached_response = json.loads(original_cached_response) + cached_response = _DECODED_DISK_VALUE.validate_python(json.loads(original_cached_response)) except Exception: cached_response = original_cached_response return cached_response diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 98f6f1409f2..a0da0f6ddca 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1200,12 +1200,16 @@ def get_usage_object( elif isinstance(usage_obj, dict): return Usage(**usage_obj) elif isinstance(usage_obj, BaseModel): - return Usage(**usage_obj.model_dump()) + return _usage_from_model(usage_obj) else: verbose_logger.debug("Unknown usage object type: %s, usage_obj: %s", type(usage_obj), usage_obj) return None +def _usage_from_model(usage_model: BaseModel) -> Usage: + return Usage(**usage_model.model_dump()) + + def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper: match usage: case DecisionsUsage(): @@ -1686,7 +1690,7 @@ def completion_cost( setattr( completion_response, "usage", - litellm.Usage(**_usage_for_dump.model_dump()), + _usage_from_model(_usage_for_dump), ) if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)): _usage = _decisions_usage(usage_obj).model_dump() @@ -3035,8 +3039,11 @@ def _batch_prompt_cost( ) +_ATTRIBUTE_VALUE: Final = TypeAdapter(object) + + def _attribute_value(obj: object, name: str) -> object: - return getattr(obj, name) + return _ATTRIBUTE_VALUE.validate_python(getattr(obj, name)) def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]: @@ -3671,7 +3678,26 @@ def _get_model_info_or_none(model: str, custom_llm_provider: str) -> ModelInfo | return None -def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[str, ...]) -> float | None: +_TranscriptionRateKey: TypeAlias = Literal[ + "input_cost_per_second", "input_cost_per_audio_token", "input_cost_per_token", "output_cost_per_token" +] + + +def _model_info_rate(info: ModelInfo, key: _TranscriptionRateKey) -> float | None: + match key: + case "input_cost_per_second": + return info.get("input_cost_per_second") + case "input_cost_per_audio_token": + return info.get("input_cost_per_audio_token") + case "input_cost_per_token": + return info.get("input_cost_per_token") + case "output_cost_per_token": + return info.get("output_cost_per_token") + case _: + assert_never(key) + + +def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[_TranscriptionRateKey, ...]) -> float | None: """First of ``keys`` this entry prices, read off the raw ``litellm.model_cost`` entry because ``get_model_info`` synthesizes zero token rates for entries that omit them.""" if info is None: @@ -3680,12 +3706,18 @@ def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[str, ...]) if declared is None: return None return next( - (float(value) for key in keys if declared.get(key) is not None and (value := info.get(key)) is not None), + ( + float(value) + for key in keys + if declared.get(key) is not None and (value := _model_info_rate(info, key)) is not None + ), None, ) -def _transcription_rate(keys: tuple[str, ...], override: ModelInfo | None, base: ModelInfo | None) -> float: +def _transcription_rate( + keys: tuple[_TranscriptionRateKey, ...], override: ModelInfo | None, base: ModelInfo | None +) -> float: rates: Final = (_declared_transcription_rate(info, keys) for info in (override, base)) return next((rate for rate in rates if rate is not None), 0.0) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 733f1aa72d5..b8e933bf046 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1273,7 +1273,7 @@ class MidStreamFallbackError(ServiceUnavailableError): self.message = _saved_message self.args = (_saved_message,) - original_headers: Final = getattr(original_exception, "litellm_response_headers", None) + original_headers: Final[object] = getattr(original_exception, "litellm_response_headers", None) if original_headers is not None: self.litellm_response_headers = original_headers diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 8015b065899..4f204b3629e 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -10,6 +10,7 @@ from mcp.types import Tool as MCPTool from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -20,6 +21,8 @@ from litellm.constants import ( from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall +_DECODED_JSON: Final = TypeAdapter(object) + ######################################################## # List MCP Tool functions @@ -200,15 +203,18 @@ async def call_mcp_tool( return tool_result +def _decode_json_arguments(arguments: str) -> object: + try: + return _DECODED_JSON.validate_python(json.loads(arguments)) + except json.JSONDecodeError: + return {} + + def _get_function_arguments(function: FunctionDefinition) -> dict: """Helper to safely get and parse function arguments.""" - arguments = function.get("arguments", {}) - if isinstance(arguments, str): - try: - arguments = json.loads(arguments) - except json.JSONDecodeError: - arguments = {} - return arguments if isinstance(arguments, dict) else {} + arguments: Final = function.get("arguments", {}) + decoded: Final = _decode_json_arguments(arguments) if isinstance(arguments, str) else arguments + return decoded if isinstance(decoded, dict) else {} def transform_openai_tool_call_request_to_mcp_tool_call_request( diff --git a/litellm/files/main.py b/litellm/files/main.py index 16ff5b02f9c..c5172bc01ae 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -9,7 +9,7 @@ import asyncio import contextvars import time import uuid as uuid_module -from collections.abc import Coroutine +from collections.abc import Awaitable, Coroutine from functools import partial from typing import Final, Literal, cast @@ -293,6 +293,14 @@ def create_file( raise e +async def _await_retrieved_file(pending: Awaitable[FileObject | OpenAIFileObject]) -> FileObject | OpenAIFileObject: + return await pending + + +async def _await_deleted_file(pending: Awaitable[object]) -> object: + return await pending + + @client async def afile_retrieve( file_id: str, @@ -325,7 +333,7 @@ async def afile_retrieve( func_with_context: Final = partial(ctx.run, func) init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): - response = await init_response + response = await _await_retrieved_file(init_response) else: response = init_response @@ -505,7 +513,7 @@ async def afile_delete( func_with_context: Final = partial(ctx.run, func) init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): - response = await init_response + response = await _await_deleted_file(init_response) else: response = init_response diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index 460050872f4..e9a030e0e32 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -1,7 +1,7 @@ import datetime import traceback from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final, Optional, cast +from typing import TYPE_CHECKING, Final, Optional, cast import anyio @@ -163,7 +163,7 @@ class FileContentStreamingResponse: merged_hidden_params: Final = cast( "StandardLoggingHiddenParams", { - **cast(dict[str, Any], payload.get("hidden_params") or {}), + **(payload.get("hidden_params") or {}), **self._hidden_params, }, ) diff --git a/litellm/harness/handlers/base.py b/litellm/harness/handlers/base.py index 18b7391e988..4721e9d6ff2 100644 --- a/litellm/harness/handlers/base.py +++ b/litellm/harness/handlers/base.py @@ -4,7 +4,6 @@ from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import AsyncIterator -from typing import Any from litellm.harness.context import SessionContext from litellm.harness.errors import CapabilityUnsupported @@ -38,5 +37,5 @@ class BaseHarnessHandler(ABC): async def resume(self, ctx: SessionContext, native_session_id: str) -> None: """Continue the runtime's own session on the next turn.""" - async def history(self, ctx: SessionContext) -> list[dict[str, Any]]: # mutable-ok: public history() API shape + async def history(self, ctx: SessionContext) -> list[dict[str, object]]: # mutable-ok: public history() API shape raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history") diff --git a/litellm/harness/handlers/cli_handler.py b/litellm/harness/handlers/cli_handler.py index 6320a877e6b..b01f449a67d 100644 --- a/litellm/harness/handlers/cli_handler.py +++ b/litellm/harness/handlers/cli_handler.py @@ -70,10 +70,10 @@ async def send_stdin(proc: Process, data: str) -> None: async def private_dir_for(sandbox: Sandbox) -> str: - tempdir = getattr(sandbox, "tempdir", None) + tempdir: Final = sandbox.tempdir if hasattr(sandbox, "tempdir") else None if tempdir is None: raise SandboxError(f"{type(sandbox).__name__} has no tempdir(); CLI harnesses need a private config dir") - path: str = await tempdir() + path: Final = await tempdir() return path diff --git a/litellm/harness/handlers/deepagents_handler.py b/litellm/harness/handlers/deepagents_handler.py index 4a36b1a08d1..f6785323292 100644 --- a/litellm/harness/handlers/deepagents_handler.py +++ b/litellm/harness/handlers/deepagents_handler.py @@ -161,7 +161,7 @@ class DeepAgentsHandler(BaseHarnessHandler): async def history( self, ctx: SessionContext - ) -> list[dict[str, Any]]: # mutable-ok: BaseHarnessHandler.history API returns OpenAI message dicts + ) -> list[dict[str, object]]: # mutable-ok: BaseHarnessHandler.history API returns OpenAI message dicts agent, deps = self._require_agent() snapshot = await agent.aget_state(self._run_config(ctx, None)) messages = (snapshot.values or MappingProxyType({})).get("messages") or () diff --git a/litellm/harness/runtime.py b/litellm/harness/runtime.py index dc24a1f0ef7..1d946521bfa 100644 --- a/litellm/harness/runtime.py +++ b/litellm/harness/runtime.py @@ -12,11 +12,10 @@ import inspect import logging import os import uuid -from collections.abc import AsyncIterator, Callable, Coroutine, Generator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Generator, Mapping, Sequence from dataclasses import dataclass, field from types import MappingProxyType from typing import ( - Any, Final, get_args, ) @@ -331,11 +330,15 @@ async def _put_end(queue: asyncio.Queue[Event | _End], end: _End) -> None: queue.put_nowait(end) +async def _await_decision(pending: Awaitable[object]) -> object: + return await pending + + async def call_approval_handler(handler: ApprovalHandler, approval: Approval) -> None: """Run on_approval (sync in a worker thread, or async) and resolve approval.""" try: if inspect.iscoroutinefunction(handler): - decision: object = await handler(approval) + decision: object = await _await_decision(handler(approval)) else: decision = await asyncio.to_thread(handler, approval) if inspect.isawaitable(decision): @@ -749,7 +752,7 @@ class AsyncSession: async def history( self, - ) -> list[dict[str, Any]]: # mutable-ok: public API returns OpenAI-format message dicts from the handler + ) -> list[dict[str, object]]: # mutable-ok: public API returns OpenAI-format message dicts from the handler if not self.harness_config.capabilities.history: raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history") await self._ensure_ready() diff --git a/litellm/images/main.py b/litellm/images/main.py index 7f94835d59b..29bac3ffe05 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -3,7 +3,7 @@ import contextvars import importlib from collections.abc import Coroutine from functools import partial -from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload +from typing import TYPE_CHECKING, Final, Literal, cast, overload from pydantic import TypeAdapter, ValidationError @@ -70,16 +70,17 @@ from litellm.utils import ( _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) # Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls -_ImageEditRequestUtils_cache: Optional["ImageEditRequestUtils"] = None +_ImageEditRequestUtils_cache: type["ImageEditRequestUtils"] | None = None -def _get_ImageEditRequestUtils() -> "ImageEditRequestUtils": +def _get_ImageEditRequestUtils() -> type["ImageEditRequestUtils"]: """Get ImageEditRequestUtils, loading it lazily if needed.""" global _ImageEditRequestUtils_cache if _ImageEditRequestUtils_cache is None: # Access via module to trigger __getattr__ if not cached - module: Final = importlib.import_module(__name__) - _ImageEditRequestUtils_cache = module.ImageEditRequestUtils + from litellm.images import main as images_main + + _ImageEditRequestUtils_cache = images_main.ImageEditRequestUtils assert _ImageEditRequestUtils_cache is not None # Type narrowing for type checker return _ImageEditRequestUtils_cache diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index 6f986144d4c..87baee91780 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -115,14 +115,16 @@ class AlertingHangingRequestCheck: if hanging_request_data.alerted: continue - request_status = await proxy_logging_obj.internal_usage_cache.async_get_cache( - key=f"request_status:{hanging_request_data.request_id}", - litellm_parent_otel_span=None, - local_only=True, - ) # this means the request status was either success or fail # and is not hanging - if request_status is not None: + if ( + await proxy_logging_obj.internal_usage_cache.async_get_cache( + key=f"request_status:{hanging_request_data.request_id}", + litellm_parent_otel_span=None, + local_only=True, + ) + is not None + ): # clear this request from hanging request cache since the request was either success or failed self.hanging_request_cache._remove_key( key=request_id, diff --git a/litellm/integrations/SlackAlerting/user_spend_alerts.py b/litellm/integrations/SlackAlerting/user_spend_alerts.py index 38794735c1b..08799a2f59f 100644 --- a/litellm/integrations/SlackAlerting/user_spend_alerts.py +++ b/litellm/integrations/SlackAlerting/user_spend_alerts.py @@ -56,8 +56,9 @@ async def fetch_user_spend_rows( today_str: Final = today.strftime("%Y-%m-%d") month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d") baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d") - raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) - return USER_SPEND_ROWS_ADAPTER.validate_python(raw) + return USER_SPEND_ROWS_ADAPTER.validate_python( + await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) + ) def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None: diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index 9043ab1853a..3364520a8ae 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -4,11 +4,16 @@ AgentOps integration for LiteLLM - Provides OpenTelemetry tracing for LLM calls import os from dataclasses import dataclass -from typing import Any, Final +from typing import Final + +from pydantic import TypeAdapter from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import get_httpx_client +_DECODED_JSON: Final = TypeAdapter(object) + @dataclass class AgentOpsConfig: @@ -58,9 +63,10 @@ class AgentOps(OpenTelemetry): project_id = None if config.api_key: try: - response: Final = self._fetch_auth_token(config.api_key, config.auth_endpoint) - jwt_token = response.get("token") - project_id = response.get("project_id") + response: Final = as_str_mapping(self._fetch_auth_token(config.api_key, config.auth_endpoint)) + if response is not None: + jwt_token = response.get("token") + project_id = response.get("project_id") except Exception: pass @@ -78,12 +84,9 @@ class AgentOps(OpenTelemetry): "telemetry.sdk.name": "agentops", } - if project_id: - resource_attrs["project.id"] = project_id + self.resource_attributes = {**resource_attrs, "project.id": project_id} if project_id else resource_attrs - self.resource_attributes = resource_attrs - - def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> dict[str, Any]: + def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> object: """ Fetch JWT authentication token from AgentOps API @@ -111,6 +114,6 @@ class AgentOps(OpenTelemetry): if response.status_code != 200: raise Exception(f"Failed to fetch auth token: {response.text}") - return response.json() + return _DECODED_JSON.validate_python(response.json()) finally: client.close() diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 24279b443aa..5a03b91d470 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -166,7 +166,7 @@ class ArizeLogger(OpenTelemetry): primary ``otel`` callback which handles proxy-level parent spans. """ - def set_attributes(self, span: Span, kwargs, response_obj: Any | None): + def set_attributes(self, span: Span, kwargs, response_obj: object): ArizeLogger.set_arize_attributes(span, kwargs, response_obj) @staticmethod diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 501f5749ea4..f72b29a6ad8 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -3,9 +3,10 @@ import os from datetime import datetime -from typing import Final +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import verbose_logger @@ -24,6 +25,18 @@ from litellm.utils import print_verbose API_BASE: Final = "https://api.braintrustdata.com/v1" +class _BraintrustProjectBody(TypedDict): + id: ReadOnly[str] + + +class _BraintrustProjectResponse(Protocol): + def json(self) -> _BraintrustProjectBody: ... + + +def _project_id(response: _BraintrustProjectResponse) -> str: + return response.json()["id"] + + def get_utc_datetime(): import datetime as dt @@ -80,8 +93,7 @@ class BraintrustLogger(CustomLogger): headers=self.headers, json={"name": project_name}, ) - project_dict: Final = response.json() - project_id: Final = project_dict["id"] + project_id: Final = _project_id(response) self._project_id_cache[project_name] = project_id return project_id except httpx.HTTPStatusError as e: @@ -100,8 +112,7 @@ class BraintrustLogger(CustomLogger): headers=self.headers, json={"name": project_name}, ) - project_dict: Final = response.json() - project_id: Final = project_dict["id"] + project_id: Final = _project_id(response) self._project_id_cache[project_name] = project_id return project_id except httpx.HTTPStatusError as e: @@ -112,18 +123,14 @@ class BraintrustLogger(CustomLogger): f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) - project_dict: Final = project.json() - - self.default_project_id = project_dict["id"] + self.default_project_id = _project_id(project) def create_sync_default_project_and_experiment(self): project: Final = self.global_braintrust_sync_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) - project_dict: Final = project.json() - - self.default_project_id = project_dict["id"] + self.default_project_id = _project_id(project) def log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 1e0413e8576..d0b25ea83e8 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -19,8 +19,9 @@ """Output modules for writing CBF data to various destinations.""" import zoneinfo +from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any, Final +from typing import Final import httpx import polars as pl @@ -188,7 +189,7 @@ class CloudZeroStreamer: ) raise - def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> dict[str, Any]: + def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> Mapping[str, object]: """Prepare batch payload according to CloudZero AnyCost API format.""" # Convert batch_date to month for the API (YYYY-MM format) try: @@ -205,7 +206,7 @@ class CloudZeroStreamer: if record: data_records.append(record) - payload: Final = {"month": month_str, "operation": operation, "data": data_records} + payload: Final[Mapping[str, object]] = {"month": month_str, "operation": operation, "data": data_records} return payload diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1c12ac395c7..9e2dc902d55 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -352,7 +352,7 @@ class CustomGuardrail(CustomLogger): self, violation_message: str, request_data: dict[str, Any], - detection_info: dict[str, Any] | None = None, + detection_info: Mapping[str, object] | None = None, original_response: object = None, ) -> None: """ @@ -1368,7 +1368,10 @@ class CustomGuardrail(CustomLogger): if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): - guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) + guardrail_mode = GuardrailMode( + tags=copy.deepcopy(self.event_hook.tags), + default=copy.deepcopy(self.event_hook.default), + ) else: guardrail_mode = self.event_hook diff --git a/litellm/integrations/custom_secret_manager.py b/litellm/integrations/custom_secret_manager.py index e59842d409a..8fc5bc551c8 100644 --- a/litellm/integrations/custom_secret_manager.py +++ b/litellm/integrations/custom_secret_manager.py @@ -37,7 +37,6 @@ Usage: """ from abc import abstractmethod -from typing import Any import httpx @@ -154,7 +153,7 @@ class CustomSecretManager(BaseSecretManager): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Asynchronously write a secret to your custom secret manager. diff --git a/litellm/integrations/dotprompt/__init__.py b/litellm/integrations/dotprompt/__init__.py index 4cf6c0de2df..0a2921db6c7 100644 --- a/litellm/integrations/dotprompt/__init__.py +++ b/litellm/integrations/dotprompt/__init__.py @@ -61,7 +61,9 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom prompt_file: Final = getattr(litellm_params, "prompt_file", None) # Handle dotprompt_content from database - dotprompt_content: Final = getattr(litellm_params, "dotprompt_content", None) + dotprompt_content: Final = ( + litellm_params.dotprompt_content if hasattr(litellm_params, "dotprompt_content") else None + ) if dotprompt_content and not prompt_data and not prompt_file: prompt_data = get_prompt_data_from_dotprompt_content(dotprompt_content) diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py index c646dbf4e2e..77d70dd761c 100644 --- a/litellm/integrations/dotprompt/dotprompt_manager.py +++ b/litellm/integrations/dotprompt/dotprompt_manager.py @@ -7,6 +7,8 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from typing_extensions import ReadOnly, TypedDict + from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.types.llms.openai import AllMessageValues @@ -18,7 +20,12 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -from .prompt_manager import PromptManager, PromptTemplate +from .prompt_manager import PromptJsonEntry, PromptManager, PromptTemplate + + +class _PromptJsonData(TypedDict, total=False): + content: ReadOnly[str] + metadata: ReadOnly[dict[str, object]] class DotpromptManager(CustomPromptManagement): @@ -342,17 +349,17 @@ class DotpromptManager(CustomPromptManagement): if self._prompt_manager: self._prompt_manager.reload_prompts() - def add_prompt_from_json(self, prompt_id: str, json_data: dict[str, Any]) -> None: + def add_prompt_from_json(self, prompt_id: str, json_data: _PromptJsonData) -> None: """Add a prompt from JSON data.""" content: Final = json_data.get("content", "") metadata: Final = json_data.get("metadata", {}) self.prompt_manager.add_prompt(prompt_id, content, metadata) - def load_prompts_from_json(self, prompts_data: dict[str, dict[str, object]]) -> None: + def load_prompts_from_json(self, prompts_data: Mapping[str, object]) -> None: """Load multiple prompts from JSON data.""" self.prompt_manager.load_prompts_from_json_data(prompts_data) - def get_prompts_as_json(self) -> dict[str, dict[str, object]]: + def get_prompts_as_json(self) -> dict[str, PromptJsonEntry]: """Get all prompts in JSON format.""" return self.prompt_manager.get_all_prompts_as_json() diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index 0eda266740c..0e52ca0b62f 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -12,6 +12,8 @@ from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import as_str_mapping + class _PromptFileJson(TypedDict): """JSON form of a .prompt file: rendered template text plus its frontmatter.""" @@ -20,6 +22,11 @@ class _PromptFileJson(TypedDict): metadata: ReadOnly[NotRequired[dict[str, object]]] +class PromptJsonEntry(TypedDict): + content: ReadOnly[str] + metadata: ReadOnly[dict[str, object]] + + def strip_version_suffix(prompt_id: str) -> str | None: base, separator, version = prompt_id.rpartition(".v") if separator and base and version.isdigit(): @@ -358,20 +365,26 @@ class PromptManager: return f"---\n{frontmatter_yaml}---\n{content}" - def get_all_prompts_as_json(self) -> dict[str, dict[str, Any]]: + def get_all_prompts_as_json(self) -> dict[str, PromptJsonEntry]: """Get all loaded prompts in JSON format. Returns: Dictionary mapping prompt_id to prompt data """ - result: Final = {} - for prompt_id, template in self.prompts.items(): - result[prompt_id] = { + return { + prompt_id: { "content": template.content, "metadata": template.metadata, } - return result + for prompt_id, template in self.prompts.items() + } - def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, object]]) -> None: - """Load additional prompts from JSON data (merges with existing prompts).""" - self._load_prompts_from_json(prompt_data) + def load_prompts_from_json_data(self, prompt_data: Mapping[str, object]) -> None: + """Load additional prompts from JSON data (merges with existing prompts). Entries that are not + mappings are skipped, the same as a malformed entry inside the loader.""" + entries: Final = { + prompt_id: dict(entry) + for prompt_id, raw in prompt_data.items() + if (entry := as_str_mapping(raw)) is not None + } + self._load_prompts_from_json(entries) diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 0ecf04326b3..cf871013265 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Mapping from typing import Any, Final from .base import FocusDestination @@ -20,7 +21,7 @@ class FocusDestinationFactory: *, provider: str, prefix: str, - config: dict[str, Any] | None = None, + config: Mapping[str, object] | None = None, ) -> FocusDestination: """Return a destination implementation for the requested provider.""" provider_lower: Final = provider.lower() @@ -39,7 +40,7 @@ class FocusDestinationFactory: def _resolve_config( *, provider: str, - overrides: dict[str, object], + overrides: Mapping[str, object], ) -> dict[str, Any]: if provider == "s3": resolved = { diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 615aba72954..5c345630984 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable -from typing import Any, Final, SupportsFloat, SupportsIndex, SupportsInt, cast +from collections.abc import Callable, Mapping +from typing import Final, SupportsFloat, SupportsIndex, SupportsInt, cast import polars as pl @@ -24,7 +24,7 @@ class FocusExportEngine: provider: str, export_format: str, prefix: str, - destination_config: dict[str, Any] | None = None, + destination_config: Mapping[str, object] | None = None, ) -> None: self.provider = provider self.export_format = export_format @@ -77,7 +77,7 @@ class FocusExportEngine: return FocusParquetSerializer() raise NotImplementedError(f"Export format '{self.export_format}' not supported. Use 'parquet' or 'csv'.") - async def dry_run_export_usage_data(self, limit: int | None) -> dict[str, Any]: + async def dry_run_export_usage_data(self, limit: int | None) -> dict[str, object]: data: Final = await self.database.get_usage_data(limit=limit) normalized: Final = self.transformer.transform(data) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index e961b53a240..71b61296cfc 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -6,6 +6,8 @@ import os from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, ClassVar, Final, cast +from typing_extensions import TypedDict, Unpack + import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -25,6 +27,10 @@ FOCUS_USAGE_DATA_JOB_NAME: Final = "focus_export_usage_data" DEFAULT_DRY_RUN_LIMIT: Final = 500 +class _CustomLoggerOptions(TypedDict, total=False, extra_items=object): + pass + + class FocusLogger(CustomLogger): """Coordinates Focus export jobs across transformer/serializer/destination layers.""" @@ -40,7 +46,7 @@ class FocusLogger(CustomLogger): interval_seconds: int | None = None, prefix: str | None = None, destination_config: dict[str, Any] | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomLoggerOptions], ) -> None: super().__init__(**kwargs) self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower() diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 5bc7247e02a..99cd8b31b63 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -3,6 +3,7 @@ import hashlib import json import os import time +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final @@ -15,6 +16,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload @@ -134,7 +136,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): """ return f"{date_str}/batch-{batch_id}.ndjson" - def _get_config_key(self, kwargs: dict[str, Any]) -> str: + def _get_config_key(self, kwargs: Mapping[str, object]) -> str: """ Extract a synchronous grouping key from kwargs to group items by GCS config. This allows us to batch items with the same bucket/credentials together. @@ -143,7 +145,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): This key may contain sensitive information (bucket names, paths) - use _sanitize_config_key() for logging purposes. """ - standard_callback_dynamic_params: Final = kwargs.get("standard_callback_dynamic_params", None) or {} + standard_callback_dynamic_params: Final = ( + as_str_mapping(kwargs.get("standard_callback_dynamic_params", None)) or {} + ) bucket_name = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME or "default" path_service_account: Final = ( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 9dfd75e5559..3f8cf93e11c 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -9,6 +9,7 @@ Usage: """ import asyncio +import os from collections.abc import AsyncIterable, Iterable from typing import Final @@ -40,7 +41,7 @@ _mocks_initialized = False # Default mock latency in seconds (simulates network round-trip) # Typical GCS API calls take 100-300ms for uploads, 50-150ms for GET/DELETE -_MOCK_LATENCY_SECONDS: Final = float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 +_MOCK_LATENCY_SECONDS: Final = float(os.getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 async def _mock_async_handler_get(self, url, params=None, headers=None, follow_redirects=None): diff --git a/litellm/integrations/gcs_pubsub/pub_sub.py b/litellm/integrations/gcs_pubsub/pub_sub.py index c799245c215..cbf6898d978 100644 --- a/litellm/integrations/gcs_pubsub/pub_sub.py +++ b/litellm/integrations/gcs_pubsub/pub_sub.py @@ -12,6 +12,8 @@ import os import traceback from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter + from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: @@ -27,6 +29,8 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) +_DECODED_JSON: Final = TypeAdapter(object) + class GcsPubSubLogger(CustomBatchLogger): def __init__( @@ -152,7 +156,7 @@ class GcsPubSubLogger(CustomBatchLogger): finally: self.log_queue.clear() - async def publish_message(self, message: SpendLogsPayload | StandardLoggingPayload) -> dict[str, Any] | None: + async def publish_message(self, message: SpendLogsPayload | StandardLoggingPayload) -> object: """ Publish message to Google Cloud Pub/Sub using REST API @@ -188,7 +192,7 @@ class GcsPubSubLogger(CustomBatchLogger): raise Exception(f"Failed to publish message: {response.text}") verbose_logger.debug("Pub/Sub response: %s", response.text) - return response.json() + return _DECODED_JSON.validate_python(response.json()) except Exception as e: verbose_logger.error("Pub/Sub publish error: %s", str(e)) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 86f80047e52..b8c6e91d0d7 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -20,6 +20,7 @@ from litellm.integrations.langfuse.langfuse_mock_client import ( should_use_langfuse_mock, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, filter_exceptions_from_params, reconstruct_model_name, safe_deep_copy, @@ -1268,11 +1269,11 @@ def log_provider_specific_information_as_span( *, tracing: "LangfuseTracing", parent: "LangfuseObservation", - enrichments: Mapping[str, Any], + enrichments: Mapping[str, object], ): """Logs provider-specific information as spans under the generation.""" - _hidden_params: Final[Mapping[str, object] | None] = enrichments.get("hidden_params", None) + _hidden_params: Final = as_str_mapping(enrichments.get("hidden_params", None)) if _hidden_params is None: return diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index f195cd97e3c..9364a3d4216 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -12,6 +12,7 @@ from litellm.integrations.langfuse.langfuse_otel_attributes import ( serialize_langfuse_datetime, ) from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.types.integrations.langfuse_otel import ( LangfuseSpanAttributes, @@ -222,7 +223,7 @@ class LangfuseOtelLogger(OpenTelemetry): ) @staticmethod - def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj): + def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, object], response_obj): """ Sets Langfuse specific metadata attributes onto the OTEL span. @@ -234,7 +235,7 @@ class LangfuseOtelLogger(OpenTelemetry): from litellm.integrations.arize._utils import safe_set_attribute from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - dynamic_params: Final = kwargs.get("standard_callback_dynamic_params") + dynamic_params: Final = as_str_mapping(kwargs.get("standard_callback_dynamic_params")) langfuse_environment: Final = ( dynamic_params.get("langfuse_environment") if dynamic_params else None ) or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT") diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 3786087ba91..50ff7002068 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -2,6 +2,7 @@ Call Hook for LiteLLM Proxy which allows Langfuse prompt management. """ +from collections.abc import Mapping from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast @@ -110,6 +111,10 @@ def _remember_trace_id(litellm_call_id: object, logged: LangfuseLoggedEvent) -> in_memory_trace_id_cache.set_cache(litellm_call_id=litellm_call_id, service_name="langfuse", trace_id=trace_id) +def _params_without_model(config: Mapping[str, object]) -> Mapping[str, object]: + return {k: v for k, v in config.items() if k != "model"} + + class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogger): def __init__( self, @@ -182,12 +187,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge return compiled_prompt def _get_optional_params_from_langfuse(self, langfuse_prompt_client: PROMPT_CLIENT) -> dict: - config: Final = langfuse_prompt_client.config - optional_params: Final = {} - for k, v in config.items(): - if k != "model": - optional_params[k] = v - return optional_params + return dict(_params_without_model(langfuse_prompt_client.config)) async def async_get_chat_completion_prompt( self, diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 01812149498..ff1f3568685 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -7,10 +7,11 @@ import random import traceback import types from datetime import datetime, timezone -from typing import Any, Final +from typing import Final import httpx from pydantic import BaseModel, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -31,6 +32,10 @@ from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPa _RUN_RESPONSE: Final = TypeAdapter(object) +class _DynamicParamsKwargs(TypedDict, total=False): + standard_callback_dynamic_params: ReadOnly[StandardCallbackDynamicParams | None] + + def is_serializable(value): non_serializable_types: Final = ( types.CoroutineType, @@ -474,7 +479,7 @@ class LangsmithLogger(CustomBatchLogger): return log_queue_by_credentials - def _get_sampling_rate_to_use_for_request(self, kwargs: dict[str, Any]) -> float: + def _get_sampling_rate_to_use_for_request(self, kwargs: _DynamicParamsKwargs) -> float: standard_callback_dynamic_params: Final[StandardCallbackDynamicParams | None] = kwargs.get( "standard_callback_dynamic_params", None ) @@ -485,7 +490,7 @@ class LangsmithLogger(CustomBatchLogger): sampling_rate = float(_sampling_rate) return sampling_rate - def _get_credentials_to_use_for_request(self, kwargs: dict[str, Any]) -> LangsmithCredentialsObject: + def _get_credentials_to_use_for_request(self, kwargs: _DynamicParamsKwargs) -> LangsmithCredentialsObject: """ Handles key/team based logging diff --git a/litellm/integrations/langtrace.py b/litellm/integrations/langtrace.py index eeff76e7063..ea411baee7e 100644 --- a/litellm/integrations/langtrace.py +++ b/litellm/integrations/langtrace.py @@ -6,7 +6,7 @@ from litellm.types.integrations.otel_span_attributes import SpanAttributes if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py index 4d2b4edf3cd..6273f39e09c 100644 --- a/litellm/integrations/levo/levo.py +++ b/litellm/integrations/levo/levo.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: Protocol = _Protocol OpenTelemetryConfig = _OpenTelemetryConfig - Span = _Span | Any + Span = _Span else: Protocol = Any OpenTelemetryConfig = Any diff --git a/litellm/integrations/logfire_logger.py b/litellm/integrations/logfire_logger.py index 848369a1cbb..862893fa883 100644 --- a/litellm/integrations/logfire_logger.py +++ b/litellm/integrations/logfire_logger.py @@ -3,8 +3,9 @@ import os import traceback +from collections.abc import Mapping from enum import Enum -from typing import Any, Final, NamedTuple +from typing import Final, NamedTuple from typing_extensions import LiteralString @@ -15,7 +16,7 @@ from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info class SpanConfig(NamedTuple): message_template: LiteralString - span_data: dict[str, Any] + span_data: Mapping[str, object] class LogfireLevel(str, Enum): diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index 9829ef4e18f..3623335f599 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -49,11 +49,12 @@ import time import uuid from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Any, Final +from typing import Final, Protocol import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.newrelic import NewRelicInitParams @@ -65,6 +66,10 @@ except ImportError: _newrelic_agent = None +class _SupportsTimestamp(Protocol): + def timestamp(self) -> float: ... + + class NewRelicLogger(CustomLogger): """ New Relic logger for LiteLLM to send AI monitoring events. @@ -126,13 +131,13 @@ class NewRelicLogger(CustomLogger): verbose_logger.error("Failed to initialize New Relic agent: %s. Integration will be disabled.", e) self.enabled = False - def _get_newrelic_params(self) -> dict: + def _get_newrelic_params(self) -> Mapping[str, object]: """ Get the newrelic_params from litellm.newrelic_params These are params specific to initializing the NewRelicLogger e.g. turn_off_message_logging """ - dict_newrelic_params: dict = {} + dict_newrelic_params: Mapping[str, object] = {} if litellm.newrelic_params is not None: if isinstance(litellm.newrelic_params, NewRelicInitParams): dict_newrelic_params = litellm.newrelic_params.model_dump() @@ -401,9 +406,9 @@ class NewRelicLogger(CustomLogger): return choices[0].get("finish_reason") or "unknown" return "unknown" - def _to_epoch_ms(self, t: Any) -> float: + def _to_epoch_ms(self, t: datetime | float) -> float: """Convert a datetime or float timestamp to epoch milliseconds.""" - if hasattr(t, "timestamp"): + if has_attributes(t, _SupportsTimestamp, "timestamp"): return t.timestamp() * 1000.0 return float(t) * 1000.0 @@ -518,13 +523,13 @@ class NewRelicLogger(CustomLogger): sequence = 0 # Extract timestamps, preferring StandardLoggingPayload - start_time = None + start_time: datetime | float | None = None if standard_logging_object: start_time = standard_logging_object.get("startTime") if not start_time: start_time = kwargs.get("start_time") - end_time = None + end_time: datetime | float | None = None if standard_logging_object: end_time = standard_logging_object.get("endTime") if not end_time: diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index a2cedeba0f5..641d1d86188 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -254,7 +254,7 @@ class NewRelicMetricsLogger(CustomBatchLogger): if not self.log_queue: return if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1: - await asyncio.sleep(2**_pass) + await asyncio.sleep(1 << _pass) async with self.flush_lock: tried_ids: Final = frozenset(id(record) for record in attempted) survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids) diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index 3aaf5bfc162..2909ab9d63c 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -57,9 +57,7 @@ def build_span_payload( span_id: Final = utils.create_uuid7() model: Final[str] = response_obj.get("model", "unknown-model") - obj_type: Final[str] = response_obj.get("object", "unknown-object") - created: Final[int] = response_obj.get("created", 0) - span_name: Final = f"{model}_{obj_type}_{created}" + span_name: Final = f"{model}_{response_obj.get('object', 'unknown-object')}_{response_obj.get('created', 0)}" _logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id) diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index cca12b885c1..dd150082252 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -130,7 +130,8 @@ class _EnvWithoutBareExcludedServices(PydanticBaseSettingsSource): return self._env_settings.get_field_value(field, field_name) def __call__(self) -> dict[str, object]: - return {key: value for key, value in self._env_settings().items() if key != "excluded_services"} + settings: Final = self._env_settings() + return {key: settings[key] for key in settings if key != "excluded_services"} class OpenTelemetryV2Config(BaseSettings): diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index 5d37894bb8f..a69c334a1fb 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -332,7 +332,7 @@ def caller_session_id(kwargs: Mapping[str, object], trace: TraceControls) -> str return (None if echoes_trace_id else explicit) or trace.session_id or None -def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: +def time_to_first_chunk_seconds(kwargs: Mapping[str, object]) -> float | None: """Seconds from the upstream request being issued (``api_call_start_time``) to the first streamed chunk (``completion_start_time``); ``None`` for non-streaming calls, where ``completion_start_time`` is backfilled with the @@ -340,8 +340,8 @@ def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: optional_params: Final = cast(Mapping[str, object], kwargs.get("optional_params") or {}) if not optional_params.get("stream"): return None - api_call_start: Final = to_seconds(kwargs.get("api_call_start_time")) - completion_start: Final = to_seconds(kwargs.get("completion_start_time")) + api_call_start: Final = _epoch_seconds(kwargs.get("api_call_start_time")) + completion_start: Final = _epoch_seconds(kwargs.get("completion_start_time")) if api_call_start is None or completion_start is None: return None return completion_start - api_call_start diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d11505ce026..9c62c45e489 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, cast from urllib.parse import urlsplit from typing_extensions import ReadOnly, TypedDict @@ -716,6 +716,15 @@ def _json_or_none(value: object) -> str | None: return None +class _EnumValue(Protocol): + @property + def value(self) -> object: ... + + +def _enum_value(item: _EnumValue) -> object: + return item.value + + def _guardrail_mode_str(value: object) -> str | None: """Normalize ``guardrail_mode`` to a stable string. @@ -732,12 +741,12 @@ def _guardrail_mode_str(value: object) -> str | None: for item in value: if item is None: continue - part = as_str(item.value) if isinstance(item, Enum) else as_str(item) + part = as_str(_enum_value(item)) if isinstance(item, Enum) else as_str(item) if part: parts.append(part) return ",".join(parts) or None if isinstance(value, Enum): - return as_str(value.value) + return as_str(_enum_value(value)) return as_str(value) diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 7ec3711ee65..72d51471ba2 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -173,6 +173,7 @@ class _MetricKwargs(TypedDict, total=False): response_cost: ReadOnly[float | None] completion_start_time: ReadOnly[datetime | float | str | None] api_call_start_time: ReadOnly[datetime | float | str | None] + end_time: ReadOnly[datetime | float | str | None] def resolve_error_type(kwargs: Mapping[str, Any]) -> str: @@ -414,15 +415,12 @@ class GenAIMetricRecorder: if generation_time > 0: self._metrics.time_per_output_token.record(generation_time / completion_tokens, attributes=common_attrs) - def _record_response_duration(self, kwargs: Mapping[str, Any], end_time: datetime, common_attrs: dict) -> None: + def _record_response_duration(self, kwargs: _MetricKwargs, end_time: datetime, common_attrs: dict) -> None: api_call_start_time: Final = kwargs.get("api_call_start_time") if api_call_start_time is None: return - _end_time = kwargs.get("end_time") or end_time - if _end_time is None: - _end_time = datetime.now() api_call_start: Final = to_seconds(api_call_start_time) - end_ts: Final = to_seconds(_end_time) + end_ts: Final = to_seconds(kwargs.get("end_time") or end_time) if api_call_start is None or end_ts is None: return duration: Final = end_ts - api_call_start diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index d787375ca3c..83f85232182 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -13,6 +13,8 @@ import re import traceback from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -32,6 +34,7 @@ from litellm.types.utils import StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger +_DECODED_JSON: Final = TypeAdapter(object) _BASE64_INLINE_PATTERN: Final = re.compile( r"data:(?:application|image|audio|video)/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=\s]+", re.MULTILINE, @@ -270,7 +273,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): if self.sqs_queue_url is None: raise ValueError("sqs_queue_url not set") - json_data: Final = json.loads(safe_dumps(payload)) + json_data: Final = _DECODED_JSON.validate_python(json.loads(safe_dumps(payload))) if self.app_crypto: aad_bytes: Final = self.sqs_app_encryption_aad.encode("utf-8") if self.sqs_app_encryption_aad else None encrypted: Final = self.app_crypto.encrypt_json(json_data, aad=aad_bytes) diff --git a/litellm/interactions/litellm_responses_transformation/handler.py b/litellm/interactions/litellm_responses_transformation/handler.py index 8fee0fcd1b5..d84f0b5ffbd 100644 --- a/litellm/interactions/litellm_responses_transformation/handler.py +++ b/litellm/interactions/litellm_responses_transformation/handler.py @@ -116,10 +116,7 @@ class LiteLLMResponsesInteractionsHandler: """Async handler for interactions API requests.""" # Call litellm.aresponses() # Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] - aresponses_fn: Final[ - Callable[..., Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] - ] = vars(litellm)["aresponses"] - responses_response: Final = await aresponses_fn( + responses_response: Final = await litellm.aresponses( **responses_request, ) diff --git a/litellm/litellm_core_utils/app_crypto.py b/litellm/litellm_core_utils/app_crypto.py index 4ef3f34f3b8..da844620dce 100644 --- a/litellm/litellm_core_utils/app_crypto.py +++ b/litellm/litellm_core_utils/app_crypto.py @@ -4,6 +4,9 @@ import os from typing import Final from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from pydantic import TypeAdapter + +_DECODED_JSON: Final = TypeAdapter(object) class AppCrypto: @@ -12,7 +15,7 @@ class AppCrypto: raise ValueError("Master key must be 32 bytes for AES-256-GCM") self.key = master_key - def encrypt_json(self, data: dict, aad: bytes | None = None) -> dict: + def encrypt_json(self, data: object, aad: bytes | None = None) -> dict: aes: Final = AESGCM(self.key) nonce: Final = os.urandom(12) plaintext: Final = json.dumps(data).encode("utf-8") @@ -24,10 +27,10 @@ class AppCrypto: "tag": base64.b64encode(tag).decode(), } - def decrypt_json(self, enc: dict, aad: bytes | None = None) -> dict: + def decrypt_json(self, enc: dict, aad: bytes | None = None) -> object: aes: Final = AESGCM(self.key) nonce: Final = base64.b64decode(enc["nonce"]) ct: Final = base64.b64decode(enc["ciphertext"]) tag: Final = base64.b64decode(enc["tag"]) data: Final = aes.decrypt(nonce, ct + tag, aad) - return json.loads(data.decode()) + return _DECODED_JSON.validate_python(json.loads(data.decode())) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index e9689bd7d8a..8e949694881 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -3,10 +3,11 @@ import re from collections.abc import Mapping, Sequence from itertools import accumulate, repeat from types import MappingProxyType -from typing import Any, Final, cast +from typing import Final from urllib.parse import quote, unquote from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping VERTEX_AI_MANAGED_GCS_PREFIX: Final = "litellm-vertex-files/" BEDROCK_MANAGED_S3_BATCH_PREFIX: Final = "litellm-bedrock-files-" @@ -141,7 +142,8 @@ def should_allow_legacy_cloud_file_ids( if isinstance(litellm_params, Mapping): trusted_model_credentials: Final = litellm_params.get("_litellm_internal_model_credentials") if isinstance(trusted_model_credentials, _MAPPING_PROXY_TYPE): - value = cast(Mapping[str, Any], trusted_model_credentials).get("allow_legacy_cloud_file_ids") + trusted_mapping: Final = as_str_mapping(trusted_model_credentials) + value = trusted_mapping.get("allow_legacy_cloud_file_ids") if trusted_mapping is not None else None if isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 1d61ef8519a..4be5cabac04 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -5,10 +5,11 @@ import logging import re from collections.abc import Collection, Iterable, Iterator, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable import httpx from pydantic import TypeAdapter, ValidationError +from typing_extensions import TypeIs # noqa: TID251 # narrows a duck-typed value to the Protocol its caller reads import litellm from litellm._logging import verbose_logger @@ -58,6 +59,13 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def decoded_json_object(value: object) -> Mapping[str, object]: + mapping: Final = as_str_mapping(value) + if mapping is not None: + return mapping + raise AttributeError(f"{type(value).__name__!r} object has no attribute 'get'", name="get", obj=value) + + @runtime_checkable class _SupportsIter(Protocol): def __iter__(self) -> Iterator[object]: ... @@ -72,6 +80,24 @@ def attribute_or_none(source: object, name: str) -> object: return getattr(source, name, None) +_Shape: Final = TypeVar("_Shape") + + +def has_attributes( + value: object, + shape: type[_Shape], + *names: str, +) -> TypeIs[_Shape]: # guard-ok: hasattr sees what __getattr__ serves, which Protocol isinstance skips on 3.12+ + return all(hasattr(value, name) for name in names) + + +def require_attributes(value: object, shape: type[_Shape], *names: str) -> _Shape: + if has_attributes(value, shape, *names): + return value + missing: Final = next(name for name in names if not hasattr(value, name)) + raise AttributeError(f"'{type(value).__name__}' object has no attribute '{missing}'", name=missing, obj=value) + + def iter_items(value: object) -> Iterable[object]: if isinstance(value, _SupportsIter): return value @@ -634,11 +660,14 @@ def preserve_upstream_non_openai_attributes( for key in candidate_keys ) if all(value is None or isinstance(value, (str, int, float, bool)) for _, value in candidate_values): - for key, value in candidate_values: - setattr(model_response, key, value) + _assign_attributes(model_response, candidate_values) return - for key, value in original_chunk.model_dump(include=frozenset(candidate_keys)).items(): - setattr(model_response, key, value) + _assign_attributes(model_response, original_chunk.model_dump(include=frozenset(candidate_keys)).items()) + + +def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) -> None: + for key, value in values: + setattr(target, key, value) def safe_deep_copy(data): @@ -764,6 +793,10 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: Filtered data structure with Exception and callable objects removed, or None if the entire input was an Exception or callable """ + return _filter_exceptions_from_params(data, max_depth) + + +def _filter_exceptions_from_params(data: object, max_depth: int) -> object: if max_depth <= 0: return data @@ -785,7 +818,7 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: if isinstance(v, Exception) or (callable(v) and not isinstance(v, type)): continue try: - filtered = filter_exceptions_from_params(v, max_depth - 1) + filtered = _filter_exceptions_from_params(v, max_depth - 1) if filtered is not None: result[k] = filtered except Exception: @@ -799,7 +832,7 @@ def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)): continue try: - filtered = filter_exceptions_from_params(item, max_depth - 1) + filtered = _filter_exceptions_from_params(item, max_depth - 1) if filtered is not None: result_list.append(filtered) except Exception: diff --git a/litellm/litellm_core_utils/cost_discount.py b/litellm/litellm_core_utils/cost_discount.py index a64ad0b20f7..bfcee7928fb 100644 --- a/litellm/litellm_core_utils/cost_discount.py +++ b/litellm/litellm_core_utils/cost_discount.py @@ -23,7 +23,7 @@ def parse_cost_discount_key(key: str) -> CostDiscountKey: def _literal_length(pattern: str) -> int: - return sum(1 for token in _GLOB_TOKEN.findall(pattern) if len(token) == 1 and token not in "*?") + return sum(1 for match in _GLOB_TOKEN.finditer(pattern) if len(match.group()) == 1 and match.group() not in "*?") def _model_names(model: str, provider_prefix: str) -> tuple[str, ...]: diff --git a/litellm/litellm_core_utils/default_encoding.py b/litellm/litellm_core_utils/default_encoding.py index c3b6a008411..970dd512a40 100644 --- a/litellm/litellm_core_utils/default_encoding.py +++ b/litellm/litellm_core_utils/default_encoding.py @@ -48,5 +48,5 @@ for attempt in range(_max_retries): # Last attempt, re-raise the exception raise # Exponential backoff with jitter to reduce collision probability - delay = _retry_delay * (2**attempt) + random.uniform(0, 0.1) + delay = _retry_delay * (2.0**attempt) + random.uniform(0, 0.1) time.sleep(delay) diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 1dac67ecbf6..09720fa413c 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -181,7 +181,7 @@ def _delete_nested_value_custom( def delete_nested_value( - data: dict[str, Any], + data: dict[str, object], path: str, depth: int = 0, max_depth: int = 20, diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index e8da59cd33a..1c816b22ef1 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -16,7 +16,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, should_report_bug, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.secret_redaction import redact_string as redact_secret_string from litellm.types.utils import LlmProviders @@ -169,7 +169,7 @@ def get_error_message(error_obj) -> str | None: return None if hasattr(error_obj, "body"): - _error_obj_body: Final = getattr(error_obj, "body") + _error_obj_body: Final = attribute_or_none(error_obj, "body") if isinstance(_error_obj_body, dict): # OpenAI-style: {"message": "...", "type": "...", ...} if _error_obj_body.get("message"): diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 1e2a511d439..7fc6eab6c60 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -43,6 +43,7 @@ from litellm.litellm_core_utils.offline_mode import is_offline_mode FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) +_DECODED_JSON: Final = TypeAdapter(object) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) @@ -314,7 +315,7 @@ def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) - def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _FetchAttemptOutcome: try: - parsed: Final[object] = json.loads(content) + parsed: Final = _DECODED_JSON.validate_python(json.loads(content)) except ValueError as e: return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}") if not isinstance(parsed, dict): @@ -383,7 +384,8 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp ) if response.is_error: return ModelCostMapReloadUnavailable(reason=f"HTTP {response.status_code} from {url}") - return _classify_fetch_content(response.content, url, response.headers.get("etag")) + etag_values: Final = response.headers.get_list("etag") + return _classify_fetch_content(response.content, url, ", ".join(etag_values) if etag_values else None) def _next_retry_wait( diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index f8d44d44cdc..a74337d2549 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -1,4 +1,5 @@ from collections.abc import Iterator, Mapping, MutableMapping +from types import MappingProxyType from typing import Final, Protocol, cast from pydantic import BaseModel @@ -52,7 +53,7 @@ class HiddenParamsModelView(MutableMapping[str, object]): def stored_attribute(obj: object, name: str) -> object: if not isinstance(obj, BaseModel): return getattr(obj, name, None) - values: Final[Mapping[str, object]] = vars(obj) + values: Final[Mapping[str, object]] = MappingProxyType[str, object](vars(obj)) if name in values: return values[name] extras: Final[Mapping[str, object]] = obj.__pydantic_extra__ or {} diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 5a533054f0b..c3769686e45 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -28,7 +28,7 @@ def inherit_message_logging_privacy(disabled: bool) -> Generator[None]: def iter_client_callback_metadata_dicts( - kwargs: dict[str, Any], + kwargs: Mapping[str, object], ) -> Iterator[tuple[str, dict[str, Any]]]: litellm_params: Final = kwargs.get("litellm_params") if isinstance(litellm_params, dict): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 41f13817d75..fc8d1ecc8c1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -556,6 +556,7 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) +_REDACTED_RESULT: Final = TypeAdapter(object) def _provider_response_id(source: object) -> str | None: @@ -1171,7 +1172,7 @@ class Logging(LiteLLMLoggingBaseClass): automatically extracts metadata/litellm_metadata from kwargs, so callers don't need to manually plumb them into litellm_params. """ - base_litellm_params: Final[dict[str, Any]] = {} + base_litellm_params: Final[dict[str, object]] = {} if isinstance(kwargs.get("metadata"), dict): base_litellm_params["metadata"] = kwargs["metadata"].copy() @@ -1188,10 +1189,11 @@ class Logging(LiteLLMLoggingBaseClass): lp_metadata: Final = litellm_params.pop("metadata", None) base_litellm_params.update(litellm_params) if lp_metadata and isinstance(lp_metadata, dict): - base_litellm_params.setdefault("metadata", {}) - for k, v in lp_metadata.items(): - if k not in base_litellm_params["metadata"]: - base_litellm_params["metadata"][k] = v + merged_metadata: Final = base_litellm_params.setdefault("metadata", {}) + if isinstance(merged_metadata, dict): + for k, v in lp_metadata.items(): + if k not in merged_metadata: + merged_metadata[k] = v self.update_environment_variables( litellm_params=base_litellm_params, @@ -1793,7 +1795,7 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.exception( "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - original_response = redact_message_input_output_from_logging( + redact_message_input_output_from_logging( model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=original_response, ) @@ -2023,10 +2025,9 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details.pop("response_cost_failure_debug_information", None) - if cache_hit is None: - cache_hit = self.model_call_details.get("cache_hit", False) + resolved_cache_hit: Final = self.model_call_details.get("cache_hit", False) if cache_hit is None else cache_hit - if cache_hit is True: + if resolved_cache_hit is True: return 0.0 if is_unbilled_non_inference_call( @@ -2078,7 +2079,7 @@ class Logging(LiteLLMLoggingBaseClass): model_value: Final = litellm_model_name or self.model cost_cache_hit_value: Final = ( - self.model_call_details.get("cache_hit", False) if cache_hit is None else cache_hit + self.model_call_details.get("cache_hit", False) if resolved_cache_hit is None else resolved_cache_hit ) cost_cache_hit: Final[bool | None] = cast( # cast-ok: callback metadata is caller-provided bool | None, cost_cache_hit_value @@ -2512,8 +2513,10 @@ class Logging(LiteLLMLoggingBaseClass): def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" - usage: Final = self.model_call_details.pop("combined_usage_object", None) - response_cost: Final = self.model_call_details.pop("response_cost", None) + usage: Final = self._model_call_details_view.get("combined_usage_object") + response_cost: Final = self._model_call_details_view.get("response_cost") + self.model_call_details.pop("combined_usage_object", None) + self.model_call_details.pop("response_cost", None) if isinstance(usage, Usage): self.model_call_details["current_attempt_api_call_start_time"] = None self.model_call_details["failed_attempts_usage"] = ( @@ -3302,7 +3305,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback_logging_payload is not None or shared_logging_payload is None else shared_logging_payload.copy() ) - call_details: dict[str, Any] # mutable-ok: legacy loggers take a plain dict + call_details: dict[str, object] # mutable-ok: legacy loggers take a plain dict call_details = ( # rebind-ok: logging hooks may return a new dict self.model_call_details if not for_async_request @@ -3323,9 +3326,11 @@ class Logging(LiteLLMLoggingBaseClass): ) ## REDACT MESSAGES ## - result = redact_message_input_output_from_logging( - model_call_details=call_details, - result=result, + result = _REDACTED_RESULT.validate_python( + redact_message_input_output_from_logging( + model_call_details=call_details, + result=result, + ) ) ## LOGGING HOOK ## hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) @@ -3591,7 +3596,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") s3Logger.log_event( kwargs=call_details, - response_obj=call_details["complete_streaming_response"], + response_obj=call_details_view["complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, @@ -3914,9 +3919,11 @@ class Logging(LiteLLMLoggingBaseClass): self._get_completed_success_callbacks(asynchronous=True) if not self.stream else callbacks ) - result = redact_message_input_output_from_logging( - model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), - result=result, + result = _REDACTED_RESULT.validate_python( + redact_message_input_output_from_logging( + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), + result=result, + ) ) ## LOGGING HOOK ## @@ -4261,9 +4268,11 @@ class Logging(LiteLLMLoggingBaseClass): result: object = None # result sent to all loggers, init this to None incase it's not created - result = redact_message_input_output_from_logging( - model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), - result=result, + result = _REDACTED_RESULT.validate_python( + redact_message_input_output_from_logging( + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), + result=result, + ) ) self.mark_logging_complete(event_type="sync_failure") for callback in callbacks: @@ -4952,7 +4961,7 @@ class Logging(LiteLLMLoggingBaseClass): """ import httpx - httpx_response: Final = self.model_call_details.get("httpx_response", None) + httpx_response: Final[Response | None] = self.model_call_details.get("httpx_response", None) if httpx_response is None: raise ValueError("Google GenAI Generate Content: httpx_response is None") dict_result: Final = httpx_response.json() @@ -4982,7 +4991,7 @@ class Logging(LiteLLMLoggingBaseClass): The response object with usage added if available """ # Get usage from model_call_details (set by asend_message) - usage: Final = self.model_call_details.get("usage") + usage: Final[Usage | None] = self.model_call_details.get("usage") if usage is None: return result @@ -5815,7 +5824,9 @@ def build_custom_logger_compatible_class( return callback # Get global BitBucket config - bitbucket_config: Final = getattr(litellm, "global_bitbucket_config", None) + bitbucket_config: Final = ( + litellm.global_bitbucket_config if hasattr(litellm, "global_bitbucket_config") else None + ) if bitbucket_config is None: raise ValueError("BitBucket configuration not found. Please set litellm.global_bitbucket_config first.") diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 04510c20db0..045869e4547 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -450,7 +450,7 @@ class StandardBuiltInToolCostTracking: def _get_code_interpreter_cost( model_info: ModelInfo | None, custom_llm_provider: str | None, - standard_built_in_tools_params: StandardBuiltInToolsParams, + standard_built_in_tools_params: Mapping[str, object], ) -> float: """Calculate code interpreter cost.""" code_interpreter_sessions: Final = standard_built_in_tools_params.get("code_interpreter_sessions", None) diff --git a/litellm/litellm_core_utils/llm_judge.py b/litellm/litellm_core_utils/llm_judge.py index 0c2d7767455..307e000119b 100644 --- a/litellm/litellm_core_utils/llm_judge.py +++ b/litellm/litellm_core_utils/llm_judge.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: JSON_FENCE_RE: Final = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE) _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) def default_router_provider() -> Router | None: @@ -40,13 +41,13 @@ def parse_json_verdict(raw: str) -> dict[str, object]: # mutable-ok: plain pars text = fenced.group(1).strip() # rebind-ok: progressively narrowed to the JSON payload parsed: object try: - parsed = json.loads(text) + parsed = _DECODED_JSON.validate_python(json.loads(text)) except json.JSONDecodeError: start: Final = text.find("{") end: Final = text.rfind("}") if start == -1 or end <= start: raise - parsed = json.loads(text[start : end + 1]) + parsed = _DECODED_JSON.validate_python(json.loads(text[start : end + 1])) if not isinstance(parsed, dict): raise ValueError("judge response is not a JSON object") return {str(k): v for k, v in parsed.items()} diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 307353d081b..56c6436eae6 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -1,5 +1,5 @@ import threading -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable, Iterable, Iterator, Mapping from types import MethodType from typing import TYPE_CHECKING, Final @@ -20,6 +20,10 @@ _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} _callback_lists_lock: Final = threading.RLock() +def _instance_attributes(target: object) -> Mapping[str, object]: + return vars(target) + + class LoggingCallbackManager: """ A centralized class that allows easy add / remove callbacks for litellm. @@ -341,7 +345,7 @@ class LoggingCallbackManager: key_parts: Final = [type(custom_logger).__name__] # Add only fundamental type instance variables to the key - for attr_name, attr_value in vars(custom_logger).items(): + for attr_name, attr_value in _instance_attributes(custom_logger).items(): if not attr_name.startswith("_"): # Skip private attributes if isinstance(attr_value, (str, bool, int)): key_parts.append(f"{attr_name}={attr_value}") diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 60be4d44706..59882dcbf2f 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -20,6 +20,7 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, @@ -1527,7 +1528,7 @@ def _grammar_branches(node: Mapping[str, object]) -> int: if properties is None: return 0 optional: Final = len(_optional_property_names(node)) - return 2**optional * (len(properties) - optional + 1) + return (1 << optional) * (len(properties) - optional + 1) def _optional_property_names(node: Mapping[str, object]) -> frozenset[str]: @@ -2839,7 +2840,7 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: break try: - obj, end_idx = decoder.raw_decode(raw, idx) + decoded = decoder.raw_decode(raw, idx) except json.JSONDecodeError as e: verbose_logger.warning( "split_concatenated_json_objects: discarding unparseable tool-call " @@ -2849,13 +2850,14 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: e, ) break - if isinstance(obj, dict): - results.append(obj) + obj = as_str_mapping(_DECODED_JSON.validate_python(decoded[0])) + if obj is not None: + results.append(dict(obj)) else: # Non-dict JSON value – wrap in empty dict (Bedrock requires # toolUse.input to be an object). results.append({}) - idx = end_idx + idx = decoded[1] return results @@ -2887,14 +2889,16 @@ def salvage_concatenated_tool_arguments(raw: str) -> tuple[dict[str, object], .. if idx >= length: break try: - obj, end_idx = decoder.raw_decode(stripped, idx) + decoded = decoder.raw_decode(stripped, idx) except json.JSONDecodeError: return () - if not isinstance(obj, dict): + decoded_object = as_str_mapping(_DECODED_JSON.validate_python(decoded[0])) + if decoded_object is None: return () - idx = end_idx + idx = decoded[1] if exceeded: continue + obj = dict(decoded_object) count += 1 if not kept: kept = (obj,) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3d865bd1819..cfb6e3ab45f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -21,6 +21,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -1021,11 +1022,12 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str: tool_function = get_attribute_or_key(tool, "function") tool_name = get_attribute_or_key(tool_function, "name") tool_arguments = get_attribute_or_key(tool_function, "arguments") - parsed_args = parse_tool_call_arguments( - tool_arguments, tool_name=tool_name, context="Anthropic XML tool invoke" + parsed_args = _DECODED_JSON.validate_python( + parse_tool_call_arguments(tool_arguments, tool_name=tool_name, context="Anthropic XML tool invoke") ) - if isinstance(parsed_args, dict): - parameters = "".join(f"<{param}>{val}\n" for param, val in parsed_args.items()) + parsed_mapping = as_str_mapping(parsed_args) + if parsed_mapping is not None: + parameters = "".join(f"<{param}>{val}\n" for param, val in parsed_mapping.items()) else: parameters = f"{parsed_args}\n" invokes += f"\n{tool_name}\n\n{parameters}\n\n" @@ -2342,7 +2344,7 @@ def _add_missing_tool_results( else: tool_name = getattr(function, "name", "unknown_tool") else: - function = getattr(tool_call, "function", None) + function = attribute_or_none(tool_call, "function") if function: tool_name = getattr(function, "name", "unknown_tool") break @@ -3184,10 +3186,14 @@ def anthropic_messages_pt( def extract_between_tags(tag: str, string: str, strip: bool = False) -> list[str]: ext_list = re.findall(f"<{tag}>(.+?)", string, re.DOTALL) if strip: - ext_list = [e.strip() for e in ext_list] + ext_list = list(_stripped(ext_list)) return ext_list +def _stripped(values: Iterable[str]) -> Iterator[str]: + return (value.strip() for value in values) + + def contains_tag(tag: str, string: str) -> bool: return bool(re.search(f"<{tag}>(.+?)", string, re.DOTALL)) @@ -4061,13 +4067,15 @@ def _convert_to_bedrock_tool_call_invoke( arguments_dict = {} else: try: - arguments_dict = json.loads(arguments) + decoded_arguments = as_str_mapping(_DECODED_JSON.validate_python(json.loads(arguments))) # Ensure arguments_dict is always a dict # (Bedrock requires toolUse.input to be an object). # Some providers return arguments: '""' which # json.loads decodes to a bare string. - if not isinstance(arguments_dict, dict): + if decoded_arguments is None: arguments_dict = {} + else: + arguments_dict = dict(decoded_arguments) except json.JSONDecodeError: # The model may return multiple JSON objects # concatenated in a single arguments string, e.g. @@ -6028,7 +6036,11 @@ def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str) return ({},) normalized_raw: Final = "{}" if raw == REDACTED_BY_LITELLM else raw try: - parsed: Final = parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context) + parsed: Final = as_str_mapping( + _DECODED_JSON.validate_python( + parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context) + ) + ) except ValueError as e: salvaged: Final = salvage_concatenated_tool_arguments(normalized_raw) if salvaged: @@ -6041,7 +6053,7 @@ def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str) return salvaged verbose_logger.warning("Failed to parse tool call arguments: %s", e) return ({},) - return (parsed,) if isinstance(parsed, dict) else ({},) + return (dict(parsed),) if parsed is not None else ({},) def _choice_tool_calls(choice: object) -> tuple[object, ...]: diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index ecf419695d0..eb3400aee8b 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -482,7 +482,7 @@ def _runs_by_anchor( ) -> Mapping[int, tuple[tuple[int, ...], ...]]: anchored: Final = tuple(sorted(((anchor, run) for run, anchor in anchors if anchor is not None), key=itemgetter(0))) return MappingProxyType( - {anchor: tuple(run for _, run in group) for anchor, group in groupby(anchored, key=itemgetter(0))} + {anchor: tuple(run for _, run in group) for anchor, group in groupby(anchored, key=lambda entry: entry[0])} ) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 52758e617d8..8c402fe9b44 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import ItemsView, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, field from typing import Any, Final @@ -30,6 +30,10 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset( ) +def _object_items(data: Mapping[str, object]) -> ItemsView[str, object]: + return data.items() + + class SensitiveDataMasker: def __init__( self, @@ -127,7 +131,7 @@ class SensitiveDataMasker: return data masked_data: Final[dict[str, object]] = {} - for k, v in data.items(): + for k, v in _object_items(data): try: key_is_sensitive = self.is_sensitive_key(k, excluded_keys) if isinstance(v, Mapping): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 39f675099d7..50789406055 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -425,7 +425,7 @@ class ChunkProcessor: if hasattr(chunk["choices"][0], "finish_reason"): chunk_finish_reason = chunk["choices"][0].finish_reason elif "finish_reason" in chunk["choices"][0]: - chunk_finish_reason = chunk["choices"][0]["finish_reason"] + chunk_finish_reason = _CHOICE_FIELD_VALUE.validate_python(chunk["choices"][0]["finish_reason"]) if chunk_finish_reason is not None: finish_reason = chunk_finish_reason @@ -1226,6 +1226,7 @@ def concatenate_base64_list(base64_strings: Sequence[str]) -> str: _CHOICE_DICTS: Final = TypeAdapter(list[dict[str, object]]) +_CHOICE_FIELD_VALUE: Final = TypeAdapter(object) def _text_completion_chunk_from_mapping(chunk: Mapping[str, object]) -> TextCompletionResponse: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f994baa52b1..fdcdca151be 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -37,7 +37,6 @@ from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CacheCreationTokenDetails, @@ -111,14 +110,15 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) -def _holds_only_scalars(model: BaseModel) -> bool: - fields: Final[Iterator[tuple[str, object]]] = iter(model) +def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) -def _delta_from_upstream(upstream_delta: BaseModel | None) -> Delta | None: +def _delta_from_upstream(upstream_delta: "BaseModel | _UpstreamDeltaLike | None") -> Delta | None: if upstream_delta is None: return Delta() + if not isinstance(upstream_delta, BaseModel): + return None try: return Delta(**dict(iter(upstream_delta))) except Exception: @@ -189,6 +189,33 @@ class _DeltaDumpDict(TypedDict): tool_calls: NotRequired[Sequence[Mapping[str, object]]] +class _UpstreamFunctionCallLike(Protocol): + arguments: str | None + + +class _UpstreamToolCallLike(Protocol): + function: _UpstreamFunctionCallLike + + +class _UpstreamDeltaLike(Protocol): + function_call: _UpstreamFunctionCallLike | None + tool_calls: Sequence[_UpstreamToolCallLike] | None + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + def model_dump(self) -> _DeltaDumpDict: ... + + +class _UpstreamChoiceLike(Protocol): + delta: _UpstreamDeltaLike | None + + +class _UpstreamChunkLike(Protocol): + id: str + choices: Sequence[_UpstreamChoiceLike] + system_fingerprint: str | None + + class _TextCompletionChoiceLike(Protocol): text: str finish_reason: str | None @@ -920,7 +947,7 @@ class CustomStreamWrapper: def copy_model_response_level_provider_specific_fields( self, - original_chunk: ModelResponseStream | OpenAIChatCompletionChunk, + original_chunk: object, model_response: ModelResponseStream, ) -> ModelResponseStream: """ @@ -949,8 +976,8 @@ class CustomStreamWrapper: or ("function_call" in completion_obj and completion_obj["function_call"] is not None) or ( "tool_calls" in model_response.choices[0].delta - and model_response.choices[0].delta["tool_calls"] is not None - and len(model_response.choices[0].delta["tool_calls"]) > 0 + and model_response.choices[0].delta.tool_calls is not None + and len(model_response.choices[0].delta.tool_calls) > 0 ) or ( "function_call" in model_response.choices[0].delta @@ -1115,7 +1142,8 @@ class CustomStreamWrapper: ## check if openai/azure chunk original_chunk: Final = response_obj.get("original_chunk", None) if original_chunk: - if len(original_chunk.choices) > 0: + upstream_chunk: Final[_UpstreamChunkLike] = original_chunk + if len(upstream_chunk.choices) > 0: original_chunk_for_reuse: Final[object] = cast( # cast-ok: runtime type-check before access object, original_chunk, @@ -1151,7 +1179,7 @@ class CustomStreamWrapper: ) else: choices: Final = [] - for choice in original_chunk.choices: + for choice in upstream_chunk.choices: try: if isinstance(choice, BaseModel): choice_json = choice.model_dump() @@ -1164,11 +1192,11 @@ class CustomStreamWrapper: setattr(model_response, "choices", choices) else: return - model_response.system_fingerprint = original_chunk.system_fingerprint + model_response.system_fingerprint = upstream_chunk.system_fingerprint setattr( model_response, "citations", - stored_attribute(original_chunk, "citations"), + stored_attribute(upstream_chunk, "citations"), ) preserve_upstream_non_openai_attributes( model_response=model_response, @@ -1611,7 +1639,9 @@ class CustomStreamWrapper: model_response.model = self.model ## FUNCTION CALL PARSING - original_chunk: Final = response_obj.get("original_chunk") if response_obj is not None else None + original_chunk: Final[_UpstreamChunkLike | None] = ( + response_obj.get("original_chunk") if response_obj is not None else None + ) if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints @@ -1623,31 +1653,19 @@ class CustomStreamWrapper: original_chunk, model_response ) if original_chunk.choices and len(original_chunk.choices) > 0: - delta = original_chunk.choices[0].delta + delta: Final = original_chunk.choices[0].delta if delta is not None and (delta.function_call is not None or delta.tool_calls is not None): try: model_response.system_fingerprint = original_chunk.system_fingerprint ## AZURE - check if arguments is not None - if original_chunk.choices[0].delta.function_call is not None: - if ( - getattr( - original_chunk.choices[0].delta.function_call, - "arguments", - ) - is None - ): - original_chunk.choices[0].delta.function_call.arguments = "" - elif original_chunk.choices[0].delta.tool_calls is not None: - if isinstance(original_chunk.choices[0].delta.tool_calls, list): - for t in original_chunk.choices[0].delta.tool_calls: - if hasattr(t, "functions") and hasattr(t.functions, "arguments"): - if ( - getattr( - t.function, - "arguments", - ) - is None - ): + if delta.function_call is not None: + if delta.function_call.arguments is None: + delta.function_call.arguments = "" + elif delta.tool_calls is not None: + if isinstance(delta.tool_calls, list): + for t in delta.tool_calls: + if hasattr(t, "functions") and hasattr(getattr(t, "functions"), "arguments"): + if t.function.arguments is None: t.function.arguments = "" _json_delta: Final[_DeltaDumpDict] = delta.model_dump() if "role" not in _json_delta or _json_delta["role"] is None: @@ -1886,7 +1904,7 @@ class CustomStreamWrapper: else: asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit)) ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler - litellm_params: Final = self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: Final[Mapping[str, object]] = self.logging_obj.model_call_details.get("litellm_params", {}) if self.logging_obj.is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 8aded69aed2..90e79d4d853 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -44,7 +44,7 @@ from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEnc if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace from litellm.litellm_core_utils.url_utils import safe_get -from litellm.llms.custom_httpx.http_handler import get_httpx_client +from litellm.llms.custom_httpx.http_handler import get_httpx_client, header_value from litellm.rust_bridge.tokenizer import get_encoding from litellm.types.llms.anthropic import ( AnthropicContentParamSource, @@ -247,7 +247,7 @@ def get_image_dimensions( client: Final = get_httpx_client() response: Final[httpx.Response] = safe_get(client, data) max_bytes: Final = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) - content_length: Final[str | None] = response.headers.get("Content-Length") + content_length: Final[str | None] = header_value(response.headers, "Content-Length") if content_length is not None and int(content_length) > max_bytes: pass # skip download; img_data stays None else: diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 4a1eaa87d3a..0db6a41b503 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -22,11 +22,11 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): import asyncio import socket from ipaddress import ip_address, ip_network -from typing import Any, Final, Protocol +from typing import Final, Protocol from urllib.parse import quote, urlparse, urlsplit, urlunparse import httpx -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict, Unpack import litellm @@ -50,6 +50,7 @@ class _UrlFetcher(Protocol): *, headers: dict[str, str] | None = None, follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, ) -> httpx.Response: ... @@ -62,6 +63,8 @@ class _AsyncUrlFetcher(Protocol): *, headers: dict[str, str] | None = None, follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ) -> httpx.Response: ... @@ -73,8 +76,14 @@ class _AsyncFetcherView(TypedDict): fetcher: ReadOnly[_AsyncUrlFetcher] -class _CallerHeadersView(TypedDict): +class _FetchOptions(TypedDict, total=False): headers: ReadOnly[dict[str, str]] + follow_redirects: ReadOnly[bool] + timeout: ReadOnly[float | httpx.Timeout | None] + + +class _AsyncFetchOptions(_FetchOptions, total=False): + max_response_bytes: ReadOnly[int | None] # Cloud-internal IPs that ``ip.is_global`` can report as public. Everything @@ -440,7 +449,7 @@ def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: +def safe_get(client: _UrlFetcher, url: str, **kwargs: Unpack[_FetchOptions]) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -460,21 +469,23 @@ def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: The final httpx.Response. """ if not getattr(litellm, "user_url_validation", True): - kwargs.setdefault("follow_redirects", True) - unvalidated: Final[_ResponseView] = {"response": client.get(url, **kwargs)} + unvalidated_options: Final[_FetchOptions] = { + **kwargs, + "follow_redirects": kwargs.get("follow_redirects", True), + } + unvalidated: Final[_ResponseView] = {"response": client.get(url, **unvalidated_options)} return unvalidated["response"] fetcher_view: Final[_FetcherView] = {"fetcher": client} fetcher: Final = fetcher_view["fetcher"] - kwargs.pop("follow_redirects", None) - headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} + caller_headers: Final = kwargs.get("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = fetcher.get( - validated_url, - headers={**headers_view["headers"], "Host": original_host}, - follow_redirects=False, + hop_options: _FetchOptions = { **kwargs, - ) + "headers": {**caller_headers, "Host": original_host}, + "follow_redirects": False, + } + response = fetcher.get(validated_url, **hop_options) if not response.is_redirect: return response # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so @@ -483,24 +494,26 @@ def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: raise SSRFError("Too many redirects") -async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Any) -> httpx.Response: +async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Unpack[_AsyncFetchOptions]) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): - kwargs.setdefault("follow_redirects", True) - unvalidated: Final[_ResponseView] = {"response": await client.get(url, **kwargs)} + unvalidated_options: Final[_AsyncFetchOptions] = { + **kwargs, + "follow_redirects": kwargs.get("follow_redirects", True), + } + unvalidated: Final[_ResponseView] = {"response": await client.get(url, **unvalidated_options)} return unvalidated["response"] fetcher_view: Final[_AsyncFetcherView] = {"fetcher": client} fetcher: Final = fetcher_view["fetcher"] - kwargs.pop("follow_redirects", None) - headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} + caller_headers: Final = kwargs.get("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = await asyncio.to_thread(validate_url, url) - response = await fetcher.get( - validated_url, - headers={**headers_view["headers"], "Host": original_host}, - follow_redirects=False, + hop_options: _AsyncFetchOptions = { **kwargs, - ) + "headers": {**caller_headers, "Host": original_host}, + "follow_redirects": False, + } + response = await fetcher.get(validated_url, **hop_options) if not response.is_redirect: return response # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 7452a4a20ef..91f33635658 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -3,6 +3,7 @@ import os from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.utils import CallTypes from . import * @@ -128,7 +129,7 @@ def discover_guardrail_translation_mappings() -> dict[CallTypes, type["BaseTrans # Check for guardrail_translation_mappings dictionary if hasattr(module, "guardrail_translation_mappings"): - mappings = getattr(module, "guardrail_translation_mappings") + mappings = attribute_or_none(module, "guardrail_translation_mappings") if isinstance(mappings, dict): discovered_mappings.update(mappings) verbose_logger.debug( diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 1818b406823..633ea73bb16 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -31,6 +31,10 @@ class A2AError(BaseLLMException): ) +def _message_role(message: Mapping[str, object]) -> object: + return message.get("role", "user") + + def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: """ Convert OpenAI messages to a single prompt string for A2A agent. @@ -51,7 +55,7 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: # Get role if isinstance(msg, BaseModel): - role = msg.model_dump().get("role", "user") + role = _message_role(msg.model_dump()) elif isinstance(msg, dict): role = msg.get("role", "user") else: diff --git a/litellm/llms/ai21/chat/transformation.py b/litellm/llms/ai21/chat/transformation.py index 2b8c36d0085..e78cc77c0fc 100644 --- a/litellm/llms/ai21/chat/transformation.py +++ b/litellm/llms/ai21/chat/transformation.py @@ -4,6 +4,8 @@ AI21 Chat Completions API this is OpenAI compatible - no translation needed / occurs """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from ...openai_like.chat.transformation import OpenAILikeChatConfig @@ -41,7 +43,7 @@ class AI21ChatConfig(OpenAILikeChatConfig): tool_choice: str | None = None, user: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/amazon_nova/chat/transformation.py b/litellm/llms/amazon_nova/chat/transformation.py index a93fbf1e933..4ef729bc76d 100644 --- a/litellm/llms/amazon_nova/chat/transformation.py +++ b/litellm/llms/amazon_nova/chat/transformation.py @@ -2,6 +2,8 @@ Translate from OpenAI's `/v1/chat/completions` to Amazon Nova's `/v1/chat/completions` """ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx @@ -38,7 +40,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): tools: list | None = None, reasoning_effort: list | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/anthropic/batches/handler.py b/litellm/llms/anthropic/batches/handler.py index 3f3716e0de6..a9d9f930b9e 100644 --- a/litellm/llms/anthropic/batches/handler.py +++ b/litellm/llms/anthropic/batches/handler.py @@ -132,7 +132,7 @@ class AnthropicBatchesHandler: max_retries: int | None, logging_obj: LiteLLMLoggingObj | None = None, litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment - ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: + ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: """ Retrieve a batch from Anthropic. diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 4c6bded68ea..3ab21099d3a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1697,14 +1697,13 @@ class AnthropicMessagesHandler(BaseTranslation): if not line.startswith("data:"): return line try: - data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads( - line[len("data:") :].strip() - ) + data: Final = _DECODED_SSE_DATA.validate_python(json.loads(line[len("data:") :].strip())) except json.JSONDecodeError: return line - if not isinstance(data, dict): + event: Final = as_str_mapping(data) + if not isinstance(event, dict): return line - rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + rewritten: Final = _rewritten_event(event, rewrite_event) return line if rewritten is data else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 9a63afd5cc0..92e7c84031c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -831,6 +831,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "else": "conditional (else): {}", "not": "must not match: {}", } + schema_fields: Final[Mapping[str, object]] = MappingProxyType(schema) min_items: Final = schema.get("minItems") enforced_min_items: Final = ( isinstance(min_items, int) and not isinstance(min_items, bool) and min_items in (0, 1) @@ -845,7 +846,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): for field, label in constraint_labels.items(): if field not in schema or field not in unsupported_fields: continue - value = schema[field] + value = schema_fields[field] # A falsy boolean constraint (e.g. ``uniqueItems: false``) imposes no # real requirement, so don't add a misleading advisory note for it. if isinstance(value, bool) and not value: @@ -869,7 +870,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): drops_conflicting_type: Final = _enum_conflicts_with_declared_type(schema) - for key, value in schema.items(): + for key, value in schema_fields.items(): if key in unsupported_fields: continue if key == "type" and drops_conflicting_type: @@ -2385,7 +2386,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["tools"] = tools ## Load Config - config: Final = litellm.AnthropicConfig.get_config(model=model) + config: Final[Mapping[str, object]] = MappingProxyType(litellm.AnthropicConfig.get_config(model=model)) for k, v in config.items(): if ( k not in optional_params diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index 73e94e85f04..5cc4e955d00 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -7,6 +7,7 @@ Litellm provider slug: `anthropic_text/` import json import time from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx @@ -112,7 +113,7 @@ class AnthropicTextConfig(BaseConfig): ) -> dict: prompt: Final = self._get_anthropic_text_prompt_from_messages(messages=messages, model=model) ## Load Config - config: Final = litellm.AnthropicTextConfig.get_config() + config: Final[Mapping[str, object]] = MappingProxyType(litellm.AnthropicTextConfig.get_config()) for k, v in config.items(): if ( k not in optional_params diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 1a793e32317..279fb92cf6c 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -5,6 +5,7 @@ from collections.abc import Coroutine, Mapping from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -24,6 +25,8 @@ from litellm.types.utils import CallTypes, LlmProviders, ModelResponse from ..chat.transformation import AnthropicConfig from ..common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP, AnthropicModelInfo +_DECODED_JSON: Final = TypeAdapter(object) + class AnthropicFilesHandler: """ @@ -194,7 +197,7 @@ class AnthropicFilesHandler: if not line.strip(): continue - anthropic_result: object = json.loads(line) + anthropic_result: object = _DECODED_JSON.validate_python(json.loads(line)) if not isinstance(anthropic_result, dict): raise TypeError( f"Anthropic batch result line is not a JSON object: {type(anthropic_result).__name__}" diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index e9fd77ef441..dadd4d8558d 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -25,7 +25,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame @@ -1367,9 +1367,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # open a block: the emitted {"type": "thinking", "thinking": ""} gets # replayed as history and Anthropic rejects it (LIT-6357). A signed # entry opens the block so the client receives the replay signature. - thinking_blocks: Final = getattr(delta, "thinking_blocks", None) - if thinking_blocks and any( - isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks + if any( + isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) + for b in attribute_items(delta, "thinking_blocks") ): return False return True diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py index 935ce77ef87..d4430ca5ab9 100644 --- a/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py +++ b/litellm/llms/anthropic/pass_through/context_management/editors/clear_tool_uses.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from functools import reduce from itertools import chain from types import MappingProxyType -from typing import Annotated, Any, Final, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeVar, cast from pydantic import BaseModel, ConfigDict, Field, StrictInt, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -32,6 +32,9 @@ from ..constants import ( from ..errors import AnthropicContextManagementError from ..placeholders import build_cleared_tool_result_content +if TYPE_CHECKING: + from litellm.types.llms.openai import ChatCompletionToolParam + T = TypeVar("T") # api.anthropic.com, measured 2026-10-19: text results of 1 to 46 bytes are never cleared or counted toward keep _TINY_RESULT_BYTES: Final = 46 @@ -216,7 +219,7 @@ def _trigger_met( current_tokens: Final = litellm.token_counter( model=model, messages=messages, - tools=cast(Any, tools), + tools=cast("list[ChatCompletionToolParam] | None", tools), ) verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens) verbose_logger.debug("context_management polyfill: threshold: %s", threshold) @@ -366,9 +369,13 @@ def apply_clear_tool_uses_20250919( tokens_before: Final = ( counted_tokens if counted_tokens is not None - else litellm.token_counter(model=model, messages=messages, tools=cast(Any, tools)) + else litellm.token_counter( + model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools) + ) + ) + tokens_after: Final = litellm.token_counter( + model=model, messages=final_messages, tools=cast("list[ChatCompletionToolParam] | None", tools) ) - tokens_after: Final = litellm.token_counter(model=model, messages=final_messages, tools=cast(Any, tools)) cleared_input_tokens: Final = max(tokens_before - tokens_after, 0) return final_messages, { diff --git a/litellm/llms/anthropic/pass_through/context_management/placeholders.py b/litellm/llms/anthropic/pass_through/context_management/placeholders.py index b3d3529f105..8a0eb6d48e1 100644 --- a/litellm/llms/anthropic/pass_through/context_management/placeholders.py +++ b/litellm/llms/anthropic/pass_through/context_management/placeholders.py @@ -1,12 +1,10 @@ """Placeholder content for cleared ``tool_result`` blocks (string or block list).""" -from typing import Any - from .constants import CLEARED_TOOL_RESULT_PLACEHOLDER def build_cleared_tool_result_content( - original_content: Any, + original_content: object, ) -> str | list[dict]: """Return a string or single text block list, matching ``original_content`` shape.""" if isinstance(original_content, list): diff --git a/litellm/llms/anthropic/pass_through/context_management/result.py b/litellm/llms/anthropic/pass_through/context_management/result.py index 5b88fe30871..4713b8d79c7 100644 --- a/litellm/llms/anthropic/pass_through/context_management/result.py +++ b/litellm/llms/anthropic/pass_through/context_management/result.py @@ -6,7 +6,7 @@ attach ``iterations`` to ``usage``. """ from dataclasses import dataclass, field -from typing import Any, Final +from typing import Final from litellm.types.llms.anthropic import ( AppliedEdit, @@ -19,8 +19,8 @@ from .constants import COMPACT_EDIT_TYPE @dataclass class PolyfillResult: - messages: list[dict[str, Any]] - system: str | list[dict[str, Any]] | None + messages: list[dict[str, object]] + system: str | list[dict[str, object]] | None applied_edits: list[AppliedEdit] = field(default_factory=list) compaction_block: CompactionBlock | None = None iterations_usage: list[UsageIteration] | None = None diff --git a/litellm/llms/anthropic/pass_through/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py index fb17910fb97..9d651f8d721 100644 --- a/litellm/llms/anthropic/pass_through/convertible_blocks.py +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -3,7 +3,6 @@ from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial from itertools import chain, groupby -from operator import itemgetter from types import MappingProxyType from typing import Final, TypeAlias @@ -279,6 +278,14 @@ def _with_kept_parts(block: object, tool_result: Mapping[str, object], parts: _K return {**tool_result, "content": [_restored_part(part, parts.get(index)) for index, part in enumerate(blocks)]} +def _entry_tool_use_id(entry: tuple[str, _KeptParts]) -> str: + return entry[0] + + +def _entry_kept_parts(entry: tuple[str, _KeptParts]) -> _KeptParts: + return entry[1] + + def _restored_block(block: object, kept: Mapping[str, Iterator[_KeptParts]]) -> object: match as_str_mapping(block): case {"type": "tool_result", "tool_use_id": str(tool_use_id)} as tool_result if ( @@ -293,13 +300,15 @@ def _restored_block(block: object, kept: Mapping[str, Iterator[_KeptParts]]) -> def keep_non_string_tool_result_texts(rewritten: Sequence[object], original: object) -> Sequence[object]: blocks: Final = chain.from_iterable(map(_user_content_blocks, _list_items(original))) - by_id: Final = tuple(sorted(chain.from_iterable(map(_tool_result_non_string_texts, blocks)), key=itemgetter(0))) + by_id: Final = tuple( + sorted(chain.from_iterable(map(_tool_result_non_string_texts, blocks)), key=_entry_tool_use_id) + ) if not any(parts for _, parts in by_id): return rewritten kept_in_order: Final = MappingProxyType( { - tool_use_id: iter(tuple(map(itemgetter(1), group))) - for tool_use_id, group in groupby(by_id, key=itemgetter(0)) + tool_use_id: iter(tuple(map(_entry_kept_parts, group))) + for tool_use_id, group in groupby(by_id, key=_entry_tool_use_id) } ) return [_with_user_blocks(message, partial(_restored_block, kept=kept_in_order)) for message in rewritten] diff --git a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py index 8d1b358c8f1..2b6161d98f6 100644 --- a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py @@ -9,8 +9,8 @@ the LLM doesn't make a tool call, and we need to return a stream to the user. """ import json -from collections.abc import Mapping -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import Final, cast from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -136,8 +136,7 @@ class FakeAnthropicMessagesStreamIterator: """Convert the non-streaming response to streaming chunks""" chunks: Final = [] - # Cast response to dict for easier access - response_dict: Final = cast(dict[str, Any], self.response) + response_dict: Final[Mapping[str, object]] = self.response # 1. message_start event usage: Final = self.response.get("usage") @@ -160,9 +159,8 @@ class FakeAnthropicMessagesStreamIterator: chunks.append(f"event: message_start\ndata: {json.dumps(message_start)}\n\n".encode()) # 2-4. For each content block, send start/delta/stop events - content_blocks: Final = response_dict.get("content", []) - for index, block in enumerate(content_blocks): - block_dict = cast(dict[str, Any], block) + content_blocks: Final = cast(Sequence[Mapping[str, object]], response_dict.get("content", [])) + for index, block_dict in enumerate(content_blocks): chunks.extend(self._create_content_block_chunks(block_dict, index)) # 5. message_delta event (with final usage and stop_reason) diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index dfe0788e50a..35db54147c5 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -722,7 +722,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): logging_obj: LiteLLMLoggingObj, ) -> AnthropicMessagesResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[AnthropicMessagesResponse] = raw_response.json() except Exception: raise AnthropicError(message=raw_response.text, status_code=raw_response.status_code) response: Final = AnthropicMessagesResponse(**raw_response_json) diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index bd0f57c3ed9..48fc987a95a 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -242,7 +242,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def _summary_part_text(part: object) -> str: if isinstance(part, Mapping): - mapping: Final = cast(Mapping[str, Any], part) # cast-ok: summary parts are untyped provider json + mapping: Final = cast(Mapping[str, object], part) # cast-ok: summary parts are untyped provider json return str(mapping.get("text") or "") return str(getattr(part, "text", None) or "") @@ -673,7 +673,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_config = anthropic_request.get("output_config") reasoning: Final = self.translate_thinking_to_reasoning( thinking, - output_config=cast(dict[str, Any] | None, output_config), + output_config=cast(dict[str, object] | None, output_config), auto_summary_level=auto_summary_level, drop_params=drop_params, ) diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 99ebd2875a7..df8a0614665 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from openai import AsyncAzureOpenAI, AzureOpenAI from pydantic import BaseModel @@ -38,7 +38,7 @@ class AzureAudioTranscription(AzureChatCompletion): atranscription: bool = False, litellm_params: dict | None = None, custom_llm_provider: str = "azure", - ) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]: + ) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]: data: Final = {"model": model, "file": audio_file, **optional_params} if atranscription is True: diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 67b25bf7d12..02f1ba3472a 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -13,7 +13,7 @@ from openai import ( BadRequestError, OpenAI, ) -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES @@ -59,6 +59,8 @@ from .image_generation import ( ) from .image_generation.http_utils import azure_deployment_image_generation_json_body +_DECODED_JSON: Final = TypeAdapter(object) + class AzureOpenAIAssistantsAPIConfig: """ @@ -961,7 +963,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): body=error_data, ) - result: Final = response.json()["result"] + result: Final = _DECODED_JSON.validate_python(response.json()["result"]) return httpx.Response( status_code=200, headers=response.headers, @@ -1078,7 +1080,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): body=error_data, ) - result: Final = response.json()["result"] + result: Final = _DECODED_JSON.validate_python(response.json()["result"]) return httpx.Response( status_code=200, headers=response.headers, @@ -1358,7 +1360,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except AzureOpenAIError as e: raise e except Exception as e: - error_code: Final = getattr(e, "status_code", None) + error_code: Final[object] = getattr(e, "status_code", None) if error_code is not None: raise AzureOpenAIError(status_code=error_code, message=str(e)) else: diff --git a/litellm/llms/azure/exception_mapping.py b/litellm/llms/azure/exception_mapping.py index dcc3601cd8e..f7146a36384 100644 --- a/litellm/llms/azure/exception_mapping.py +++ b/litellm/llms/azure/exception_mapping.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from litellm.exceptions import ContentPolicyViolationError @@ -54,7 +54,7 @@ class AzureOpenAIExceptionMapping: @staticmethod def _extract_azure_error( original_exception: Exception, - ) -> tuple[dict[str, object], dict | None]: + ) -> tuple[dict[str, object], object]: """Extract Azure OpenAI error payload and inner error details. Azure error formats can vary by endpoint/version. Common shapes: @@ -67,13 +67,13 @@ class AzureOpenAIExceptionMapping: return {}, None # Some SDKs place the payload under "error". - azure_error: dict[str, Any] + azure_error: dict[str, object] if isinstance(body_dict.get("error"), dict): azure_error = body_dict.get("error", {}) else: azure_error = body_dict - inner_error: Final = ( + inner_error: Final[object] = ( azure_error.get("inner_error") or azure_error.get("innererror") or body_dict.get("innererror") diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 8cae4bf284b..6d904503df8 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -95,7 +95,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def get_stripped_model_name(self, model: str) -> str: return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "") - def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]: + def _handle_reasoning_item(self, item: dict[str, Any]) -> Mapping[str, object]: """ Handle reasoning items to filter out the status field. Issue: https://github.com/BerriAI/litellm/issues/13484 diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 83d00a621b4..3b541310d28 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -397,7 +397,7 @@ class AzureAIStudioConfig(OpenAIConfig): return data def _drop_tool_level_extra_fields(self, request_data: dict, error_text: str) -> dict: - fields_to_drop: Final = set(re.findall(r"tools\[\d+\]\.([\w-]+)", error_text)) + fields_to_drop: Final[Collection[str]] = set(re.findall(r"tools\[\d+\]\.([\w-]+)", error_text)) tools: Final = request_data.get("tools") if fields_to_drop and isinstance(tools, list): for tool in tools: diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index b2eb4cde5ce..91d0d52a629 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -94,9 +94,10 @@ class BaseConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("_") and not isinstance( v, @@ -513,7 +514,7 @@ class BaseConfig(ABC): def apply_assembled_streaming_response_metadata( self, response: "ModelResponse", - chunks: list[Any], + chunks: list[object], ) -> None: """Hook for providers to merge chunk metadata into assembled streaming responses.""" return diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 5fe81a7f1b7..6399ded9555 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -121,7 +121,7 @@ class BaseTranslation(ABC): @staticmethod def merge_user_api_key_metadata_into_request( - request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place + request_data: dict[str, object], # mutable-ok: proxy hooks share and mutate the request payload dict in place user_api_key_dict: Optional["UserAPIKeyAuth"], ) -> None: """ diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 55a4b3da0b3..725146fe894 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -9,6 +9,8 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage +_DECODED_JSON: Final = TypeAdapter(object) + def _anthropic_stream_chunk_events(item: object) -> list[dict]: if isinstance(item, dict): @@ -30,7 +32,7 @@ def _anthropic_stream_chunk_events(item: object) -> list[dict]: if not payload or payload == "[DONE]": continue try: - parsed = json.loads(payload) + parsed = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: continue if isinstance(parsed, dict): diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 4faf0aaaf30..35b4bb90025 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -1,7 +1,7 @@ import types from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles @@ -32,9 +32,10 @@ class BaseImageEditConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index b7bbf7bde73..39ca1d8f7d7 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -11,7 +11,8 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): import types from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final import httpx @@ -60,9 +61,10 @@ class BaseInteractionsAPIConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index d3b05716452..2663a3c42e7 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -141,7 +141,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): self, llm_router: Router, model: str, - request_data: dict[str, Any], + request_data: dict[str, object], litellm_parent_otel_span: Span, ) -> ResourceObjectType: """ @@ -399,7 +399,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): async def create_resource_for_each_model( self, llm_router: Router, - request_data: dict[str, Any], + request_data: dict[str, object], target_model_names_list: list[str], litellm_parent_otel_span: Span, ) -> list[ResourceObjectType]: diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index ac5032c08f1..7e58cc6e71b 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class PassthroughMetadata(LiteLLMBaseModel): @@ -71,7 +72,7 @@ def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None def relayed_body(httpx_response: Response) -> str | dict: try: - body: Final[object] = httpx_response.json() + body: Final = _DECODED_JSON.validate_python(httpx_response.json()) except ValueError: return httpx_response.text return body if isinstance(body, dict) else httpx_response.text diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 5d2f92b5e82..7f246bf33da 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -80,7 +80,7 @@ class BaseRerankConfig(ABC): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index fb85cdb7687..be49098b62b 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -1,6 +1,7 @@ import types from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, TypedDict +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, TypedDict import httpx @@ -40,9 +41,10 @@ class BaseTextToSpeechConfig(ABC): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 28f6348e4d9..138700b3212 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -152,7 +152,7 @@ class BaseVectorStoreConfig: api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: pass @@ -164,7 +164,7 @@ class BaseVectorStoreConfig: api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Optional async version of transform_search_vector_store_request. diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py index 381d487c36d..c165831b90a 100644 --- a/litellm/llms/baseten/chat.py +++ b/litellm/llms/baseten/chat.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import litellm @@ -42,7 +43,7 @@ class BasetenConfig(OpenAIGPTConfig): frequency_penalty: int | None = None, stream_options: dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 48e79d39b8e..f37abacc883 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -174,12 +174,13 @@ class BedrockBatchesHandler: ) def job_status() -> "LiteLLMBatch": + auth_kwargs: Final[Mapping[str, object]] = dict(auth_params.model_dump()) return BedrockBatchesHandler.handle_model_invocation_job_status( batch_id=batch_id, aws_region_name=region, logging_obj=logging_obj, aws_bedrock_control_plane_endpoint=control_plane_endpoint, - **auth_params.model_dump(), + **auth_kwargs, ) try: diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index bb5bb246a93..18212b7390f 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, Union from urllib.parse import quote import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -51,6 +52,8 @@ else: HTTPHandler = Any AsyncHTTPHandler = Any +_DECODED_JSON: Final = TypeAdapter(object) + class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): def __init__(self, **kwargs): @@ -314,7 +317,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return None try: - data: Final[object] = json.loads(json_str) + data: Final = _DECODED_JSON.validate_python(json.loads(json_str)) # Skip non-dict data (some lines contain JSON strings) return data if isinstance(data, dict) else None except json.JSONDecodeError: @@ -482,7 +485,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): if "application/json" in content_type: # Direct JSON response verbose_logger.debug("Parsing JSON response") - response_json: Final[object] = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) verbose_logger.debug("Response JSON: %s", response_json) return self._parse_json_response(response_json) else: @@ -566,7 +569,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): continue try: - data_obj = json.loads(json_str) + data_obj = _DECODED_JSON.validate_python(json.loads(json_str)) if not isinstance(data_obj, dict): continue @@ -706,7 +709,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) try: body: Final = response.read() - response_json: Final[object] = json.loads(body) + response_json: Final = _DECODED_JSON.validate_python(json.loads(body)) except (json.JSONDecodeError, Exception) as e: raise BedrockError( status_code=response.status_code, @@ -788,7 +791,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): continue try: - data_obj = json.loads(json_str) + data_obj = _DECODED_JSON.validate_python(json.loads(json_str)) if not isinstance(data_obj, dict): continue @@ -928,7 +931,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) try: body: Final = await response.aread() - response_json: Final[object] = json.loads(body) + response_json: Final = _DECODED_JSON.validate_python(json.loads(body)) except (json.JSONDecodeError, Exception) as e: raise BedrockError( status_code=response.status_code, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index cd3727401ef..4d7ec663d09 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -12,7 +12,8 @@ from itertools import chain from typing import TYPE_CHECKING, Final, Literal, cast, overload import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -21,6 +22,7 @@ from litellm.constants import ( RESPONSE_FORMAT_TOOL_NAME, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, filter_exceptions_from_params, filter_internal_params, map_finish_reason, @@ -160,6 +162,14 @@ _TOOLS_AS_SENT: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _JSON_VALUE: Final = TypeAdapter(object) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) _TOOLS_AS_RECEIVED: Final = TypeAdapter(list[object]) + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _ToolCallInContent(TypedDict): + name: ReadOnly[str] + + +_TOOL_CALL_IN_CONTENT: Final = TypeAdapter(_ToolCallInContent) CLAUDE_ON_BEDROCK_INVOKE_REMEDY: Final = " This model runs tool search on Bedrock through `bedrock/invoke/`." _DEPLOYMENT_PARAMS: Final = TypeAdapter(Mapping[str, object]) @@ -2453,9 +2463,17 @@ class AmazonConverseConfig(BaseConfig): if message.content is not None: try: tool_call_names: Final = self.get_tool_call_names(tools) - json_content: Final = json.loads(message.content) - if json_content.get("type") == "function" and json_content.get("name") in tool_call_names: - tool_calls: Final = [ChatCompletionMessageToolCall(function=Function(**json_content))] + json_content: Final = as_str_mapping(_JSON_VALUE.validate_python(json.loads(message.content))) + if ( + json_content is not None + and json_content.get("type") == "function" + and json_content.get("name") in tool_call_names + ): + tool_calls: Final = [ + ChatCompletionMessageToolCall( + function=Function(**_TOOL_CALL_IN_CONTENT.validate_python(json_content)) + ) + ] message.tool_calls = tool_calls message.content = None diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py index 3005bc82299..777237b4f18 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_ai21_transformation.py @@ -56,9 +56,10 @@ class AmazonAI21Config(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py index 00c2eaf0f8f..2118d183b22 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py @@ -37,9 +37,10 @@ class AmazonCohereConfig(AmazonInvokeConfig, CohereChatConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py index d2eb6eef9ef..cb867a09f8f 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_llama_transformation.py @@ -37,9 +37,10 @@ class AmazonLlamaConfig(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py index 44511eeaa3e..4420bb3a418 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py @@ -47,9 +47,10 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py index a18e97c63a6..736b3ea3cd2 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py @@ -43,9 +43,10 @@ class AmazonTitanConfig(AmazonInvokeConfig, BaseConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not k.startswith("_abc") and not isinstance( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py index b29ddf01d1e..2655c13be44 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py @@ -44,9 +44,10 @@ class AmazonAnthropicConfig(AmazonInvokeConfig): @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index e48196946cd..d61541aab8a 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -21,6 +21,7 @@ from litellm.types.utils import ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.anthropic.chat.handler import ModelResponseIterator LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -134,7 +135,7 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "ModelResponseIterator": return self.get_model_response_iterator_for_request( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, optional_params={} ) @@ -145,7 +146,7 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): sync_stream: bool, json_mode: bool | None, optional_params: Mapping[str, object], - ) -> object: + ) -> "ModelResponseIterator": from litellm.llms.anthropic.chat.handler import ModelResponseIterator return ModelResponseIterator( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4954e646f58..2b946d9664c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -2021,8 +2021,13 @@ def get_bedrock_response_stream_shape(): _BEDROCK_STREAM_OUTPUT_SHAPES: Final = ("ConverseStreamOutput", "ResponseStream") +class _ModeledErrorMetadata(TypedDict, total=False): + httpStatusCode: ReadOnly[int] + + def _modeled_error_status(member: Shape) -> int | None: - status: Final = (member.metadata or {}).get("error", {}).get("httpStatusCode") + error: Final[_ModeledErrorMetadata] = (member.metadata or {}).get("error", {}) + status: Final = error.get("httpStatusCode") return None if status is None else int(status) diff --git a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py index 581b89369f7..fba2673b5aa 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py @@ -35,9 +35,10 @@ class AmazonTitanG1Config: @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py index 3b594662a9b..7d319ca6b06 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py @@ -39,9 +39,10 @@ class AmazonTitanV2Config: @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 3d23b69f846..37f95ff7b9b 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -1,6 +1,6 @@ import asyncio from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final import httpx @@ -139,7 +139,7 @@ class BedrockFilesHandler(BaseAWSLLM): optional_params: dict, timeout: float | httpx.Timeout, max_retries: int | None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: """ Download file content from S3 bucket for Bedrock files. Supports both sync and async operations. diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 466993de9d1..bdfe6f5a8d2 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -10,6 +10,7 @@ import json from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -39,6 +40,8 @@ if TYPE_CHECKING: else: AWSPreparedRequest = Any +_DECODED_JSON: Final = TypeAdapter(object) + class BedrockImageEditPreparedRequest(LiteLLMBaseModel): """ @@ -317,7 +320,7 @@ class BedrockImageEdit(BaseAWSLLM): additional_args={"complete_input_dict": data}, ) verbose_logger.debug("raw model_response: %s", response.text) - response_dict: Final = response.json() + response_dict: Final = _DECODED_JSON.validate_python(response.json()) if response_dict is None: raise ValueError("Error in response object format, got None") diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index acc847e0e3c..e5c5edba750 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final from openai.types.image import Image @@ -20,9 +21,10 @@ class AmazonStability3Config: @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index 145e6a0781f..5731541da56 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -3,6 +3,7 @@ Transformation logic for Amazon Titan Image Generation. """ import types +from collections.abc import Mapping from typing import Final from openai.types.image import Image @@ -35,16 +36,17 @@ class AmazonTitanImageGenerationConfig: width: int | None = None, height: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 12421aa94fd..805652630ed 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -1,4 +1,4 @@ -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast @@ -145,7 +145,7 @@ class AmazonAnthropicClaudeMessagesConfig( self, headers: dict, model: str, - messages: list[object], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index c71721d800e..d4a2d2eabe9 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -7,7 +7,7 @@ stripping, and the anthropic-version / anthropic-beta placement (headers, never the body) that are specific to the bedrock-mantle endpoint. """ -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -90,7 +90,7 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/bedrock/model_listing.py b/litellm/llms/bedrock/model_listing.py index 99ff48f9400..409bf2c7b97 100644 --- a/litellm/llms/bedrock/model_listing.py +++ b/litellm/llms/bedrock/model_listing.py @@ -4,7 +4,7 @@ from typing import Final, TypeAlias from urllib.parse import quote, urlencode import httpx -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token @@ -15,6 +15,7 @@ from litellm.types.llms.bedrock import AwsAuthParams LIST_MODELS_TIMEOUT: Final = 10.0 INFERENCE_PROFILES_PAGE_SIZE: Final = 1000 INFERENCE_PROFILES_PAGE_CAP: Final = 20 +_DECODED_JSON: Final = TypeAdapter(object) QueryPairs: TypeAlias = tuple[tuple[str, str], ...] @@ -103,7 +104,7 @@ class BedrockModelLister(BaseAWSLLM): f"Status code: {response.status_code}, Response: {response.text}" ), ) - return response.json() + return _DECODED_JSON.validate_python(response.json()) def _auth_headers(self, url: str) -> Mapping[str, str]: if self._bearer_token is not None: diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index a90b4829634..1b01b0ec6d0 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -544,7 +544,7 @@ class _StreamFrame(TypedDict): def _unpack_uint32(buffer: bytes) -> int: - return struct.unpack("!I", buffer)[0] + return int.from_bytes(buffer, "big") def _collect_stream_delta_text_holders(delta: object) -> list[_DeltaHolder]: @@ -1255,7 +1255,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): headers_bytes = frame_raw[12 : 12 + orig_hdrs_len] try: - payload_dict = json.loads(frame_raw[12 + orig_hdrs_len : orig_total - 4]) + payload_dict = _BODY_BY_REFERENCE.validate_python( + json.loads(frame_raw[12 + orig_hdrs_len : orig_total - 4]) + ) for local_idx, (_, container, key) in enumerate( _collect_stream_delta_text_holders(payload_dict.get("delta")) ): diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 7cfb0daf334..38bcbb0d905 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -51,6 +51,7 @@ class _BflPolledResult(TypedDict): _BFL_POLLED_RESULT_ADAPTER: Final = TypeAdapter(_BflPolledResult) +_DECODED_JSON: Final = TypeAdapter(object) _BFL_REQUEST_PARAMS: Final = ( @@ -316,7 +317,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): The response contains: {"status": "Ready", "result": {"sample": "https://..."}} """ try: - raw_response_data: Final[object] = raw_response.json() + raw_response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise BlackForestLabsError( status_code=raw_response.status_code, diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index 928af314506..52f2c9a7c50 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -40,6 +40,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): @@ -272,7 +273,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): The response contains: {"status": "Ready", "result": {"sample": "https://..."}} """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise BlackForestLabsError( status_code=raw_response.status_code, diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py index c3aa26ade35..9fc99500fc5 100644 --- a/litellm/llms/cerebras/chat.py +++ b/litellm/llms/cerebras/chat.py @@ -4,6 +4,7 @@ Cerebras Chat Completions API this is OpenAI compatible - no translation needed / occurs """ +from collections.abc import Mapping from typing import Final from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig @@ -41,7 +42,7 @@ class CerebrasConfig(OpenAIGPTConfig): user: str | None = None, reasoning_effort: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index ee3120bacea..49be583f6d1 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -7,6 +7,8 @@ Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. from collections.abc import Awaitable from typing import Final, Protocol +from pydantic import TypeAdapter + from litellm.types.utils import ( ChatCompletionDeltaCustomToolCall, ChatCompletionDeltaToolCall, @@ -23,6 +25,9 @@ class ChatGPTChunkStream(Protocol): def __anext__(self) -> Awaitable[ModelResponseStream]: ... +_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) + + def _first_choice_delta(chunk: ModelResponseStream) -> Delta | None: return chunk.choices[0].delta @@ -47,7 +52,7 @@ class ChatGPTToolCallNormalizer: self._last_id: str | None = None # tracks which tool call the next delta belongs to def __getattr__(self, name: str) -> object: - return getattr(self._stream, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self._stream, name)) def __iter__(self) -> "ChatGPTToolCallNormalizer": return self diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index 8c2fd9e0aff..d07ae5e4215 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -23,6 +23,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class ClarifaiConfig(OpenAIGPTConfig): @@ -114,7 +115,7 @@ class ClarifaiConfig(OpenAIGPTConfig): ) ## Reponse try: - completion_response: Final = raw_response.json() + completion_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OpenAIError( status_code=raw_response.status_code, diff --git a/litellm/llms/claude_code/harness/transformation.py b/litellm/llms/claude_code/harness/transformation.py index f75d86c7e63..8252d2dd461 100644 --- a/litellm/llms/claude_code/harness/transformation.py +++ b/litellm/llms/claude_code/harness/transformation.py @@ -30,6 +30,7 @@ from litellm.harness.types import ( ToolCall, ToolResult, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, HarnessSessionSetup, @@ -231,12 +232,12 @@ def _system_events(event: Mapping[str, object], state: ClaudeCodeStreamState) -> return event_list() -def _record_result(event: Mapping[str, Any], state: ClaudeCodeStreamState) -> Sequence[Event]: +def _record_result(event: Mapping[str, object], state: ClaudeCodeStreamState) -> Sequence[Event]: state.result_seen = True state.is_error = bool(event.get("is_error", False)) result = event.get("result") state.result_text = result if isinstance(result, str) else None - state.errors = [str(e) for e in event.get("errors") or ()] # mutable-ok: mirrors the JSON errors array + state.errors = [str(e) for e in iter_items(event.get("errors") or ())] # mutable-ok: mirrors the JSON errors array state.structured_output = event.get("structured_output") if event.get("session_id"): state.session_id = str(event["session_id"]) diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 08799c3a0a5..6bebfe2d444 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -2,7 +2,7 @@ ## handler file for TextCompletionCodestral Integration - https://codestral.com/ import json -from collections.abc import Callable +from collections.abc import Callable, Mapping from functools import partial from typing import Final, Protocol @@ -268,7 +268,7 @@ class CodestralTextCompletion: prompt = prompt_factory(model=model, messages=messages) ## Load Config - config: Final = litellm.CodestralTextCompletionConfig.get_config() + config: Final[Mapping[str, object]] = litellm.CodestralTextCompletionConfig.get_config() for k, v in config.items(): if ( k not in optional_params diff --git a/litellm/llms/codex/harness/transformation.py b/litellm/llms/codex/harness/transformation.py index 4512e39189e..d7388bd23fc 100644 --- a/litellm/llms/codex/harness/transformation.py +++ b/litellm/llms/codex/harness/transformation.py @@ -15,7 +15,7 @@ import re from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from pydantic import ConfigDict, TypeAdapter @@ -31,6 +31,7 @@ from litellm.harness.types import ( ToolCall, ToolResult, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, HarnessSessionSetup, @@ -77,13 +78,13 @@ class CodexStreamState: started: set[str] = field(default_factory=set) # mutable-ok: parser records announced tool items -def _tool_input(item: Mapping[str, Any]) -> tuple[str, str, Mapping[str, object], bool]: +def _tool_input(item: Mapping[str, object]) -> tuple[str, str, Mapping[str, object], bool]: """(normalized name, native name, input, builtin) for a tool-like item.""" item_type = item.get("type") if item_type == "command_execution": return "bash", "command_execution", MappingProxyType({"command": item.get("command", "")}), True if item_type == "file_change": - changes: Final = list(item.get("changes") or ()) # mutable-ok: JSON array, as codex reports it + changes: Final = list(iter_items(item.get("changes") or ())) # mutable-ok: JSON array, as codex reports it return "edit", "apply_patch", MappingProxyType({"changes": changes}), True if item_type == "web_search": return "web_search", "web_search", MappingProxyType({"query": item.get("query", "")}), True diff --git a/litellm/llms/cohere/chat/transformation.py b/litellm/llms/cohere/chat/transformation.py index fdf44d5333e..e7fadd7f375 100644 --- a/litellm/llms/cohere/chat/transformation.py +++ b/litellm/llms/cohere/chat/transformation.py @@ -1,6 +1,6 @@ import json import time -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -221,7 +221,8 @@ class CohereChatConfig(BaseConfig): headers: dict, ) -> dict: ## Load Config - for k, v in litellm.CohereChatConfig.get_config().items(): + config: Final[Mapping[str, object]] = litellm.CohereChatConfig.get_config() + for k, v in config.items(): if ( k not in optional_params ): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 8b641c24e9b..7d9c8de7a07 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -15,6 +15,7 @@ from litellm.types.rerank import OptionalRerankParams, RerankRequest, RerankResp from ..common_utils import CohereError _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class CohereRerankConfig(BaseRerankConfig): @@ -144,7 +145,7 @@ class CohereRerankConfig(BaseRerankConfig): No transformation required, litellm follows cohere API response format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise CohereError(message=raw_response.text, status_code=raw_response.status_code) diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 009a220534c..83596f90d1f 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.cohere.rerank.transformation import CohereRerankConfig @@ -39,7 +39,7 @@ class CohereRerankV2Config(CohereRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/cometapi/chat/transformation.py b/litellm/llms/cometapi/chat/transformation.py index 2bf119a0a4a..ae04889bc97 100644 --- a/litellm/llms/cometapi/chat/transformation.py +++ b/litellm/llms/cometapi/chat/transformation.py @@ -6,7 +6,7 @@ Documentation: [CometAPI Documentation Link] """ from collections.abc import AsyncIterator, Iterator -from typing import Any, Final +from typing import Final import httpx @@ -137,7 +137,7 @@ class CometAPIConfig(OpenAIGPTConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "CometAPIChatCompletionStreamingHandler": """ Get model response iterator for streaming responses """ diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index a02db2338d8..87fc61985e6 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -3,10 +3,10 @@ CompactifAI chat completion transformation """ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import OpenAIError @@ -25,11 +25,32 @@ else: LiteLLMLoggingObj = Any +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + +class _CompactifAIToolCall(TypedDict): + function: ReadOnly[Mapping[str, str]] + + +class _CompactifAIMessage(TypedDict, total=False): + content: str | None # writable-ok: json mode moves a lone tool call's arguments into the content in place + tool_calls: Sequence[_CompactifAIToolCall] | None # writable-ok: json mode clears the tool call it moved + + +class _CompactifAIChoice(TypedDict, total=False): + message: ReadOnly[_CompactifAIMessage] + + class CompactifAIResponseFields(TypedDict, total=False): """The chat completion fields of a CompactifAI response body.""" id: ReadOnly[str] - choices: ReadOnly[Sequence[Mapping[str, object]]] + choices: ReadOnly[Required[Sequence[_CompactifAIChoice]]] created: ReadOnly[int] model: ReadOnly[str] system_fingerprint: ReadOnly[str | None] @@ -82,22 +103,19 @@ class CompactifAIChatConfig(OpenAIGPTConfig): ) ## RESPONSE OBJECT - response_json: Final = raw_response.json() + response_json: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + CompactifAIResponseFields, _json_body(raw_response) + ) # Handle JSON mode if needed if json_mode: for choice in response_json["choices"]: message = choice.get("message") - if message and message.get("tool_calls"): - # Convert tool calls to content for JSON mode - tool_calls = message.get("tool_calls", []) - if len(tool_calls) == 1: - message["content"] = tool_calls[0]["function"].get("arguments", "") - message["tool_calls"] = None + if message and (tool_calls := message.get("tool_calls")) and len(tool_calls) == 1: + message["content"] = tool_calls[0]["function"].get("arguments", "") + message["tool_calls"] = None - response_fields: Final[CompactifAIResponseFields] = response_json - - returned_response: Final = ModelResponse(**response_fields) + returned_response: Final = ModelResponse(**response_json) # Set model name with provider prefix returned_response.model = f"compactifai/{model}" diff --git a/litellm/llms/custom_httpx/async_client_cleanup.py b/litellm/llms/custom_httpx/async_client_cleanup.py index 6ef61b6f58f..3c0cc90f30f 100644 --- a/litellm/llms/custom_httpx/async_client_cleanup.py +++ b/litellm/llms/custom_httpx/async_client_cleanup.py @@ -20,7 +20,7 @@ async def close_litellm_async_clients(): cache_dict: Final = getattr(litellm.in_memory_llm_clients_cache, "cache_dict", {}) - for key, handler in cache_dict.items(): + for handler in cache_dict.values(): # Handle BaseLLMAIOHTTPHandler instances (aiohttp_openai provider) if isinstance(handler, BaseLLMAIOHTTPHandler) and hasattr(handler, "close"): try: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a692f3f964e..55b2c21718d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -49,7 +49,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) -from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header @@ -412,6 +412,27 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False +def _run_responses_deployment_hook( + *, + model: str, + request_input: str | ResponseInputParam, + custom_llm_provider: str, + optional_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, +) -> Mapping[str, object]: + return run_async_function( + async_pre_call_deployment_hook, + { + **dict(litellm_params), + **optional_params, + "model": model, + "input": request_input, + "custom_llm_provider": custom_llm_provider, + }, + CallTypes.responses.value, + ) + + def _mask_presigned_request_headers(transformed_request: bytes | str | dict) -> bytes | str | dict: """A pre-signed request carries its auth inside its own ``headers`` key, which logging treats as request body (only the top-level headers channel gets masked), @@ -477,6 +498,7 @@ _DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encodin _OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_OBJECTS: Final = TypeAdapter(list[dict[str, object]]) +_DECODED_JSON: Final = TypeAdapter(object) def _decoded_body_headers(response: httpx.Response) -> httpx.Headers: @@ -2765,16 +2787,12 @@ class BaseLLMHTTPHandler: litellm_params, ) - modified_kwargs: Final = run_async_function( - async_pre_call_deployment_hook, - { - **dict(litellm_params), - **response_api_optional_request_params, - "model": model, - "input": input, - "custom_llm_provider": custom_llm_provider, - }, - CallTypes.responses.value, + modified_kwargs: Final = _run_responses_deployment_hook( + model=model, + request_input=input, + custom_llm_provider=custom_llm_provider, + optional_params=response_api_optional_request_params, + litellm_params=litellm_params, ) if modified_kwargs is None: return ( @@ -3753,7 +3771,7 @@ class BaseLLMHTTPHandler: response: httpx.Response, upload_url_location: str, upload_url_key: str = "upload_url", - ) -> tuple[str | None, dict | None]: + ) -> tuple[object, Mapping[str, object] | None]: """ Extract upload URL from initial file creation response. @@ -3773,8 +3791,10 @@ class BaseLLMHTTPHandler: else: # Response body style (e.g., Manus, S3 presigned URLs) try: - response_data: Final = response.json() - upload_url = response_data.get(upload_url_key) + response_data: Final = as_str_mapping(_DECODED_JSON.validate_python(response.json())) + if response_data is None: + return None, None + upload_url: Final = response_data.get(upload_url_key) return upload_url, response_data if upload_url else None except Exception: return None, None @@ -3940,7 +3960,7 @@ class BaseLLMHTTPHandler: elif isinstance(transformed_request, dict) and "file" in transformed_request: # Handle multipart form-data uploads (e.g., Anthropic Files API) # The dict contains tuples suitable for httpx's `files` parameter - file_request: Final = cast(dict[str, Any], transformed_request) + file_request: Final = cast(dict[str, object], transformed_request) upload_response = sync_httpx_client.post( url=api_base, headers=headers, diff --git a/litellm/llms/custom_httpx/mock_transport.py b/litellm/llms/custom_httpx/mock_transport.py index 0322cceec0f..6acadb8e98a 100644 --- a/litellm/llms/custom_httpx/mock_transport.py +++ b/litellm/llms/custom_httpx/mock_transport.py @@ -12,6 +12,9 @@ import uuid from typing import Final import httpx +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import decoded_json_object # --------------------------------------------------------------------------- # Pre-built response templates @@ -22,7 +25,7 @@ def _mock_id() -> str: return f"chatcmpl-mock-{uuid.uuid4().hex[:8]}" -def _chat_completion_json(model: str) -> dict: +def _chat_completion_json(model: object) -> dict: """Return a minimal valid ChatCompletion object.""" return { "id": _mock_id(), @@ -55,6 +58,8 @@ _JSON_HEADERS: Final = { "content-type": "application/json", } +_DECODED_JSON: Final = TypeAdapter(object) + class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): """ @@ -64,13 +69,13 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): """ @staticmethod - def _parse_request(request: httpx.Request) -> tuple[str, bool]: + def _parse_request(request: httpx.Request) -> tuple[object, bool]: """Extract model from the request body.""" try: - body: Final = json.loads(request.content) + body: Final = _DECODED_JSON.validate_python(json.loads(request.content)) except (json.JSONDecodeError, ValueError): return ("mock-model", False) - model: Final = body.get("model", "mock-model") + model: Final = decoded_json_object(body).get("model", "mock-model") return (model, False) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 3f388637a83..55d6fb1a0e2 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -55,10 +55,9 @@ def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: ) -def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) -> float: - value: Final = model_info.get(cost_key) +def _flat_rate(value: float | None, fallback_value: float | None) -> float: if value is None: - return float(model_info.get(fallback_cost_key) or 0.0) + return float(fallback_value or 0.0) return float(value) @@ -66,8 +65,12 @@ def _flat_rates(model_info: ModelInfo) -> TokenRates: reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token") return TokenRates( input_rate=float(model_info.get("input_cost_per_token") or 0.0), - cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"), - cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"), + cache_read_rate=_flat_rate( + model_info.get("cache_read_input_token_cost"), model_info.get("input_cost_per_token") + ), + cache_creation_rate=_flat_rate( + model_info.get("cache_creation_input_token_cost"), model_info.get("input_cost_per_token") + ), output_rate=float(model_info.get("output_cost_per_token") or 0.0), reasoning_rate=None if reasoning_rate is None else float(reasoning_rate), ) diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index 6a172e1e074..d1faa6c6cd0 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -49,6 +49,7 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" @@ -187,7 +188,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): ) try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Failed to parse DashScope image generation response: {e}", diff --git a/litellm/llms/databricks/embed/transformation.py b/litellm/llms/databricks/embed/transformation.py index ddaf2bbb621..0feb2aea5a0 100644 --- a/litellm/llms/databricks/embed/transformation.py +++ b/litellm/llms/databricks/embed/transformation.py @@ -3,6 +3,7 @@ Translates from OpenAI's `/v1/embeddings` to Databricks' `/embeddings` """ import types +from collections.abc import Mapping from typing import Final @@ -16,16 +17,17 @@ class DatabricksEmbeddingConfig: ) def __init__(self, instruction: str | None = None) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/databricks/messages/transformation.py b/litellm/llms/databricks/messages/transformation.py index cb8d2b652a2..f5ccbd92fb4 100644 --- a/litellm/llms/databricks/messages/transformation.py +++ b/litellm/llms/databricks/messages/transformation.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterator, Mapping from typing import Final import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items @@ -25,6 +26,7 @@ DROP_CONTAINER_WARNING: Final = ( "and the chat completions bridge it replaces never forwarded it." ) BEDROCK_INVOCATION_METRICS_FIELD: Final = "amazon-bedrock-invocationMetrics" +_DECODED_JSON: Final = TypeAdapter(object) def _is_unsupported_server_tool(tool: object) -> bool: @@ -57,7 +59,7 @@ def _without_unsupported_server_tools(params: Mapping[str, object]) -> dict[str, def _without_invocation_metrics(sse_line: bytes) -> bytes: if not sse_line.startswith(b"data:") or BEDROCK_INVOCATION_METRICS_FIELD.encode() not in sse_line: return sse_line - event: Final = as_str_mapping(json.loads(sse_line.removeprefix(b"data:"))) + event: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(sse_line.removeprefix(b"data:")))) if event is None: return sse_line kept_fields: Final = {key: value for key, value in event.items() if key != BEDROCK_INVOCATION_METRICS_FIELD} diff --git a/litellm/llms/deepagents/harness/sandbox_backend.py b/litellm/llms/deepagents/harness/sandbox_backend.py index d39690ccdd1..8ad90faa2b4 100644 --- a/litellm/llms/deepagents/harness/sandbox_backend.py +++ b/litellm/llms/deepagents/harness/sandbox_backend.py @@ -22,7 +22,7 @@ import shlex import uuid from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import Any, Final, TypeVar +from typing import Final, TypeVar from deepagents.backends.protocol import ( DeleteResult, @@ -171,7 +171,7 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa # -- sync bridge (deepagents only calls these outside the event loop) --- - def _sync(self, coro: Coroutine[Any, Any, T]) -> T: + def _sync(self, coro: Coroutine[object, object, T]) -> T: try: running = asyncio.get_running_loop() except RuntimeError: diff --git a/litellm/llms/deepagents/harness/transformation.py b/litellm/llms/deepagents/harness/transformation.py index cb80e3e12f7..54a6a5573ac 100644 --- a/litellm/llms/deepagents/harness/transformation.py +++ b/litellm/llms/deepagents/harness/transformation.py @@ -11,11 +11,13 @@ from __future__ import annotations import itertools import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol +from pydantic import TypeAdapter + from litellm.harness.options import DeepAgentsOptions from litellm.harness.types import ( Capabilities, @@ -41,6 +43,7 @@ DEEPAGENTS_DEFAULT_RECURSION_LIMIT: Final = 1000 _MODEL_NODE: Final = "model" # Only these graph nodes produce new messages; middleware hooks may re-emit history. _EVENT_NODES: Final = frozenset({"model", "tools"}) +_TOOL_CALLS: Final = TypeAdapter(Iterable[Mapping[str, object]]) NORMALIZED_TO_NATIVE: Final[Mapping[str, str]] = MappingProxyType( { @@ -208,7 +211,7 @@ def update_events( def _message_events(message: object, skip_tools: frozenset[str]) -> tuple[Event, ...]: kind = getattr(message, "type", None) if kind == "ai": - calls = getattr(message, "tool_calls", None) or () + calls: Final = _TOOL_CALLS.validate_python(getattr(message, "tool_calls", None) or ()) return tuple(tool_call_event(call) for call in calls if call.get("name") not in skip_tools) if kind == "tool" and getattr(message, "name", None) not in skip_tools: status: Final[object] = getattr(message, "status", None) diff --git a/litellm/llms/deepgram/audio_transcription/transformation.py b/litellm/llms/deepgram/audio_transcription/transformation.py index a1c3021c9f5..1c8a056116f 100644 --- a/litellm/llms/deepgram/audio_transcription/transformation.py +++ b/litellm/llms/deepgram/audio_transcription/transformation.py @@ -2,12 +2,13 @@ Translates from OpenAI's `/v1/audio/transcriptions` to Deepgram's `/v1/listen` """ -from collections.abc import Iterable, Mapping -from typing import Final +from collections.abc import Iterable, Mapping, Sequence +from typing import Final, Protocol, cast from urllib.parse import urlencode from httpx import Headers, Response -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, InstanceOf, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.litellm_core_utils.url_utils import join_api_base @@ -27,6 +28,31 @@ from ..common_utils import DeepgramException _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_RESPONSE_BODY: Final = TypeAdapter(InstanceOf[dict[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_WORD_LIST: Final = TypeAdapter(InstanceOf[list[object]], config=ConfigDict(hide_input_in_errors=True)) + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + +class _DeepgramAlternative(TypedDict): + transcript: ReadOnly[str] + words: ReadOnly[Sequence[Mapping[str, object]]] + paragraphs: ReadOnly[NotRequired[Mapping[str, str]]] + + +class _DeepgramChannel(TypedDict): + alternatives: ReadOnly[Sequence[_DeepgramAlternative]] + + +class _DeepgramResponse(TypedDict): + results: ReadOnly[Mapping[str, Sequence[_DeepgramChannel]]] + metadata: ReadOnly[Mapping[str, float]] class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @@ -82,7 +108,10 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): Transforms the raw response from Deepgram to the TranscriptionResponse format """ try: - response_json: Final = raw_response.json() + body: Final = _json_body(raw_response) + response_json: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + _DeepgramResponse, body + ) # Get the first alternative from the first channel first_channel: Final = response_json["results"]["channels"][0] @@ -102,7 +131,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): text = first_alternative["paragraphs"]["transcript"] else: # Diarization without paragraphs: reconstruct from words - text = self._reconstruct_diarized_transcript(first_alternative["words"]) + text = self._reconstruct_diarized_transcript(_WORD_LIST.validate_python(first_alternative["words"])) # Create TranscriptionResponse object response: Final = TranscriptionResponse(text=text) @@ -124,7 +153,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ] # Store full response in hidden params - response.hidden_params = response_json + response.hidden_params = _RESPONSE_BODY.validate_python(body) return response diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 5196c44d504..62a780396a1 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Final, Protocol import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._uuid import uuid @@ -24,6 +25,8 @@ from litellm.types.rerank import ( RerankTokens, ) +_DECODED_JSON: Final = TypeAdapter(object) + class _DeepinfraInferenceStatus(TypedDict, total=False): """The ``inference_status`` block of a DeepInfra rerank response.""" @@ -251,7 +254,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): try: error_data = error_message if isinstance(error_message, str): - error_data = json.loads(error_message) + error_data = _DECODED_JSON.validate_python(json.loads(error_message)) if isinstance(error_data, dict): # Check for {"detail": {"error": "..."}} detail: Final = error_data.get("detail") diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index 6ddee99aea3..3b2dd5af260 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -2,8 +2,10 @@ DeepSeek Anthropic-compatible messages transformation config. """ -from collections.abc import Sequence -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.url_utils import join_api_base @@ -13,6 +15,24 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +_TOOL_LIST: Final = TypeAdapter(list[object]) +_TOOL_FIELDS: Final = TypeAdapter(dict[object, object]) + + +def _custom_tool_without_type(tool: object) -> Mapping[object, object] | None: + if not isinstance(tool, dict): + return None + fields: Final = _TOOL_FIELDS.validate_python(tool) + if fields.get("type") == "custom": + fields.pop("type", None) + return fields + return None + + +def _sanitized_tool(tool: object) -> object: + stripped: Final = _custom_tool_without_type(tool) + return tool if stripped is None else stripped + class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): """ @@ -96,19 +116,11 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): return join_api_base(base_url, "v1/messages") @staticmethod - def _sanitize_tools_for_deepseek(tools: Any) -> object: + def _sanitize_tools_for_deepseek(tools: object) -> object: if not isinstance(tools, list): return tools - sanitized_tools: Final = [] - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "custom": - sanitized_tool = dict(tool) - sanitized_tool.pop("type", None) - sanitized_tools.append(sanitized_tool) - else: - sanitized_tools.append(tool) - return sanitized_tools + return [_sanitized_tool(tool) for tool in _TOOL_LIST.validate_python(tools)] def transform_anthropic_messages_request( self, diff --git a/litellm/llms/deprecated_providers/palm.py b/litellm/llms/deprecated_providers/palm.py index 785cffa48ea..b4a32630c42 100644 --- a/litellm/llms/deprecated_providers/palm.py +++ b/litellm/llms/deprecated_providers/palm.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final @@ -41,16 +42,17 @@ class PalmConfig: top_p: float | None = None, max_output_tokens: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/edenai/embedding/transformation.py b/litellm/llms/edenai/embedding/transformation.py index c79a6839434..13fa7142916 100644 --- a/litellm/llms/edenai/embedding/transformation.py +++ b/litellm/llms/edenai/embedding/transformation.py @@ -5,7 +5,7 @@ per-request cost at the top level of the body. Docs: https://www.edenai.co/docs/v3/llms/embeddings """ -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol, cast import httpx @@ -24,6 +24,14 @@ if TYPE_CHECKING: _SUPPORTED_PARAMS: Final = ("dimensions", "encoding_format", "user") +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class EdenAIEmbeddingConfig(BaseEmbeddingConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract return list(_SUPPORTED_PARAMS) @@ -80,10 +88,14 @@ class EdenAIEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict[str, object], # mutable-ok: inherited contract litellm_params: dict[str, object], # mutable-ok: inherited contract ) -> EmbeddingResponse: - body: Final = raw_response.json() + body: Final = _json_body(raw_response) logging_obj.post_call(original_response=body) response: Final[EmbeddingResponse] = convert_to_model_response_object( - response_object=body, model_response_object=model_response, response_type="embedding" + response_object=cast( # cast-ok: the body passes through unvalidated, as before it was typed + dict[str, object], body + ), + model_response_object=model_response, + response_type="embedding", ) set_response_cost_in_hidden_params(response, reported_cost(body)) return response diff --git a/litellm/llms/edenai/image_generation/transformation.py b/litellm/llms/edenai/image_generation/transformation.py index e2a0b8684f6..a0c4e1b0458 100644 --- a/litellm/llms/edenai/image_generation/transformation.py +++ b/litellm/llms/edenai/image_generation/transformation.py @@ -5,7 +5,7 @@ for every image model in the catalog with the real per-request cost at the top l Docs: https://www.edenai.co/docs/v3/llms/image-generation """ -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol, cast import httpx @@ -36,6 +36,14 @@ _SUPPORTED_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( ) +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class EdenAIImageGenerationConfig(BaseImageGenerationConfig): def get_supported_openai_params( self, model: str @@ -97,10 +105,14 @@ class EdenAIImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - body: Final = raw_response.json() + body: Final = _json_body(raw_response) logging_obj.post_call(original_response=body) response: Final[ImageResponse] = convert_to_model_response_object( - response_object=body, model_response_object=model_response, response_type="image_generation" + response_object=cast( # cast-ok: the body passes through unvalidated, as before it was typed + dict[str, object], body + ), + model_response_object=model_response, + response_type="image_generation", ) set_response_cost_in_hidden_params(response, reported_cost(body)) return response diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index 5a673b5ac41..3dc42f35d3c 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/audio/transcriptions` to ElevenLabs's `/v1/speech- """ from collections.abc import Iterable, Mapping -from typing import Final +from typing import Final, Protocol from httpx import Headers, Response from pydantic import ConfigDict, TypeAdapter @@ -27,6 +27,15 @@ from ..common_utils import ElevenLabsException _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_HIDDEN_PARAMS: Final = TypeAdapter(dict[str, object]) + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @@ -120,7 +129,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): Transforms the raw response from ElevenLabs to the TranscriptionResponse format """ try: - response_json: Final = raw_response.json() + response_json: Final = _json_body(raw_response) response_object: Final = _JSON_OBJECT.validate_python(response_json) # Extract the main transcript text @@ -135,20 +144,18 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # Map ElevenLabs words to OpenAI format if "words" in response_object: - response["words"] = [] - for word_data in _JSON_OBJECTS.validate_python(response_object["words"]): - # Only include actual words, skip spacing and audio events - if word_data.get("type") == "word": - response["words"].append( - { - "word": word_data.get("text", ""), - "start": word_data.get("start", 0), - "end": word_data.get("end", 0), - } - ) + response["words"] = [ + { + "word": word_data.get("text", ""), + "start": word_data.get("start", 0), + "end": word_data.get("end", 0), + } + for word_data in _JSON_OBJECTS.validate_python(response_object["words"]) + if word_data.get("type") == "word" + ] # Store full response in hidden params - response.hidden_params = response_json + response.hidden_params = _HIDDEN_PARAMS.validate_python(response_json) return response diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 0fdd89aec91..fc215e3d94e 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -19,6 +19,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class FalAIFluxProV11UltraConfig(FalAIBaseConfig): @@ -222,7 +223,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index 038a06c5749..26b257f5c89 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -18,6 +18,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class FalAIRecraftV3Config(FalAIBaseConfig): @@ -195,7 +196,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index 30d51cc8084..fdd552efd73 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -19,6 +19,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class FalAIStableDiffusionConfig(FalAIBaseConfig): @@ -236,7 +237,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/featherless_ai/chat/transformation.py b/litellm/llms/featherless_ai/chat/transformation.py index a68261c58e9..4bf0486f723 100644 --- a/litellm/llms/featherless_ai/chat/transformation.py +++ b/litellm/llms/featherless_ai/chat/transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import litellm @@ -42,7 +44,7 @@ class FeatherlessAIConfig(OpenAIGPTConfig): tool_choice: str | None = None, tools: list | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 964f85c821e..84c93974cd7 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -44,6 +44,7 @@ class _FireworksAIResultFields(LiteLLMBaseModel): _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str) +_DECODED_JSON: Final = TypeAdapter(object) class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): @@ -200,7 +201,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): Transform Fireworks AI rerank response to LiteLLM RerankResponse format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Failed to parse response: {e}", diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index cd9ffe9e03e..680db52b283 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -10,13 +10,17 @@ from collections.abc import Callable from typing import Final, Protocol from urllib.parse import urlsplit +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict, Unpack import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues +_DECODED_JSON: Final = TypeAdapter(object) + class _OpenAIGPTConfigOptions(TypedDict, total=False): """The sampling defaults ``OpenAIGPTConfig.__init__`` accepts and stashes on the class.""" @@ -203,10 +207,10 @@ class GDCGeminiConfig(OpenAILikeChatConfig): import google.auth try: - json_obj: Final = json.loads(api_key) + json_obj: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(api_key))) except json.JSONDecodeError: return None, False - if not isinstance(json_obj, dict) or json_obj.get("type") != self._GDCH_CREDENTIAL_TYPE: + if json_obj is None or json_obj.get("type") != self._GDCH_CREDENTIAL_TYPE: raise ValueError( "GDC only accepts a GDCH service account credential as a JSON api_key " '(expected "type": "gdch_service_account"). Other Google credential types are ' diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index c19bd9ce87c..f7352da5fb6 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -10,11 +10,12 @@ Proxies the Gemini v1beta Agents API: """ from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo @@ -29,6 +30,15 @@ from litellm.types.agents import ( # create-agent body verbatim. _GEMINI_AGENT_BODY_KEYS: Final = ("base_agent", "instructions", "base_environment") + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + # LiteLLM-internal keys that must never be forwarded to Gemini. _LITELLM_INTERNAL_KEYS: Final = frozenset( { @@ -156,18 +166,19 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): """ self._raise_for_status(raw_response) try: - data: dict[str, Any] = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: verbose_logger.warning( "GeminiAgentsConfig: non-JSON create response (status=%d).", raw_response.status_code, ) - data = {"id": name} + decoded = {"id": name} + data: Final = dict(decoded_json_object(decoded)) # Gemini uses "id" as the identifier; normalise to both fields. data.setdefault("id", name) data.setdefault("name", data["id"]) verbose_logger.debug("GeminiAgentsConfig create response: %s", data) - return AgentCreateResponse(**data) + return AgentCreateResponse.model_validate(data) # ------------------------------------------------------------------ # # LIST # @@ -192,13 +203,16 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): ) -> AgentListResponse: self._raise_for_status(raw_response) try: - data = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: - data = {} + decoded = {} + data: Final = decoded_json_object(decoded) verbose_logger.debug("GeminiAgentsConfig list response: %s", data) - return AgentListResponse( - agents=data.get("agents", []), - next_page_token=data.get("nextPageToken"), + return AgentListResponse.model_validate( + { + "agents": data.get("agents", []), + "next_page_token": data.get("nextPageToken"), + } ) # ------------------------------------------------------------------ # @@ -222,13 +236,14 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): """Same shape as create response — Gemini returns "id" as identifier.""" self._raise_for_status(raw_response) try: - data = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: - data = {"id": name} + decoded = {"id": name} + data: Final = dict(decoded_json_object(decoded)) data.setdefault("id", name) data.setdefault("name", data["id"]) verbose_logger.debug("GeminiAgentsConfig get response: %s", data) - return AgentCreateResponse(**data) + return AgentCreateResponse.model_validate(data) # ------------------------------------------------------------------ # # DELETE # @@ -285,11 +300,14 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): """ self._raise_for_status(raw_response) try: - data = raw_response.json() + decoded: object = _json_body(raw_response) except Exception: - data = {} + decoded = {} + data: Final = decoded_json_object(decoded) verbose_logger.debug("GeminiAgentsConfig list_versions response for '%s': %s", name, data) - return AgentVersionsResponse( - agent_versions=data.get("agentVersions", []), - next_page_token=data.get("nextPageToken"), + return AgentVersionsResponse.model_validate( + { + "agent_versions": data.get("agentVersions", []), + "next_page_token": data.get("nextPageToken"), + } ) diff --git a/litellm/llms/gemini/audio_transcription/transformation.py b/litellm/llms/gemini/audio_transcription/transformation.py index ea116a7e9de..3457b22faa4 100644 --- a/litellm/llms/gemini/audio_transcription/transformation.py +++ b/litellm/llms/gemini/audio_transcription/transformation.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from typing import Final from httpx import Headers, Response +from pydantic import TypeAdapter from litellm.litellm_core_utils.audio_utils.subtitle_utils import SUBTITLE_RESPONSE_FORMATS from litellm.litellm_core_utils.audio_utils.utils import ( @@ -36,6 +37,7 @@ from litellm.types.utils import ( INTERACTIONS_API_REVISION: Final = "2026-05-20" WORD_INFO_ANNOTATION_TYPE: Final = "word_info" +_DECODED_JSON: Final = TypeAdapter(object) class GeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @@ -132,7 +134,7 @@ class GeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig): raw_response: Response, ) -> TranscriptionResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except ValueError: raise GeminiError( status_code=raw_response.status_code, diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 13df3921d13..9634c580858 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -2,9 +2,9 @@ import base64 import datetime import json import math -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, cast import httpx from pydantic import JsonValue, TypeAdapter @@ -389,7 +389,7 @@ class GeminiModelInfo(BaseLLMModelInfo): def get_base_model(model: str) -> str | None: return model.replace("gemini/", "") - def process_model_name(self, models: list[dict[str, str]]) -> list[str]: + def process_model_name(self, models: Iterable[Mapping[str, str]]) -> list[str]: litellm_model_names: Final = [] for model in models: stripped_model_name = model["name"].replace("models/", "") @@ -416,9 +416,11 @@ class GeminiModelInfo(BaseLLMModelInfo): f"Failed to fetch models from Gemini. Status code: {response.status_code}, Response: {response.json()}" ) - models: Final[list[dict[str, str]]] = response.json()["models"] + listing: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + Mapping[str, Iterable[Mapping[str, str]]], _JSON_VALUE.validate_python(response.json()) + ) - litellm_model_names: Final = self.process_model_name(models) + litellm_model_names: Final = self.process_model_name(listing["models"]) return litellm_model_names def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py index 44cdac11325..34007bc8fd4 100644 --- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -15,7 +15,10 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamTransformSink, @@ -28,6 +31,7 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({}) +_DECODED_JSON: Final = TypeAdapter(object) def _field(container: object, name: str) -> object | None: @@ -108,12 +112,10 @@ def _write_back_texts(text_parts: Sequence[object], guardrailed_texts: Sequence[ def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: try: - parsed: Final = json.loads(payload) + parsed: Final = _DECODED_JSON.validate_python(json.loads(payload)) except json.JSONDecodeError: return None - if isinstance(parsed, dict): - return parsed - return None + return as_str_mapping(parsed) def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index bc3180a2129..0e29df36e80 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -3,7 +3,7 @@ Transformation for Calling Google models in their native format. """ from copy import deepcopy -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast import httpx @@ -36,6 +36,14 @@ else: from ..common_utils import get_api_key_from_env +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ Configuration for calling Google models in their native format. @@ -382,7 +390,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): from litellm.types.google_genai.main import GenerateContentResponse try: - response = raw_response.json() + decoded: Final = _json_body(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error transforming generate content response: {e}", @@ -391,7 +399,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): ) logging_obj.model_call_details["httpx_response"] = raw_response - response = self.convert_citation_sources_to_citations(response) + response: Final = self.convert_citation_sources_to_citations( + cast(dict[str, object], decoded) # cast-ok: the body passes through unvalidated, as before it was typed + ) return GenerateContentResponse(**response) diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 450c38a2206..d48a66f1f93 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -1,12 +1,13 @@ import base64 from collections.abc import Mapping from io import BufferedReader, BytesIO -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx from httpx._types import RequestFiles from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.common_utils import ( @@ -18,6 +19,7 @@ from litellm.llms.gemini.image_usage_transformation import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.llms.vertex_ai import GenerateContentResponseBody from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( FileTypes, @@ -34,6 +36,14 @@ else: LiteLLMLoggingObj = Any +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_body(response: _JsonResponse) -> object: + return response.json() + + class GeminiImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" SUPPORTED_PARAMS: list[str] = ["n", "size", "imageConfig"] @@ -125,7 +135,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_json: Final = raw_response.json() + response_json: Final = _json_body(raw_response) except Exception as exc: raise self.get_error_class( error_message=f"Error transforming image edit response: {exc}", @@ -133,7 +143,10 @@ class GeminiImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) - candidates: Final = response_json.get("candidates", []) + payload: Final = cast( # cast-ok: the body passes through unvalidated, as before it was typed + GenerateContentResponseBody, response_json + ) + candidates: Final = payload.get("candidates", []) data_list: Final[list[ImageObject]] = [] for candidate in candidates: @@ -141,17 +154,18 @@ class GeminiImageEditConfig(BaseImageEditConfig): parts = content.get("parts", []) for part in parts: inline_data = part.get("inlineData") - if inline_data and inline_data.get("data"): + image_data = inline_data.get("data") if inline_data else None + if image_data: data_list.append( ImageObject( - b64_json=inline_data["data"], + b64_json=image_data, url=None, ) ) model_response.data = cast(list[OpenAIImage], data_list) - if "usageMetadata" in response_json: - model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) + if "usageMetadata" in payload: + model_response.usage = transform_gemini_image_usage(dict(decoded_json_object(payload["usageMetadata"]))) return model_response def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, object]]: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 1aa3d074b1a..1f986b11af2 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -37,6 +37,7 @@ else: _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) ) @@ -188,7 +189,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): Transform Google AI Imagen response to litellm ImageResponse format """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -226,9 +227,9 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): ) # Extract usage metadata for Gemini models - if "usageMetadata" in response_data: + if "usageMetadata" in response_object: model_response.usage = transform_gemini_image_usage( - _JSON_DICT.validate_python(response_data["usageMetadata"]) + _JSON_DICT.validate_python(response_object["usageMetadata"]) ) web_search_requests: Final = get_gemini_image_web_search_requests(response_object) if web_search_requests and model_response.usage is not None: diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 9af7edeac86..6f778ab2761 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -5,13 +5,14 @@ This file contains the transformation logic for the Gemini realtime API. import json from collections import OrderedDict from collections.abc import Mapping, Sequence -from typing import Any, Final, cast +from typing import Final, cast from typing_extensions import ReadOnly, Required, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig @@ -167,16 +168,17 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return True @staticmethod - def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: - if not isinstance(details, dict): + def _usage_detail_alias(details: object, defaults: dict[str, int]) -> dict[str, object]: + detail_mapping: Final = as_str_mapping(details) + if detail_mapping is None or not isinstance(details, dict): return dict(defaults) return { **defaults, - **{key: value for key, value in details.items() if value is not None}, + **{key: value for key, value in detail_mapping.items() if value is not None}, } @staticmethod - def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, Any]) -> dict[str, object]: + def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, object]) -> dict[str, object]: usage_dict.setdefault( "input_token_details", GeminiRealtimeConfig._usage_detail_alias( diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 0656095927c..3f6b900fd01 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles +from pydantic import TypeAdapter import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS @@ -37,6 +38,8 @@ else: LiteLLMLoggingObj = Any BaseLLMException = Any +_DECODED_JSON: Final = TypeAdapter(object) + def _convert_image_to_gemini_format(image_file) -> dict[str, str]: """ @@ -60,7 +63,7 @@ def _convert_image_to_gemini_format(image_file) -> dict[str, str]: def _json_payload(raw_response: httpx.Response) -> object: """Read an HTTP response body as an opaque JSON payload.""" - return raw_response.json() + return _DECODED_JSON.validate_python(raw_response.json()) def _usage_video_resolution_from_parameters( diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 936e6339865..65921b628fb 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -273,7 +273,7 @@ class GigaChatConfig(BaseConfig): # Default to None (don't set function_call) return None - def _upload_image(self, image_url: str) -> str | None: + def _upload_image(self, image_url: str) -> object: """ Upload image to GigaChat and return file_id. @@ -293,7 +293,7 @@ class GigaChatConfig(BaseConfig): verbose_logger.error("Failed to upload image: %s", e) return None - def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[str]]: + def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[object]]: """ Extract text and image attachments from a multimodal message content list. @@ -304,7 +304,7 @@ class GigaChatConfig(BaseConfig): Tuple of (combined text, list of attachment file ids) """ texts: Final[list[str]] = [] # mutable-ok: accumulator - attachments: Final[list[str]] = [] # mutable-ok: accumulator + attachments: Final[list[object]] = [] # mutable-ok: accumulator for part in content: if isinstance(part, dict): if part.get("type") == "text": diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index aff0465d1ec..d53ae3af334 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -12,6 +12,8 @@ import uuid from collections.abc import Mapping from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( @@ -19,12 +21,14 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.llms.gigachat.utils import get_api_base +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.utils import LlmProviders from .authenticator import get_access_token, get_access_token_async # Simple in-memory cache for file IDs -_file_cache: Final[dict[str, str]] = {} +_file_cache: Final[dict[str, object]] = {} +_DECODED_JSON: Final = TypeAdapter(object) def _get_url_hash(url: str) -> str: @@ -91,7 +95,7 @@ def upload_file_sync( credentials: str | None = None, api_base: str | None = None, litellm_params: Mapping[str, object] | None = None, -) -> str | None: +) -> object: """ Upload file to GigaChat and return file_id (sync). @@ -137,9 +141,9 @@ def upload_file_sync( timeout=60, ) response.raise_for_status() - result: Final = response.json() + result: Final = _DECODED_JSON.validate_python(response.json()) - file_id: Final = result.get("id") + file_id: Final = decoded_json_object(result).get("id") if file_id: _file_cache[url_hash] = file_id verbose_logger.debug("File uploaded successfully, file_id: %s", file_id) @@ -156,7 +160,7 @@ async def upload_file_async( credentials: str | None = None, api_base: str | None = None, litellm_params: Mapping[str, object] | None = None, -) -> str | None: +) -> object: """ Upload file to GigaChat and return file_id (async). @@ -205,9 +209,9 @@ async def upload_file_async( timeout=60, ) response.raise_for_status() - result: Final = response.json() + result: Final = _DECODED_JSON.validate_python(response.json()) - file_id: Final = result.get("id") + file_id: Final = decoded_json_object(result).get("id") if file_id: _file_cache[url_hash] = file_id verbose_logger.debug("File uploaded successfully, file_id: %s", file_id) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 872eee6dbf9..27f75eece50 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -5,11 +5,13 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import httpx +from pydantic import TypeAdapter from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.gigachat.authenticator import get_access_token from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator from litellm.llms.gigachat.utils import GIGACHAT_BASE_URL +from litellm.llms.openai.common_utils import decoded_json_object from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse @@ -20,6 +22,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import CostResponseTypes +_DECODED_JSON: Final = TypeAdapter(object) + class GigaChatPassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: @@ -161,7 +165,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): continue chunk = chunk.removeprefix("data: ") try: - message = json.loads(chunk) + message = _DECODED_JSON.validate_python(json.loads(chunk)) except json.JSONDecodeError: continue @@ -169,7 +173,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): streaming_response=None, sync_stream=False, ) - translated_chunk = gigachat_iterator.chunk_parser(chunk=message) + translated_chunk = gigachat_iterator.chunk_parser(chunk=decoded_json_object(message)) if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields( # pyright: ignore[reportUnnecessaryIsInstance] # runtime guard for patched chunk_parser dict(translated_chunk) diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 9553327db97..83bc8df23a6 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -4,9 +4,11 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # spread dict loses the role literal without a cast import httpx +from pydantic import TypeAdapter from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.openai.common_utils import decoded_json_object from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ModelResponse @@ -23,6 +25,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding +_DECODED_JSON: Final = TypeAdapter(object) + class GithubCopilotConfig(OpenAIConfig): def __init__( @@ -309,7 +313,7 @@ class GithubCopilotConfig(OpenAIConfig): json_mode: bool | None = None, ) -> "ModelResponse": try: - response_json = raw_response.json() + decoded_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: return super().transform_response( model=model, @@ -325,12 +329,13 @@ class GithubCopilotConfig(OpenAIConfig): json_mode=json_mode, ) + response_json: Final = decoded_json_object(decoded_response) if not response_json.get("choices"): - response_json = self._synthesize_choices_for_anthropic_native(response_json) + synthesized: Final = self._synthesize_choices_for_anthropic_native(dict(response_json)) raw_response = httpx.Response( status_code=raw_response.status_code, headers=raw_response.headers, - content=json.dumps(response_json).encode(), + content=json.dumps(synthesized).encode(), ) return super().transform_response( diff --git a/litellm/llms/github_copilot/messages/transformation.py b/litellm/llms/github_copilot/messages/transformation.py index 56c3693a02c..613b7e42100 100644 --- a/litellm/llms/github_copilot/messages/transformation.py +++ b/litellm/llms/github_copilot/messages/transformation.py @@ -1,4 +1,5 @@ -from typing import Any, Final, cast # noqa: TID251 # narrows untyped request dicts at the session boundary +from collections.abc import Sequence +from typing import Final, cast # noqa: TID251 # narrows untyped request dicts at the session boundary from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.url_utils import ensure_api_route @@ -54,7 +55,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index bf0b8c79068..b055d388dd5 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -9,6 +9,7 @@ https://github.com/caozhiyuan/copilot-api """ import os +from collections.abc import Mapping from typing import ( TYPE_CHECKING, Any, @@ -16,6 +17,8 @@ from typing import ( cast, # noqa: TID251 # narrows untyped request dicts at the session boundary ) +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -47,6 +50,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_OBJECT_LIST: Final = TypeAdapter(list[object]) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[object, object]) + def github_copilot_supports_responses_api(model: str) -> bool: """ @@ -387,7 +393,9 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ return self._contains_vision_content(input_param) - def _contains_vision_content(self, value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bool: + def _contains_vision_content( + self, value: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH + ) -> bool: """ Recursively check if a value contains vision content. @@ -404,22 +412,28 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Check arrays if isinstance(value, list): - return any(self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in value) + items: Final = _OBJECT_LIST.validate_python(value) + return any(self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in items) # Only check dict/object types if not isinstance(value, dict): return False + fields: Final = _OBJECT_MAPPING.validate_python(value) + # Check if this item is an input_image - item_type: Final = value.get("type") + item_type: Final = fields.get("type") if isinstance(item_type, str) and item_type.lower() == "input_image": return True # Check content field recursively - if "content" in value and isinstance(value["content"], list): - return any( - self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in value["content"] - ) + if "content" in fields: + content: Final = fields["content"] + if isinstance(content, list): + content_items: Final = _OBJECT_LIST.validate_python(content) + return any( + self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in content_items + ) return False diff --git a/litellm/llms/gradient_ai/chat/transformation.py b/litellm/llms/gradient_ai/chat/transformation.py index 0ba5ca45c2a..2a2368ad8de 100644 --- a/litellm/llms/gradient_ai/chat/transformation.py +++ b/litellm/llms/gradient_ai/chat/transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final, Literal from litellm.litellm_core_utils.url_utils import join_api_base @@ -43,7 +45,7 @@ class GradientAIConfig(OpenAILikeChatConfig): include_guardrails_info: bool | None = None, provide_citations: bool | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 8c269cb5dfd..349a4b4ab87 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -40,6 +40,7 @@ class GroqExecutedToolIdentity(LiteLLMBaseModel): _EXECUTED_TOOLS_ADAPTER: Final = TypeAdapter(tuple[GroqExecutedToolIdentity, ...]) +_ATTRIBUTE_VALUE: Final = TypeAdapter(object) class GroqChatConfig(OpenAILikeChatConfig): @@ -324,7 +325,7 @@ class GroqChatConfig(OpenAILikeChatConfig): json_mode=json_mode, ) - original_service_tier: Final[object] = getattr(model_response, "service_tier") + original_service_tier: Final = _ATTRIBUTE_VALUE.validate_python(getattr(model_response, "service_tier")) mapped_service_tier: Final[Literal["auto", "default", "flex"]] = self._map_groq_service_tier( original_service_tier=original_service_tier ) diff --git a/litellm/llms/groq/cost_calculator.py b/litellm/llms/groq/cost_calculator.py index bcebb347aea..7f75a607410 100644 --- a/litellm/llms/groq/cost_calculator.py +++ b/litellm/llms/groq/cost_calculator.py @@ -19,7 +19,10 @@ if TYPE_CHECKING: def cost_per_web_search_request(usage: Usage, model_info: "ModelInfo") -> float: search_costs: Final = model_info.get("search_context_cost_per_query") cost_per_search: Final = search_costs.get("search_context_size_medium", 0.0) if search_costs else 0.0 - server_tool_use: Final = getattr(usage, "server_tool_use", None) + try: + server_tool_use: Final = usage.server_tool_use + except AttributeError: + return 0.0 if server_tool_use is None: return 0.0 searches: Final = server_tool_use.web_search_requests or 0 diff --git a/litellm/llms/groq/stt/transformation.py b/litellm/llms/groq/stt/transformation.py index 93d36dd7b0f..33e01ab7c0e 100644 --- a/litellm/llms/groq/stt/transformation.py +++ b/litellm/llms/groq/stt/transformation.py @@ -3,6 +3,7 @@ Translate from OpenAI's `/v1/audio/transcriptions` to Groq's `/v1/audio/transcri """ import types +from collections.abc import Mapping from typing import Final import litellm @@ -39,16 +40,17 @@ class GroqSTTConfig: tools: list | None = None, tool_choice: str | dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 4a5729945b4..0b77a04c69a 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -7,7 +7,7 @@ from types import MappingProxyType from typing import Final import httpx -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from litellm._uuid import uuid from litellm.exceptions import UnsupportedParamsError @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.llms.openai.common_utils import decoded_json_object from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( HostedVLLMRerankTruncationParams, @@ -28,6 +29,8 @@ from litellm.types.rerank import ( RerankTokens, ) +_DECODED_JSON: Final = TypeAdapter(object) + class HostedVLLMRerankError(BaseLLMException): def __init__( @@ -193,11 +196,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig): Process response from Hosted VLLM rerank API """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise ValueError(f"Error parsing response: {raw_response.text}, status_code={raw_response.status_code}") - return self._transform_response(raw_response_json) + return self._transform_response(dict(decoded_json_object(raw_response_json))) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/hosted_vllm/videos/transformation.py b/litellm/llms/hosted_vllm/videos/transformation.py index 30ae6835263..2045d2c8988 100644 --- a/litellm/llms/hosted_vllm/videos/transformation.py +++ b/litellm/llms/hosted_vllm/videos/transformation.py @@ -8,6 +8,7 @@ from typing import Final from urllib.parse import urlparse from httpx._types import FileTypes, RequestFiles +from pydantic import TypeAdapter from litellm.images.utils import ImageEditRequestUtils from litellm.litellm_core_utils.url_utils import SSRFError, join_api_base, validate_url @@ -16,6 +17,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoCreateOptionalRequestParams +_DECODED_JSON: Final = TypeAdapter(object) _EXCLUDED_FORM_KEYS: Final = frozenset( { "model", @@ -82,7 +84,7 @@ def _maybe_json(value: object) -> object: stripped: Final = value.strip() if not stripped or stripped[0] not in "{[": return value - return json.loads(stripped) + return _DECODED_JSON.validate_python(json.loads(stripped)) def _reject_unsafe_media_url(url: str) -> None: diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 60917c68221..a1d1a3f6a19 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -3,6 +3,7 @@ import os from typing import Final, Literal, Protocol, get_args import httpx +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -11,6 +12,7 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, ) +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.utils import EmbeddingResponse from ...base import BaseLLM @@ -21,6 +23,8 @@ config: Final = HuggingFaceEmbeddingConfig() HF_HUB_URL: Final = "https://huggingface.co" +_DECODED_JSON: Final = TypeAdapter(object) + hf_tasks_embeddings: Final = ( Literal[ # pipeline tags + hf tei endpoints - https://huggingface.github.io/text-embeddings-inference/#/ "sentence-similarity", "feature-extraction", "rerank", "embed", "similarity" @@ -34,7 +38,7 @@ class _SupportsTokenEncode(Protocol): def encode(self, text: str) -> list[int]: ... -def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> str | None: +def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -44,14 +48,14 @@ def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}") - model_info_dict: Final = model_info.json() + model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) - pipeline_tag: Final[str | None] = model_info_dict.get("pipeline_tag", None) + pipeline_tag: Final = model_info_dict.get("pipeline_tag", None) return pipeline_tag -async def async_get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> str | None: +async def async_get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -63,9 +67,9 @@ async def async_get_hf_task_embedding_for_model(model: str, task_type: str | Non model_info: Final = await http_client.get(url=f"{api_base}/api/models/{model}") - model_info_dict: Final = model_info.json() + model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) - pipeline_tag: Final[str | None] = model_info_dict.get("pipeline_tag", None) + pipeline_tag: Final = model_info_dict.get("pipeline_tag", None) return pipeline_tag @@ -77,7 +81,7 @@ class HuggingFaceEmbedding(BaseLLM): def __init__(self) -> None: super().__init__() - def _transform_input_on_pipeline_tag(self, input: list, pipeline_tag: str | None) -> dict: + def _transform_input_on_pipeline_tag(self, input: list, pipeline_tag: object) -> dict: if pipeline_tag is None: return {"inputs": input} if pipeline_tag == "sentence-similarity" or pipeline_tag == "similarity": diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 0883a9a8077..1b7d6483739 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -447,14 +447,13 @@ class HuggingFaceEmbeddingConfig(BaseConfig): except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails pass - output_text: Final = model_response["choices"][0]["message"].get("content", "") + first_message: Final = model_response.choices[0].message + output_text: Final = first_message.content if output_text is not None and len(output_text) > 0: completion_tokens = 0 try: if encoding is not None: - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) - ) ##[TODO] use the llama2 tokenizer here + completion_tokens = len(encoding.encode(output_text)) ##[TODO] use the llama2 tokenizer here except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails pass diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 1217a2540d4..1f130f84283 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -86,7 +86,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index 0938cc0d07b..3807be562cf 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -1,9 +1,11 @@ """LangFlow run API: POST {api_base}/api/v1/run/{flow_id}""" +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from urllib.parse import quote import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -11,6 +13,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -27,6 +30,8 @@ else: AsyncHTTPHandler = Any CustomStreamWrapper = Any +_DECODED_JSON: Final = TypeAdapter(object) + class LangFlowError(BaseLLMException): """Exception class for LangFlow API errors.""" @@ -179,7 +184,7 @@ class LangFlowConfig(BaseConfig): verbose_logger.debug("LangFlow request payload: %s", payload) return payload - def _extract_content_from_response(self, response_json: dict) -> str | None: + def _extract_content_from_response(self, response_json: Mapping[str, object]) -> str | None: """ Extract the assistant text from a LangFlow run response. @@ -236,7 +241,7 @@ class LangFlowConfig(BaseConfig): json_mode: bool | None = None, ) -> ModelResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise LangFlowError( message=f"LangFlow returned a non-JSON response: {e}", @@ -245,7 +250,7 @@ class LangFlowConfig(BaseConfig): verbose_logger.debug("LangFlow response: %s", response_json) - content: Final = self._extract_content_from_response(response_json) + content: Final = self._extract_content_from_response(decoded_json_object(response_json)) if content is None: raise LangFlowError( message=( diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 499cbab3b7b..575820bceee 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -9,10 +9,13 @@ import uuid from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices +_DECODED_JSON: Final = TypeAdapter(object) + class LangGraphSSEStreamIterator: """ @@ -60,7 +63,7 @@ class LangGraphSSEStreamIterator: return None try: - data: Final = json.loads(json_str) + data: Final = _DECODED_JSON.validate_python(json.loads(json_str)) return self._process_data(data) except json.JSONDecodeError: verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100]) diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index ce624eecc0e..603d5a07292 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -13,6 +13,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional, Union import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -21,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator +from litellm.llms.openai.common_utils import decoded_json_object from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -37,6 +39,8 @@ else: AsyncHTTPHandler = Any CustomStreamWrapper = Any +_DECODED_JSON: Final = TypeAdapter(object) + class LangGraphError(BaseLLMException): """Exception class for LangGraph API errors.""" @@ -233,7 +237,7 @@ class LangGraphConfig(BaseConfig): verbose_logger.debug("LangGraph request payload: %s", payload) return payload - def _extract_content_from_response(self, response_json: dict) -> str: + def _extract_content_from_response(self, response_json: Mapping[str, object]) -> str: """ Extract content from LangGraph non-streaming response. @@ -434,10 +438,10 @@ class LangGraphConfig(BaseConfig): Transform the LangGraph response to LiteLLM ModelResponse format. """ try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) verbose_logger.debug("LangGraph response: %s", response_json) - content: Final = self._extract_content_from_response(response_json) + content: Final = self._extract_content_from_response(decoded_json_object(response_json)) # Create the message message: Final = Message(content=content, role="assistant") diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index 01fdc0b8bf8..42bd63b0160 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -3,7 +3,7 @@ Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completio """ from collections.abc import Iterable, Mapping -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import quote import httpx @@ -27,6 +27,7 @@ if TYPE_CHECKING: _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str) +_OPAQUE_VALUE: Final = TypeAdapter(object) class _JsonResponse(Protocol): @@ -123,7 +124,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): return ["lemonade/" + _STR.validate_python(model["id"]) for model in model_list] @staticmethod - def _get_positive_int(value: Any) -> int | None: + def _get_positive_int(value: object) -> int | None: if isinstance(value, bool): return None if isinstance(value, int) and value > 0: @@ -163,7 +164,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): provider_specific_entry.get("context_window"), provider_specific_entry.get("max_context_window"), ): - parsed = self._get_positive_int(value) + parsed = self._get_positive_int(_OPAQUE_VALUE.validate_python(value)) if parsed is not None: return parsed return None diff --git a/litellm/llms/litellm_proxy/skills/prompt_injection.py b/litellm/llms/litellm_proxy/skills/prompt_injection.py index 3377338a03e..e8e1c68a112 100644 --- a/litellm/llms/litellm_proxy/skills/prompt_injection.py +++ b/litellm/llms/litellm_proxy/skills/prompt_injection.py @@ -9,9 +9,10 @@ import posixpath import zipfile from collections.abc import Mapping from io import BytesIO -from typing import Any, Final +from typing import Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.models.skills import LiteLLM_SkillsTable @@ -251,27 +252,30 @@ class SkillPromptInjectionHandler: if len(description) > max_desc_length: description = description[: max_desc_length - 3] + "..." - tool: Final[dict[str, Any]] = { + params: Final = ( + as_str_mapping(skill.metadata.get("parameters")) + if skill.metadata and isinstance(skill.metadata, dict) + else None + ) + parameters: Final[Mapping[str, object]] = ( + params + if params and isinstance(params, dict) + else { + "type": "object", + "properties": {}, + "required": [], + } + ) + + return { "type": "function", "function": { "name": func_name, "description": description, - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, + "parameters": parameters, }, } - # If skill has metadata with parameter definitions, use them - if skill.metadata and isinstance(skill.metadata, dict): - params: Final = skill.metadata.get("parameters") - if params and isinstance(params, dict): - tool["function"]["parameters"] = params - - return tool - def convert_skill_to_anthropic_tool(self, skill: LiteLLM_SkillsTable) -> Mapping[str, object]: """ Convert a LiteLLM skill to an Anthropic-style tool (messages API format). diff --git a/litellm/llms/litellm_proxy/skills/sandbox_executor.py b/litellm/llms/litellm_proxy/skills/sandbox_executor.py index 3e38dd81905..8e5c9effa17 100644 --- a/litellm/llms/litellm_proxy/skills/sandbox_executor.py +++ b/litellm/llms/litellm_proxy/skills/sandbox_executor.py @@ -7,7 +7,8 @@ Supports Docker, Podman, and Kubernetes backends. import base64 import os -from typing import Any, Final, Protocol, TypedDict +from collections.abc import Sequence +from typing import Final, Protocol, TypedDict from typing_extensions import ReadOnly @@ -41,6 +42,13 @@ class _GeneratedFile(TypedDict): mime_type: ReadOnly[str] +class _SandboxExecution(TypedDict): + success: ReadOnly[bool] + output: ReadOnly[str] + error: ReadOnly[str] + files: ReadOnly[Sequence[_GeneratedFile]] + + class SkillsSandboxExecutor: """ Executes skill code in llm-sandbox Docker container. @@ -76,7 +84,7 @@ class SkillsSandboxExecutor: code: str, skill_files: dict[str, bytes], requirements: str | None = None, - ) -> dict[str, Any]: + ) -> _SandboxExecution: """ Execute code with skill files in sandbox. diff --git a/litellm/llms/lm_studio/embed/transformation.py b/litellm/llms/lm_studio/embed/transformation.py index 66b2dbb534d..d6264ceb6e0 100644 --- a/litellm/llms/lm_studio/embed/transformation.py +++ b/litellm/llms/lm_studio/embed/transformation.py @@ -7,6 +7,7 @@ Docs - https://lmstudio.ai/docs/basics/server """ import types +from collections.abc import Mapping from typing import Final @@ -18,16 +19,17 @@ class LmStudioEmbeddingConfig: def __init__( self, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/maritalk.py b/litellm/llms/maritalk.py index 094dee4cee1..36bde5847fb 100644 --- a/litellm/llms/maritalk.py +++ b/litellm/llms/maritalk.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final from httpx._models import Headers @@ -32,7 +33,7 @@ class MaritalkConfig(OpenAIGPTConfig): tools: list[dict] | None = None, tool_choice: str | dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 67de3a06af8..14470922e57 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -4,12 +4,14 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) +from litellm.llms.openai.common_utils import decoded_json_object from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -30,6 +32,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + MILVUS_OPTIONAL_PARAMS: Final = { "annsField", "limit", @@ -218,15 +222,17 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): results: Final = response_json.get("data", []) # Try to get text_field from optional_params first, then litellm_params - optional_params: Final = litellm_logging_obj.model_call_details.get("optional_params", {}) + optional_params: Final = decoded_json_object( + _DECODED_JSON.validate_python(litellm_logging_obj.model_call_details.get("optional_params", {})) + ) text_field = optional_params.get("milvus_text_field", "") # Fallback to litellm_params if not in optional_params if not text_field: - text_field = litellm_logging_obj.model_call_details.get("litellm_params", {}).get( - "milvus_text_field", "" - ) + text_field = decoded_json_object( + _DECODED_JSON.validate_python(litellm_logging_obj.model_call_details.get("litellm_params", {})) + ).get("milvus_text_field", "") # Transform results to standard format search_results: Final[list[VectorStoreSearchResult]] = [] diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index f9f3502237c..f7a263a467c 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -2,8 +2,6 @@ MiniMax Anthropic transformation config - extends AnthropicConfig for MiniMax's Anthropic-compatible API """ -from typing import Any # noqa: TID251 # override below must mirror the legacy base signature - import litellm from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.messages.transformation import ( @@ -54,7 +52,7 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): self, headers: dict, # mutable-ok: mirrors the legacy base override signature model: str, - messages: list[Any], # mutable-ok: mirrors the legacy base override signature + messages: list[object], # mutable-ok: mirrors the legacy base override signature optional_params: dict, # mutable-ok: mirrors the legacy base override signature litellm_params: dict, # mutable-ok: mirrors the legacy base override signature api_key: str | None = None, diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 36c8705f256..c308de92d1d 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -33,6 +33,7 @@ else: LiteLLMLoggingObj = object _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) @@ -173,7 +174,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): {"created": timestamp, "data": [{"url": "..."}]} """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing ModelScope response: {e}", diff --git a/litellm/llms/nadir/chat/transformation.py b/litellm/llms/nadir/chat/transformation.py index d22b1a51bab..c8bde7d9fad 100644 --- a/litellm/llms/nadir/chat/transformation.py +++ b/litellm/llms/nadir/chat/transformation.py @@ -2,8 +2,9 @@ import math from typing import Final import httpx +from pydantic import TypeAdapter -from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_response_cost_in_hidden_params from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -19,13 +20,17 @@ _SUPPORTED_OPENAI_PARAMS: Final = ( "temperature", "top_p", ) +_DECODED_JSON: Final = TypeAdapter(object) def _reported_cost_usd(raw_response: httpx.Response) -> float | None: try: - cost: Final = raw_response.json()["nadir_metadata"]["cost"]["total_cost_usd"] + body: Final = as_str_mapping(_DECODED_JSON.validate_python(raw_response.json())) except (ValueError, KeyError, TypeError): return None + metadata: Final = as_str_mapping(body.get("nadir_metadata")) if body is not None else None + cost_block: Final = as_str_mapping(metadata.get("cost")) if metadata is not None else None + cost: Final = cost_block.get("total_cost_usd") if cost_block is not None else None if isinstance(cost, bool) or not isinstance(cost, (int, float)): return None if not math.isfinite(cost) or cost < 0: diff --git a/litellm/llms/nvidia_nim/embed.py b/litellm/llms/nvidia_nim/embed.py index 5a0105a3776..d2ef446299e 100644 --- a/litellm/llms/nvidia_nim/embed.py +++ b/litellm/llms/nvidia_nim/embed.py @@ -9,6 +9,7 @@ API calling is done using the OpenAI SDK with an api_base """ import types +from collections.abc import Mapping from typing import Final @@ -32,16 +33,17 @@ class NvidiaNimEmbeddingConfig: input_type: str | None = None, truncate: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index 748203b22d9..cafeece0a4e 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -28,7 +28,7 @@ without the optional STT extras installed. import asyncio import inspect -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol @@ -48,6 +48,8 @@ from litellm.llms.nvidia_riva.audio_transcription.transformation import ( RIVA_TARGET_NUM_CHANNELS, RIVA_TARGET_SAMPLE_RATE_HZ, NvidiaRivaAudioTranscriptionConfig, + RivaFinalResult, + RivaWordTiming, ) from litellm.llms.nvidia_riva.common_utils import ( NvidiaRivaException, @@ -535,14 +537,14 @@ class NvidiaRivaAudioTranscription: yield chunk @staticmethod - def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]: + def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> Sequence[RivaFinalResult]: """ Walk the gRPC responses, ignore empty / non-final chunks, and return a list of normalized final-result dicts. Matching the user's note: the ``id`` blocks with no ``results`` are streaming heartbeats and must be skipped. """ - final_results: Final[list[dict[str, object]]] = [] + final_results: Final[list[RivaFinalResult]] = [] for response in stream: results = getattr(response, "results", None) or [] for result in results: @@ -554,16 +556,19 @@ class NvidiaRivaAudioTranscription: top = alternatives[0] transcript = getattr(top, "transcript", "") or "" words_proto = getattr(top, "words", None) or [] - words = [] - for word in words_proto: - words.append( - { - "word": getattr(word, "word", ""), - "start_time_ms": int(getattr(word, "start_time", 0) or 0), - "end_time_ms": int(getattr(word, "end_time", 0) or 0), - } - ) - final_results.append({"transcript": transcript, "words": words}) + final_results.append( + { + "transcript": transcript, + "words": [ + RivaWordTiming( + word=getattr(word, "word", ""), + start_time_ms=int(getattr(word, "start_time", 0) or 0), + end_time_ms=int(getattr(word, "end_time", 0) or 0), + ) + for word in words_proto + ], + } + ) return final_results diff --git a/litellm/llms/nvidia_riva/audio_transcription/transformation.py b/litellm/llms/nvidia_riva/audio_transcription/transformation.py index 262dfb7bd78..5862a5d10b1 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/transformation.py +++ b/litellm/llms/nvidia_riva/audio_transcription/transformation.py @@ -11,9 +11,11 @@ dict at call time. Reference: https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-overview.html """ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final, TypedDict from httpx import Headers, Response +from typing_extensions import ReadOnly from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import ( @@ -34,6 +36,17 @@ RIVA_TARGET_NUM_CHANNELS: Final = 1 RIVA_TARGET_ENCODING: Final = "LINEAR_PCM" +class RivaWordTiming(TypedDict): + word: ReadOnly[str] + start_time_ms: ReadOnly[int] + end_time_ms: ReadOnly[int] + + +class RivaFinalResult(TypedDict): + transcript: ReadOnly[str] + words: ReadOnly[Sequence[RivaWordTiming]] + + class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """ Config for NVIDIA Riva ASR (gRPC). @@ -219,7 +232,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): @staticmethod def build_transcription_response( - final_results: list[dict[str, Any]], + final_results: Sequence[RivaFinalResult], response_format: str, duration_seconds: float | None, timestamp_granularities: list[str] | None, diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 8e1502b2610..e8d37f857e6 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -189,6 +189,7 @@ def _normalize_tool_choice(selected_params: dict) -> None: _TOOL_CHOICE_MODES: Final = frozenset({"auto", "none", "required", "any"}) _TOOLS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) _JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) def _as_json_object(value: object) -> Mapping[str, object] | None: @@ -701,7 +702,7 @@ class OCIChatConfig(BaseConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) if not isinstance(response_json, dict): raise OCIError( diff --git a/litellm/llms/oci/embed/transformation.py b/litellm/llms/oci/embed/transformation.py index ffcb156bb16..6dee64bf19a 100644 --- a/litellm/llms/oci/embed/transformation.py +++ b/litellm/llms/oci/embed/transformation.py @@ -56,6 +56,7 @@ else: LiteLLMLoggingObj = Any _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) # OCI sends up to 96 texts per embedText request (Cohere limit). OCI_EMBED_BATCH_LIMIT: Final = 96 @@ -272,7 +273,7 @@ class OCIEmbedConfig(BaseEmbeddingConfig): ) try: - json_response: Final = raw_response.json() + json_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OCIError( status_code=raw_response.status_code, diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 7894904c870..f313c62f026 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -13,6 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ModelResponse, ModelResponseStream _PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) class OllamaResponseIdNormalizer: @@ -152,7 +153,7 @@ class OllamaModelInfo(BaseLLMModelInfo): try: resp: Final = httpx.get(join_api_base(base, "api/tags"), headers=headers) resp.raise_for_status() - data: Final = resp.json() + data: Final = _DECODED_JSON.validate_python(resp.json()) # Expecting a dict with a 'models' list models_list = [] if isinstance(data, dict) and "models" in data and isinstance(data["models"], list): diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 30de22b17ff..239b4e32645 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -547,7 +547,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): held: Final = self.held_content self.held_content = "" try: - parsed: Final[object] = json.loads(_strip_json_code_fence(held)) + parsed: Final = _JSON_VALUE.validate_python(json.loads(_strip_json_code_fence(held))) except json.JSONDecodeError: return GenericStreamingChunk(text=held, is_finished=True, finish_reason="stop", usage=usage) if isinstance(parsed, dict) and "name" in parsed and "arguments" in parsed: diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index e472661efb9..12ead582975 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -23,13 +23,14 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta from openai.types.completion_usage import CompletionUsage -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter if TYPE_CHECKING: from aiohttp import ClientSession import litellm from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object from litellm.litellm_core_utils.token_counter import token_counter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -134,6 +135,9 @@ class OpenAIError(BaseLLMException): ) +_DECODED_JSON: Final = TypeAdapter(object) + + ####### Error Handling Utils for OpenAI API ####################### ################################################################### def drop_params_from_unprocessable_entity_error( @@ -152,19 +156,20 @@ def drop_params_from_unprocessable_entity_error( """ invalid_params: Final[list[str]] = [] if isinstance(e, httpx.HTTPStatusError): - error_json: Final = e.response.json() + error_json: Final = decoded_json_object(_DECODED_JSON.validate_python(e.response.json())) error_message: Final = error_json.get("error", {}) error_body = error_message else: error_body = e.body - if error_body is not None and isinstance(error_body, dict) and error_body.get("message"): - message = error_body.get("message", {}) + error_mapping: Final = as_str_mapping(error_body) + if isinstance(error_mapping, dict) and error_mapping.get("message"): + message = error_mapping.get("message", {}) if isinstance(message, str): try: - message = json.loads(message) + message = _DECODED_JSON.validate_python(json.loads(message)) except json.JSONDecodeError: message = {"detail": message} - detail: Final = message.get("detail") + detail: Final = decoded_json_object(message).get("detail") if isinstance(detail, list) and len(detail) > 0 and isinstance(detail[0], dict): for error_dict in detail: diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index 581e5e8ad87..73d211e1613 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -2,10 +2,10 @@ OpenAI Evals API configuration and transformations """ -from collections.abc import Mapping from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base @@ -32,9 +32,11 @@ from litellm.types.llms.openai_evals import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +_DECODED_JSON: Final = TypeAdapter(object) -def _parsed_response_json(raw_response: httpx.Response) -> Mapping[str, object]: - return raw_response.json() + +def _parsed_response_json(raw_response: httpx.Response) -> object: + return _DECODED_JSON.validate_python(raw_response.json()) class OpenAIEvalsConfig(BaseEvalsAPIConfig): diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 437f19a4dc6..deee502c9f2 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -20,7 +20,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.transformation impor LiteLLMAnthropicToResponsesAPIAdapter, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials +from litellm.llms.openai.common_utils import OpenAIError, decoded_json_object, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) @@ -244,7 +244,7 @@ class OpenAITokenCounter(BaseTokenCounter): if result is not None: return TokenCountResponse.model_validate( { - "total_tokens": result.get("input_tokens", 0), + "total_tokens": decoded_json_object(result).get("input_tokens", 0), "request_model": request_model, "model_used": model_to_use, "tokenizer_type": "openai_api", diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index c1ee42444c5..0f885146e36 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -6,12 +6,13 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp import copy from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal, TypeAlias +from typing import Final, Literal, Protocol, TypeAlias from openai.types.responses.input_token_count_params import InputTokenCountParams from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.convertible_blocks import malformed_tool_result_text_error @@ -61,6 +62,26 @@ class _AnthropicCountMessage(TypedDict): _ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) + + +class _FieldReader(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +class _FunctionFields(_FieldReader, Protocol): + def __contains__(self, key: object, /) -> bool: ... + + def __getitem__(self, key: str, /) -> object: ... + + +def _field_reader(value: object) -> _FieldReader: + return require_attributes(value, _FieldReader, "get") + + +def _function_fields(value: object) -> _FunctionFields: + return require_attributes(value, _FunctionFields, "get", "__contains__", "__getitem__") + + _ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) @@ -196,7 +217,7 @@ class OpenAICountTokensConfig: @staticmethod def _transform_tools_for_responses_api( - tools: list[dict[str, Any]], + tools: list[dict[str, object]], ) -> list[dict[str, object]]: """ Transform OpenAI chat and Anthropic tools to Responses API tools format. @@ -208,7 +229,7 @@ class OpenAICountTokensConfig: transformed: Final[list[dict[str, object]]] = [] for tool in tools: if tool.get("type") == "function" and "function" in tool: - func = tool["function"] + func = _function_fields(tool["function"]) item: dict[str, object] = { "type": "function", "name": func.get("name", ""), @@ -235,7 +256,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( - messages: list[dict[str, Any]], + messages: Sequence[Mapping[str, object]], ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -276,14 +297,17 @@ class OpenAICountTokensConfig: if content: input_items.append({"role": "assistant", "content": content}) if tool_calls: - for tc in tool_calls: - func = tc.get("function", {}) + for tc in iter_items(tool_calls): + tool_call = _field_reader(tc) + function = tool_call.get("function", {}) + call_id = tool_call.get("id", "") + function_fields = _field_reader(function) input_items.append( { "type": "function_call", - "call_id": tc.get("id", ""), - "name": func.get("name", ""), - "arguments": func.get("arguments", ""), + "call_id": call_id, + "name": function_fields.get("name", ""), + "arguments": function_fields.get("arguments", ""), } ) elif not content: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 3ea126c5a2a..d37a12a91d6 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -36,7 +36,7 @@ from collections.abc import Mapping, MutableSequence, Sequence from dataclasses import dataclass from itertools import accumulate, chain, repeat from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, Union, cast from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -47,6 +47,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i OpenAiResponsesToChatCompletionStreamIterator, tool_call_dict_from_output_item, ) +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -237,6 +238,10 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +class _StreamEventFields(Protocol): + def get(self, key: str, /) -> object: ... + + _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( { ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, @@ -1264,7 +1269,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _sync_stream_events_with_rewrites( self, - stream_events: Sequence[Any], + stream_events: Sequence[object], rewrites_by_position: Mapping[tuple[int, int], str], ) -> None: """Sync pre-completion stream events with the rewritten completed @@ -1278,7 +1283,7 @@ class OpenAIResponsesHandler(BaseTranslation): {position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()} ) for event in stream_events: - if not (isinstance(event, dict) or hasattr(event, "get")): + if not has_attributes(event, _StreamEventFields, "get"): continue event_type = event.get("type") output_index = event.get("output_index") diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 12c8b1da166..0ab55d653e0 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final import litellm from litellm.litellm_core_utils.url_utils import join_api_base @@ -41,7 +41,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): self, headers: dict[str, str], model: str, - messages: list[Any], + messages: list[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -151,7 +151,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): self, headers: dict[str, str], model: str, - messages: list[Any], + messages: list[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/opencode/harness/transformation.py b/litellm/llms/opencode/harness/transformation.py index 46e8b750ca5..e801b66c448 100644 --- a/litellm/llms/opencode/harness/transformation.py +++ b/litellm/llms/opencode/harness/transformation.py @@ -26,10 +26,10 @@ import itertools import json import secrets import string -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError @@ -45,6 +45,7 @@ from litellm.harness.types import ( ToolCall, ToolResult, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.base_llm.harness.transformation import ( BaseCLIHarnessConfig, HarnessSessionCopy, @@ -269,10 +270,17 @@ def session_copy(exported: str, path: str, workdir: str) -> HarnessSessionCopy | ) -def _as_dict(value: object) -> Mapping[str, Any]: +def _as_dict(value: object) -> Mapping[str, object]: return value if isinstance(value, dict) else MappingProxyType({}) +def _spread_items(value: object) -> Iterable[object]: + try: + return iter_items(value) + except TypeError: + raise TypeError(f"Value after * must be an iterable, not {type(value).__name__}") from None + + def _tool_events(part: Mapping[str, object]) -> Sequence[Event]: native = str(part.get("tool") or "") call_id = str(part.get("callID") or part.get("id") or "") @@ -346,7 +354,7 @@ def build_opencode_config( token_path: str, permissions: PermissionMode, disable_tools: Sequence[str] = (), - user_config: Mapping[str, Any] | None = None, + user_config: Mapping[str, object] | None = None, instructions_path: str | None = None, skills_path: str | None = None, ) -> Mapping[str, object]: @@ -355,10 +363,16 @@ def build_opencode_config( validate_user_config(user) qualified = f"{OPENCODE_PROVIDER_ID}/{model}" extra_instructions: Final = (instructions_path,) if instructions_path else () - instructions: Final = [*(user.get("instructions") or ()), *extra_instructions] # mutable-ok: opencode config JSON + instructions: Final = [ + *_spread_items(user.get("instructions") or ()), + *extra_instructions, + ] # mutable-ok: opencode config JSON user_skills = _as_dict(user.get("skills")) extra_skills: Final = (skills_path,) if skills_path else () - skill_paths: Final = [*(user_skills.get("paths") or ()), *extra_skills] # mutable-ok: opencode config JSON + skill_paths: Final = [ + *_spread_items(user_skills.get("paths") or ()), + *extra_skills, + ] # mutable-ok: opencode config JSON options: Final = {"baseURL": base_url, "apiKey": "{file:" + token_path + "}"} # mutable-ok: opencode config JSON models: Final[dict[str, Mapping[str, object]]] = {model: {}} # mutable-ok: opencode config JSON provider: Final = { # mutable-ok: opencode config JSON diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index 9644ee3f280..c2d8398aead 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -47,7 +47,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles -from pydantic import ConfigDict, RootModel, with_config +from pydantic import ConfigDict, RootModel, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -74,6 +74,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) + @with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) class _ImageUrl(TypedDict): @@ -248,7 +250,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: try: - response_json: Final[object] = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OpenRouterException( message=f"Error parsing OpenRouter response: {e}", diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index ce6e62aaa88..a9e07555d6a 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -59,6 +59,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) @@ -359,7 +360,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ImageResponse: Populated image response """ try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise OpenRouterException( message=f"Error parsing OpenRouter response: {e}", diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index 574fbfa9e5b..d30a9ff534e 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -20,6 +20,7 @@ from litellm.constants import ( OPEN_SANDBOX_POLL_INTERVAL, OPEN_SANDBOX_READY_TIMEOUT, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.sandbox.transformation import ( @@ -44,6 +45,13 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_VALUE: Final = TypeAdapter(object) +def _decoded_json_object(value: object) -> Mapping[str, object]: + mapping: Final = as_str_mapping(value) + if mapping is not None: + return mapping + raise AttributeError(f"{type(value).__name__!r} object has no attribute 'get'", name="get", obj=value) + + class OpenSandboxSandboxConfig(BaseSandboxConfig): def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler: if client is not None: @@ -321,7 +329,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): headers=headers, params={"use_server_proxy": use_server_proxy}, ) - data: Final = response.json() + data: Final = _decoded_json_object(_JSON_VALUE.validate_python(response.json())) endpoint: Final = data.get("endpoint") if not endpoint: raise ValueError(f"OpenSandbox did not return an execd endpoint for {sandbox_id}") @@ -492,7 +500,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): if not data: return None try: - parsed: Final = json.loads(data) + parsed: Final = _JSON_VALUE.validate_python(json.loads(data)) except json.JSONDecodeError: return None if not isinstance(parsed, dict): diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 885fc98ffea..3ce8d7ddad5 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -6,11 +6,14 @@ It uses the field targeting configuration from litellm_logging_obj to extract specific fields for guardrail processing. """ -from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -215,6 +218,19 @@ class _EventStreamEndpointChecker(Protocol): def event_stream_endpoint_is_de_anonymizable(self, endpoint: str, model: str, body: object) -> bool: ... +class _EventStreamDeAnonymizer(Protocol): + async def de_anonymize_event_stream( + self, + body_bytes: bytes, + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: "UserAPIKeyAuth", + data: Mapping[str, object], + ) -> bytes: ... + + +_TRANSLATION_RESULT: Final = TypeAdapter(object) + + _PROVIDER_HANDLERS: dict[str, type[BaseTranslation]] = {} @@ -260,10 +276,12 @@ class LlmPassthroughRouteHandler(BaseTranslation): provider, ) return data - return await handler_cls().process_input_messages( - data=data, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, + return _TRANSLATION_RESULT.validate_python( + await handler_cls().process_input_messages( + data=data, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + ) ) async def process_output_response( @@ -282,12 +300,14 @@ class LlmPassthroughRouteHandler(BaseTranslation): provider, ) return response - return await handler_cls().process_output_response( - response=response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=request_data, + return _TRANSLATION_RESULT.validate_python( + await handler_cls().process_output_response( + response=response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) ) @staticmethod @@ -399,9 +419,11 @@ class LlmPassthroughRouteHandler(BaseTranslation): return handler_cls.event_stream_media_type() @staticmethod - def _resolve_event_stream_de_anonymizer(provider: str | None): + def _resolve_event_stream_de_anonymizer(provider: str | None) -> Callable[..., Awaitable[bytes]] | None: handler_cls: Final = _get_provider_handlers().get(provider or "") - return getattr(handler_cls, "de_anonymize_event_stream", None) + if not has_attributes(handler_cls, _EventStreamDeAnonymizer, "de_anonymize_event_stream"): + return None + return handler_cls.de_anonymize_event_stream @staticmethod def supports_event_stream_relay(provider: object, endpoint: object, model: object, body: object) -> bool: diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index a2b7eceea40..77817a5e5bb 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -12,6 +12,7 @@ Ref: https://docs.perplexity.ai/api-reference/responses-post from typing import Any, Final import httpx +from pydantic import TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base @@ -22,6 +23,8 @@ from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +_DECODED_JSON: Final = TypeAdapter(object) + class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def get_supported_openai_params(self, model: str) -> list: @@ -104,7 +107,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): ) -> ResponsesAPIResponse: """Check for Perplexity's status:'failed' on HTTP 200 before delegating to base.""" try: - raw_response_json = raw_response.json() + raw_response_json = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raw_response_json = None diff --git a/litellm/llms/petals/completion/transformation.py b/litellm/llms/petals/completion/transformation.py index ee20c2b12d5..a875992448a 100644 --- a/litellm/llms/petals/completion/transformation.py +++ b/litellm/llms/petals/completion/transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final from httpx import Headers, Response @@ -57,7 +59,7 @@ class PetalsConfig(BaseConfig): top_p: float | None = None, repetition_penalty: float | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index eeecd02653d..89851f60d76 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -76,7 +76,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = join_api_base(api_base, f"{encoded_vector_store_id}/search") diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index 7b1a7714ec6..49a2938e1c0 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -26,6 +26,7 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class RecraftImageGenerationConfig(BaseImageGenerationConfig): @@ -134,7 +135,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): https://www.recraft.ai/docs#generate-image """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index e721ac678a1..23c8c86e7ca 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -1,9 +1,10 @@ import asyncio import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -11,6 +12,7 @@ from litellm.constants import ( RUNWAYML_DEFAULT_API_VERSION, RUNWAYML_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.core_helpers import decoded_json_object from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -36,9 +38,25 @@ class _RunwayMLTask(TypedDict, total=False): id: ReadOnly[str] status: ReadOnly[str] - output: ReadOnly[Sequence[str | Mapping[str, str]]] - failure: ReadOnly[str] - failureCode: ReadOnly[str] + output: ReadOnly[Sequence[str | Mapping[str, str | None]] | None] + failure: ReadOnly[str | None] + failureCode: ReadOnly[str | None] + + +_JSON_VALUE: Final = TypeAdapter(object) + + +def _decoded_json(response: httpx.Response) -> object: + return _JSON_VALUE.validate_python(response.json()) + + +def _runway_task(decoded: object) -> _RunwayMLTask: + task: Final = decoded_json_object(decoded) + return cast(_RunwayMLTask, task) # cast-ok: the payload passes through unvalidated, as before it was typed + + +def _parse_runway_task(response: httpx.Response) -> _RunwayMLTask: + return _runway_task(_decoded_json(response)) class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): @@ -234,7 +252,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayMLTask = response.json() + response_data: _RunwayMLTask = _parse_runway_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -282,7 +300,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayMLTask = response.json() + response_data: _RunwayMLTask = _parse_runway_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -328,13 +346,14 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): } """ try: - response_data: _RunwayMLTask = raw_response.json() + decoded: Final = _decoded_json(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) + response_data: _RunwayMLTask = _runway_task(decoded) verbose_logger.debug("RunwayML starting polling...") @@ -358,7 +377,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) # Update response_data with polled result - response_data = raw_response.json() + response_data = _parse_runway_task(raw_response) verbose_logger.debug("RunwayML polling complete, transforming to OpenAI format") @@ -388,13 +407,14 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): We need to poll the task until it completes (status SUCCEEDED) using async polling. """ try: - response_data: _RunwayMLTask = raw_response.json() + decoded: Final = _decoded_json(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) + response_data: _RunwayMLTask = _runway_task(decoded) verbose_logger.debug("RunwayML starting polling (async)...") @@ -418,7 +438,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) # Update response_data with polled result - response_data = raw_response.json() + response_data = _parse_runway_task(raw_response) verbose_logger.debug("RunwayML polling complete (async), transforming to OpenAI format") diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 1f052648f13..ce1ec1714f1 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -41,6 +41,11 @@ class _RunwayTtsTaskResponse(TypedDict, total=False): failureCode: ReadOnly[str] +def _parse_runway_tts_task(response: httpx.Response) -> _RunwayTtsTaskResponse: + task: Final[_RunwayTtsTaskResponse] = response.json() + return task + + class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Configuration for RunwayML Text-to-Speech @@ -321,7 +326,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayTtsTaskResponse = response.json() + response_data: _RunwayTtsTaskResponse = _parse_runway_tts_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -368,7 +373,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data: _RunwayTtsTaskResponse = response.json() + response_data: _RunwayTtsTaskResponse = _parse_runway_tts_task(response) # Check task status status = self._check_task_status(response_data=response_data) @@ -459,7 +464,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -489,7 +494,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(polled_response) verbose_logger.debug("RunwayML TTS polling complete, downloading audio") @@ -528,7 +533,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(raw_response) except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -558,7 +563,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = _parse_runway_tts_task(polled_response) verbose_logger.debug("RunwayML TTS polling complete (async), downloading audio") diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 106567f6fba..ad8a24c1535 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -119,7 +119,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): @staticmethod def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse: - return raw_response.json() + return _parse_runway_task_response(raw_response) def get_supported_openai_params(self, model: str) -> list: """ diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 14f0e3ead62..eaa805e7769 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -1,3 +1,4 @@ +import builtins import codecs import functools import json @@ -34,6 +35,14 @@ class _SageMakerMessagesChunkData(TypedDict, total=False): choices: ReadOnly[Required[list[Mapping[str, object]]]] +class _SageMakerPayloadPart(TypedDict, total=False): + bytes: ReadOnly[builtins.bytes] + + +class _SageMakerStreamEvent(TypedDict, total=False): + chunk: ReadOnly[_SageMakerPayloadPart] + + class _SageMakerChunkDataRoot(RootModel[_SageMakerChunkData]): model_config = ConfigDict(frozen=True, hide_input_in_errors=True) @@ -345,7 +354,7 @@ class AWSEventStreamDecoder: ), ) response_dict: Final = event.to_response_dict() - parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) + parsed_response: Final[_SageMakerStreamEvent] = self.parser.parse(response_dict, response_stream_shape) body: Final = response_dict.get("body") if response_dict["status_code"] != 200: diff --git a/litellm/llms/sagemaker/embedding/cohere_transformation.py b/litellm/llms/sagemaker/embedding/cohere_transformation.py index 35352f65dda..8c0f4437edf 100644 --- a/litellm/llms/sagemaker/embedding/cohere_transformation.py +++ b/litellm/llms/sagemaker/embedding/cohere_transformation.py @@ -10,7 +10,8 @@ be of type Object`. Reference: https://docs.cohere.com/v2/reference/embed """ -from typing import TYPE_CHECKING, Any, cast +from collections.abc import Sequence +from typing import TYPE_CHECKING, cast if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -25,6 +26,7 @@ from litellm.llms.bedrock.embed.cohere_transformation import ( BedrockCohereEmbeddingConfig, ) from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse from ..common_utils import SagemakerError @@ -123,7 +125,7 @@ class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/sagemaker/embedding/openai_transformation.py b/litellm/llms/sagemaker/embedding/openai_transformation.py index a3f0ebfbf59..1222c697832 100644 --- a/litellm/llms/sagemaker/embedding/openai_transformation.py +++ b/litellm/llms/sagemaker/embedding/openai_transformation.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllEmbeddingInputValues from httpx._models import Headers, Response -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig @@ -25,6 +25,7 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import STR_KEYED_MAPPING, SagemakerError SUPPORTED_OPENAI_PARAMS: Final = ("dimensions", "encoding_format", "user") +_DECODED_JSON: Final = TypeAdapter(object) class OpenAIEmbeddingUsage(BaseModel): @@ -94,7 +95,7 @@ class SagemakerOpenAIEmbeddingConfig(BaseEmbeddingConfig): litellm_params: Mapping[str, object] | None = None, ) -> "EmbeddingResponse": try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except ValueError as e: raise SagemakerError( message=f"Failed to parse response: {e}", diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 1395f9b98ec..2cd02683167 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -4,7 +4,8 @@ Translate from OpenAI's `/v1/embeddings` to Sagemaker's `/invoke` In the Huggingface TGI format. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Sequence +from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -15,6 +16,7 @@ from httpx._models import Headers, Response from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SagemakerError @@ -158,7 +160,7 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 8a7baf18e8f..622c28f3c90 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -25,7 +25,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -from ..credentials import get_token_creator +from ..credentials import AICoreConfiguration, AICoreDeploymentList, get_token_creator from .handler import ( AsyncSAPStreamIterator, GenAIHubOrchestrationError, @@ -186,11 +186,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): # Keep a short, tight client lifecycle here to avoid fd leaks client: Final = litellm.module_level_client # with httpx.Client(timeout=30) as client: - deployments: Final = client.get(join_api_base(self.base_url, "lm/deployments"), headers=self.headers).json() + deployments: Final[AICoreDeploymentList] = client.get( + join_api_base(self.base_url, "lm/deployments"), headers=self.headers + ).json() valid: Final[list[tuple[str, str]]] = [] for dep in deployments.get("resources", []): if dep.get("scenarioId") == "orchestration": - cfg = client.get( + cfg: AICoreConfiguration = client.get( join_api_base(self.base_url, f"lm/configurations/{dep['configurationId']}"), headers=self.headers, ).json() diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index 2622d0fa9b6..74d065956cc 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path from threading import Lock -from typing import Any, Final, Protocol +from typing import Final, Protocol import httpx from pydantic import TypeAdapter @@ -34,6 +34,21 @@ DEFAULT_HOME_PATH: Final = os.path.join(os.path.expanduser("~"), ".aicore") _DECODED_JSON: Final = TypeAdapter(object) +class AICoreDeployment(TypedDict): + configurationId: ReadOnly[str] + scenarioId: ReadOnly[str] + deploymentUrl: ReadOnly[str] + createdAt: ReadOnly[str] + + +class AICoreDeploymentList(TypedDict, total=False): + resources: ReadOnly[Sequence[AICoreDeployment]] + + +class AICoreConfiguration(TypedDict): + executableId: ReadOnly[str] + + def _get_home() -> str: return os.getenv(HOME_PATH_ENV_VAR, DEFAULT_HOME_PATH) @@ -59,7 +74,7 @@ def _get_nested(d: object, path: Sequence[str]) -> object: return cur -def _load_json_env(var_name: str) -> dict[str, object] | None: +def _load_json_env(var_name: str) -> Mapping[str, Sequence[Mapping[str, object]]] | None: raw: Final = os.environ.get(var_name) if not raw: return None @@ -76,7 +91,7 @@ def _str_or_none(value: object) -> str | None: return None -def _load_vcap() -> dict[str, Any]: +def _load_vcap() -> Mapping[str, Sequence[Mapping[str, object]]]: return _load_json_env(VCAP_SERVICES_ENV_VAR) or {} @@ -196,7 +211,7 @@ def resolve_resource_group(sources: list[Source]) -> str | None: def _parse_service_key_once( service_key: str | dict | None, -) -> dict[str, object] | None: +) -> object: """ Pre-parse service_key if it's a string to avoid repeated JSON parsing. @@ -208,7 +223,7 @@ def _parse_service_key_once( return service_key if isinstance(service_key, str): try: - return json.loads(service_key) + return _DECODED_JSON.validate_python(json.loads(service_key)) except json.JSONDecodeError: verbose_logger.warning("SAP service key is a string but not valid JSON. Skipping this source.") return None @@ -218,7 +233,7 @@ def _parse_service_key_once( return None -def _resolve_credential_from_service_key(service_key: str | dict | None, cv: CredentialsValue) -> str | None: +def _resolve_credential_from_service_key(service_key: object, cv: CredentialsValue) -> str | None: if service_key is None: return None val: Final = _str_or_none(_get_nested(service_key, (("credentials",) + cv.vcap_key) if cv.vcap_key else (cv.name,))) @@ -253,14 +268,16 @@ def fetch_credentials( """ config: Final = init_conf(profile) - service_key = _parse_service_key_once(service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR)) + parsed_service_key: Final = _parse_service_key_once( + service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR) + ) vcap_service: Final = _get_vcap_service(VCAP_AICORE_SERVICE_NAME) sources: Final = [ Source("kwargs", lambda cv: _str_or_none(kwargs.get(cv.name))), Source( "service key", - lambda cv: _resolve_credential_from_service_key(service_key, cv), + lambda cv: _resolve_credential_from_service_key(parsed_service_key, cv), ), Source( "environment variables", diff --git a/litellm/llms/sap/embed/transformation.py b/litellm/llms/sap/embed/transformation.py index 107fef0b355..a59b7e1e00d 100644 --- a/litellm/llms/sap/embed/transformation.py +++ b/litellm/llms/sap/embed/transformation.py @@ -19,7 +19,7 @@ from litellm.types.llms.openai import AllEmbeddingInputValues from litellm.types.utils import EmbeddingResponse from ..chat.handler import GenAIHubOrchestrationError -from ..credentials import get_token_creator +from ..credentials import AICoreConfiguration, AICoreDeploymentList, get_token_creator class Usage(LiteLLMBaseModel): @@ -102,10 +102,12 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): def deployment_url(self) -> str: with httpx.Client(timeout=30) as client: valid_deployments: Final = [] - deployments: Final = client.get(join_api_base(self.base_url, "lm/deployments"), headers=self.headers).json() + deployments: Final[AICoreDeploymentList] = client.get( + join_api_base(self.base_url, "lm/deployments"), headers=self.headers + ).json() for deployment in deployments.get("resources", []): if deployment["scenarioId"] == "orchestration": - config_details = client.get( + config_details: AICoreConfiguration = client.get( join_api_base(self.base_url, f"lm/configurations/{deployment['configurationId']}"), headers=self.headers, ).json() diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 5d5fad6128e..974996c5e48 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -6,10 +6,12 @@ Handles transformation between OpenAI-compatible format and Stability AI API for API Reference: https://platform.stability.ai/docs/api-reference """ +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -30,6 +32,13 @@ else: LiteLLMLoggingObj = Any +class _StabilityImageEditResponse(TypedDict, total=False): + image: ReadOnly[str] + finish_reason: ReadOnly[str] + seed: ReadOnly[int] + errors: ReadOnly[Sequence[str]] + + class StabilityImageEditConfig(BaseImageEditConfig): """ Configuration for Stability AI image edit. @@ -259,7 +268,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): OpenAI expects: {"data": [{"b64_json": "base64..."}], "created": timestamp} """ try: - response_data: Final = raw_response.json() + response_data: Final[_StabilityImageEditResponse] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing Stability AI response: {e}", diff --git a/litellm/llms/tencent/messages/transformation.py b/litellm/llms/tencent/messages/transformation.py index 1e52bc37d47..30c3ecaf3a9 100644 --- a/litellm/llms/tencent/messages/transformation.py +++ b/litellm/llms/tencent/messages/transformation.py @@ -5,8 +5,6 @@ Tencent TokenHub exposes an Anthropic-compatible Messages API endpoint alongside its standard OpenAI-compatible chat completions endpoint. """ -from typing import Any - import litellm from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.messages.transformation import ( @@ -48,7 +46,7 @@ class TencentAnthropicMessagesConfig(AnthropicMessagesConfig): self, headers: dict, model: str, - messages: list[Any], + messages: list[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index b37bbd78f2b..fe0f96cb805 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -7,6 +7,7 @@ from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Final, Literal from httpx import Headers, Response +from typing_extensions import ReadOnly, TypedDict from litellm.constants import DEFAULT_MAX_TOKENS_FOR_TRITON from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory @@ -32,6 +33,10 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +class _TritonGenerateResponse(TypedDict): + text_output: ReadOnly[str] + + class TritonConfig(BaseConfig): """ Base class for Triton configurations. @@ -220,7 +225,7 @@ class TritonGenerateConfig(TritonConfig): json_mode: bool | None = None, ) -> ModelResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[_TritonGenerateResponse] = raw_response.json() except Exception: raise TritonError(message=raw_response.text, status_code=raw_response.status_code) model_response.choices = [Choices(index=0, message=Message(content=raw_response_json["text_output"]))] diff --git a/litellm/llms/triton/embedding/transformation.py b/litellm/llms/triton/embedding/transformation.py index 8c6e0b832fc..3a735f741fd 100644 --- a/litellm/llms/triton/embedding/transformation.py +++ b/litellm/llms/triton/embedding/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.base_llm.chat.transformation import AllMessageValues, BaseLLMException from litellm.llms.base_llm.embedding.transformation import ( @@ -14,6 +16,16 @@ from litellm.utils import token_counter from ..common_utils import TritonError +class _TritonEmbeddingOutput(TypedDict): + shape: ReadOnly[Sequence[int]] + data: ReadOnly[Sequence[float]] + + +class _TritonEmbeddingResponse(TypedDict): + outputs: ReadOnly[Sequence[_TritonEmbeddingOutput]] + model_name: NotRequired[ReadOnly[str]] + + class TritonEmbeddingConfig(BaseEmbeddingConfig): """ Transformations for triton /embeddings endpoint (This is a trtllm model) @@ -79,7 +91,7 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict = {}, ) -> EmbeddingResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[_TritonEmbeddingResponse] = raw_response.json() except Exception: raise TritonError(message=raw_response.text, status_code=raw_response.status_code) @@ -134,7 +146,7 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): return TritonError(message=error_message, status_code=status_code, headers=headers) @staticmethod - def split_embedding_by_shape(data: list[float], shape: list[int]) -> list[list[float]]: + def split_embedding_by_shape(data: Sequence[float], shape: Sequence[int]) -> list[Sequence[float]]: if len(shape) != 2: raise ValueError("Shape must be of length 2.") embedding_size: Final = shape[1] diff --git a/litellm/llms/vercel_ai_gateway/chat/transformation.py b/litellm/llms/vercel_ai_gateway/chat/transformation.py index 3fc74b6acc1..cf633358a41 100644 --- a/litellm/llms/vercel_ai_gateway/chat/transformation.py +++ b/litellm/llms/vercel_ai_gateway/chat/transformation.py @@ -6,9 +6,11 @@ Calls done in OpenAI/openai.py as Vercel AI Gateway is openai-compatible. Docs: https://vercel.com/docs/ai-gateway """ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.url_utils import join_api_base @@ -20,6 +22,14 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig from ..common_utils import VercelAIGatewayException +class _VercelAIGatewayModel(TypedDict): + id: ReadOnly[str] + + +class _VercelAIGatewayModelList(TypedDict): + data: ReadOnly[Sequence[_VercelAIGatewayModel]] + + class VercelAIGatewayConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: @@ -99,5 +109,5 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): if response.status_code != 200: raise Exception(f"Failed to get models: {response.text}") - models: Final = response.json()["data"] - return [model["id"] for model in models] + models: Final[_VercelAIGatewayModelList] = response.json() + return [model["id"] for model in models["data"]] diff --git a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py index 39d2e9b4c49..26de755d5ca 100644 --- a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from typing import Final from httpx import Headers, Response +from pydantic import TypeAdapter import litellm from litellm.exceptions import UnsupportedParamsError @@ -45,6 +46,7 @@ from litellm.types.utils import ( DEFAULT_GEMINI_TRANSCRIBE_LOCATION: Final = "global" AUDIO_MODALITY: Final = "AUDIO" +_DECODED_JSON: Final = TypeAdapter(object) class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase): @@ -177,7 +179,7 @@ class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexB raw_response: Response, ) -> TranscriptionResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except ValueError: raise VertexAIError( status_code=raw_response.status_code, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a8d02045740..99ef750b5fc 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -202,7 +202,7 @@ class VertexAIBatchTransformation: @classmethod def transform_vertex_ai_batch_list_response_to_openai_list_response( cls, response: dict[str, Any] - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Transforms Vertex AI batch list response into OpenAI-compatible list response. """ diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index f2da04a7db7..54c0fa97550 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -3,7 +3,7 @@ import json import os import time from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final from urllib.parse import unquote import httpx @@ -197,7 +197,7 @@ class VertexAIFilesHandler(GCSBucketBase): timeout: float | httpx.Timeout, max_retries: int | None, litellm_params: dict | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: """ Download file content from GCS bucket for VertexAI files. Supports both sync and async operations. diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index e36aba00293..410f29ef88a 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,7 +5,7 @@ import json import os import re import time -from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping +from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping, Sequence from contextlib import aclosing from dataclasses import dataclass from functools import partial @@ -16,6 +16,7 @@ from urllib.parse import quote, unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from pydantic import TypeAdapter from typing_extensions import ReadOnly, Required import litellm @@ -31,7 +32,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, @@ -96,6 +97,7 @@ _JSONL_NEWLINE: Final = b"\n" _BATCH_OUTPUT_FIRST_ROW_PEEK_LIMIT_BYTES: Final = 32 * 1024 * 1024 _PASSTHROUGH_MANAGED_GCS_PREFIX: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}passthrough/" _RAW_UPLOAD_CHUNK_BYTES: Final = 1024 * 1024 +_DECODED_JSON: Final = TypeAdapter(object) class _GcsObjectMetadataJson(TypedDict, total=False): @@ -827,7 +829,7 @@ def _iter_openai_jsonl_entries( def _load_openai_jsonl_row(lineno: int, line: str) -> object: try: - return json.loads(line) + return _DECODED_JSON.validate_python(json.loads(line)) except json.JSONDecodeError as e: raise VertexAIError( status_code=400, @@ -967,7 +969,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[Mapping[str, object]], deployment_model: str | None = None, ) -> str: """ @@ -985,8 +987,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): `projects/../locations/../endpoints/` batch job model instead of a nonexistent publisher model. """ - first_entry: Final = openai_jsonl_content[0] - body_model: Final = first_entry["body"].get("model") if isinstance(first_entry.get("body"), dict) else None + body: Final = openai_jsonl_content[0].get("body") + body_fields: Final = as_str_mapping(body) + body_model: Final = body_fields.get("model") if body_fields is not None and isinstance(body, dict) else None jsonl_model: Final = body_model if isinstance(body_model, str) else "" raw_model: Final = deployment_model.removeprefix("vertex_ai/") if deployment_model else jsonl_model return _managed_batch_object_name(raw_model, passthrough=False) diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index 40564d536a7..7cb52281a6a 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -27,6 +27,11 @@ from litellm.types.utils import LiteLLMFineTuningJob _JSON_VALUE: Final = TypeAdapter(object) +def _response_tuning_job(response: httpx.Response) -> ResponseTuningJob: + tuning_job: Final[ResponseTuningJob] = response.json() + return ResponseTuningJob(**tuning_job) + + class VertexFineTuningAPI(VertexLLM): """ Vertex methods to support for batches @@ -187,9 +192,7 @@ class VertexFineTuningAPI(VertexLLM): "got response from creating fine tuning job: %s", _JSON_VALUE.validate_python(response.json()) ) - vertex_response: Final = ResponseTuningJob( - **response.json(), - ) + vertex_response: Final = _response_tuning_job(response) verbose_logger.debug("vertex_response %s", vertex_response) open_ai_response: Final = self.convert_vertex_response_to_open_ai_response(vertex_response) @@ -272,9 +275,7 @@ class VertexFineTuningAPI(VertexLLM): verbose_logger.debug( "got response from creating fine tuning job: %s", _JSON_VALUE.validate_python(response.json()) ) - vertex_response: Final = ResponseTuningJob( - **response.json(), - ) + vertex_response: Final = _response_tuning_job(response) verbose_logger.debug("vertex_response %s", vertex_response) open_ai_response: Final = self.convert_vertex_response_to_open_ai_response(vertex_response) diff --git a/litellm/llms/vertex_ai/gemini/long_decode.py b/litellm/llms/vertex_ai/gemini/long_decode.py index 571fa94d4a5..25aae70a853 100644 --- a/litellm/llms/vertex_ai/gemini/long_decode.py +++ b/litellm/llms/vertex_ai/gemini/long_decode.py @@ -332,7 +332,8 @@ def with_initial_prompt_tokens(usage: Usage, initial_prompt_tokens: int | None) def parse_body(response: httpx.Response) -> GenerateContentResponseBody: try: - return GenerateContentResponseBody(**response.json()) + body: Final[GenerateContentResponseBody] = response.json() + return GenerateContentResponseBody(**body) except Exception as e: raise VertexAIError( message=( diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 0a800497de1..f15c6393cb2 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -6,6 +6,7 @@ import json from typing import TYPE_CHECKING, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.llms.custom_httpx.http_handler import ( @@ -35,6 +36,11 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +class _GeminiFileResource(TypedDict, total=False): + mimeType: ReadOnly[str] + uri: ReadOnly[str] + + class GoogleBatchEmbeddings(VertexLLM): @staticmethod def _flatten_and_detect_file_refs( @@ -74,7 +80,7 @@ class GoogleBatchEmbeddings(VertexLLM): if response.status_code != 200: raise Exception(f"Error fetching file {element}: {response.status_code} {response.text}") - file_data = response.json() + file_data: _GeminiFileResource = response.json() resolved_files[element] = { "mime_type": file_data.get("mimeType", ""), "uri": file_data.get("uri", element), @@ -111,7 +117,7 @@ class GoogleBatchEmbeddings(VertexLLM): if response.status_code != 200: raise Exception(f"Error fetching file {element}: {response.status_code} {response.text}") - file_data = response.json() + file_data: _GeminiFileResource = response.json() resolved_files[element] = { "mime_type": file_data.get("mimeType", ""), "uri": file_data.get("uri", element), diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 725a7f39917..6d071363823 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -190,7 +190,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) - return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files)) + return cast(tuple[dict[str, object], RequestFiles | None], (payload, empty_files)) def transform_image_edit_response( self, diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index 085ff715c30..deaa32f7308 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles from pydantic import ConfigDict, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -34,6 +35,10 @@ _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object] ) +class _VertexImagenPredictResponse(TypedDict, total=False): + predictions: ReadOnly[Sequence[object]] + + class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Imagen Image Edit Configuration @@ -207,7 +212,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_json: Final = raw_response.json() + response_json: Final[_VertexImagenPredictResponse] = raw_response.json() except Exception as exc: raise self.get_error_class( error_message=f"Error transforming image edit response: {exc}", diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index 7ba7b060991..ec915e10fcd 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -168,7 +168,7 @@ class VertexImageGeneration(VertexLLM): if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") - json_response: Final = response.json() + json_response: Final[Mapping[str, object]] = response.json() return self.process_image_generation_response(json_response, model_response, model) async def aimage_generation( @@ -272,7 +272,7 @@ class VertexImageGeneration(VertexLLM): if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") - json_response: Final = response.json() + json_response: Final[Mapping[str, object]] = response.json() return self.process_image_generation_response(json_response, model_response, model) def is_image_generation_response(self, json_response: dict[str, Any]) -> bool: diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 916e25a027e..f6df772dd56 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_DECODED_JSON: Final = TypeAdapter(object) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( @@ -299,7 +300,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): Transform Gemini image generation response to litellm ImageResponse format """ try: - response_data: Final = raw_response.json() + response_data: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index da77fa338db..c282ae22ea6 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -208,7 +208,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): if raw_response.status_code != 200: raise Exception(f"Error: {raw_response.status_code} {raw_response.text}") - _json_response: Final = raw_response.json() + _json_response: Final[MultimodalPredictions] = raw_response.json() if "predictions" not in _json_response: raise InternalServerError( message=f"embedding response does not contain 'predictions', got {_json_response}", diff --git a/litellm/llms/vertex_ai/rag_engine/transformation.py b/litellm/llms/vertex_ai/rag_engine/transformation.py index a34ef519578..61cb083fa4c 100644 --- a/litellm/llms/vertex_ai/rag_engine/transformation.py +++ b/litellm/llms/vertex_ai/rag_engine/transformation.py @@ -4,7 +4,10 @@ Transformation utilities for Vertex AI RAG Engine. Handles transforming LiteLLM's unified formats to Vertex AI RAG Engine API format. """ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE @@ -13,6 +16,15 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.rag import RAGChunkingStrategy +class VertexRAGChunkingConfig(TypedDict): + chunk_size: ReadOnly[int] + chunk_overlap: ReadOnly[int] + + +class VertexRAGTransformationConfig(TypedDict): + chunking_config: ReadOnly[VertexRAGChunkingConfig] + + class VertexAIRAGTransformation(VertexBase): """ Transformation class for Vertex AI RAG Engine API. @@ -55,7 +67,7 @@ class VertexAIRAGTransformation(VertexBase): def transform_chunking_strategy_to_vertex_format( self, chunking_strategy: RAGChunkingStrategy | None, - ) -> dict[str, Any]: + ) -> VertexRAGTransformationConfig: """ Transform LiteLLM's unified chunking_strategy to Vertex AI RAG format. @@ -105,7 +117,7 @@ class VertexAIRAGTransformation(VertexBase): self, gcs_uri: str, chunking_strategy: RAGChunkingStrategy | None = None, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Build the request payload for importing RAG files. diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index dce4d2f2a87..398471a7c9b 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -6,10 +6,11 @@ Why separate file? Make it easy to see how transformation works import math import uuid -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -27,6 +28,17 @@ from litellm.types.rerank import ( ) +class _VertexRankedRecord(TypedDict): + id: ReadOnly[str] + score: NotRequired[ReadOnly[float]] + title: NotRequired[ReadOnly[str]] + content: NotRequired[ReadOnly[str]] + + +class _VertexRankResponse(TypedDict, total=False): + records: ReadOnly[Sequence[_VertexRankedRecord]] + + class VertexAIRerankConfig(BaseRerankConfig, VertexBase): """ Configuration for Vertex AI Discovery Engine Rerank API @@ -172,7 +184,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): Transform Vertex AI Discovery Engine response to Cohere format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[_VertexRankResponse] = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse response: {e}") diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py index 45cc5320401..3873a4be0b5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final import litellm @@ -18,16 +19,17 @@ class VertexAIAi21Config(OpenAIGPTConfig): self, max_tokens: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 5d748ad6e74..cfd5b5d3980 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,4 +1,5 @@ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.pass_through.messages.transformation import ( @@ -29,7 +30,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert self, headers: dict, model: str, - messages: list[Any], + messages: Sequence[object], optional_params: dict, litellm_params: dict, api_key: str | None = None, diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index ddc83bc53a5..96fccf007a5 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -1,4 +1,5 @@ import types +from collections.abc import Mapping from typing import Final, Literal from pydantic import TypeAdapter @@ -70,16 +71,17 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): | None = None, title: str | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = types.MappingProxyType(locals().copy()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + class_attributes: Final[Mapping[str, object]] = types.MappingProxyType(cls.__dict__) return { k: v - for k, v in cls.__dict__.items() + for k, v in class_attributes.items() if not k.startswith("__") and not isinstance( v, diff --git a/litellm/llms/volcengine/chat/transformation.py b/litellm/llms/volcengine/chat/transformation.py index 52b3242c186..4ccab38e761 100644 --- a/litellm/llms/volcengine/chat/transformation.py +++ b/litellm/llms/volcengine/chat/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig @@ -34,7 +35,7 @@ class VolcEngineChatConfig(OpenAILikeChatConfig): top_p: int | None = None, response_format: dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 74f6182bd89..747ea21e4fb 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -45,6 +45,11 @@ class _EventModelClass(Protocol): def model_validate(self, obj: Mapping[str, object]) -> ResponsesAPIStreamingResponse: ... +class _FieldDefaultCarrier(Protocol): + @property + def default(self) -> object: ... + + class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): _SUPPORTED_OPTIONAL_PARAMS: list[str] = [ # Doc-listed knobs @@ -430,6 +435,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): annotation: Final[object] = field.annotation return annotation + @staticmethod + def _field_default(field: _FieldDefaultCarrier) -> object: + return field.default + @staticmethod def _fill_missing_fields(chunk: Mapping[str, object], event_model: object | None) -> Mapping[str, object]: """ @@ -451,7 +460,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): continue # Explicit default or factory - field_default: object = field.default + field_default: object = VolcEngineResponsesAPIConfig._field_default(field) if field_default is not pyd_fields.PydanticUndefined and field_default is not None: patched[name] = field_default continue diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 7fe643870d5..21c6e6a82b9 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -29,6 +29,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) _STR: Final = TypeAdapter(str) +_DECODED_JSON: Final = TypeAdapter(object) class VoyageRerankConfig(BaseRerankConfig): @@ -112,7 +113,7 @@ class VoyageRerankConfig(BaseRerankConfig): logging_obj.post_call(original_response=raw_response.text) try: - _json_response: Final = raw_response.json() + _json_response: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception: raise VoyageError( message=f"Failed to parse response: {raw_response.text}", diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index e5da7dd412f..35fda05c028 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -29,6 +29,7 @@ from ...openai.transcriptions.whisper_transformation import ( from ..common_utils import IBMWatsonXMixin _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTranscriptionConfig): @@ -172,7 +173,7 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran removed before creating the TranscriptionResponse object. """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise ValueError(f"Error transforming response to json: {e}\nResponse: {raw_response.text}") diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 591ddda3cbd..87f30483196 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -30,6 +30,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _OPTIONAL_INT: Final[TypeAdapter[int | None]] = TypeAdapter(int | None) _STR: Final = TypeAdapter(str) +_DECODED_JSON: Final = TypeAdapter(object) class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): @@ -170,7 +171,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): Transform IBM watsonx.ai rerank response to LiteLLM RerankResponse format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _DECODED_JSON.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Failed to parse response: {e}", diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 984b62eeb59..2b943fa4107 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -32,6 +32,7 @@ from ...openai.chat.gpt_transformation import ( ) _RESPONSE_BODY: Final = TypeAdapter(dict[object, object], config=ConfigDict(hide_input_in_errors=True)) +_FILTERED_TOOL: Final = TypeAdapter(object) def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None: @@ -204,11 +205,13 @@ class XAIChatConfig(OpenAIGPTConfig): if param == "max_completion_tokens": optional_params["max_tokens"] = value elif param == "tools" and value is not None: - tools = [] - for tool in value: - tool = filter_value_from_dict(tool, "strict") - if tool is not None: - tools.append(tool) + tools = [ + filtered_tool + for filtered_tool in ( + _FILTERED_TOOL.validate_python(filter_value_from_dict(tool, "strict")) for tool in value + ) + if filtered_tool is not None + ] if len(tools) > 0: optional_params["tools"] = tools elif param in supported_openai_params: diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index de2cbf5f9e3..aec535c2086 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -120,12 +120,11 @@ def _web_search_cost_per_call_from_model_info(model_info: "ModelInfo") -> float: search_costs: Final = model_info.get("search_context_cost_per_query") if not isinstance(search_costs, Mapping): return _DEFAULT_WEB_SEARCH_COST_PER_CALL - for key in ( - "search_context_size_medium", - "search_context_size_low", - "search_context_size_high", + for value in ( + search_costs.get("search_context_size_medium"), + search_costs.get("search_context_size_low"), + search_costs.get("search_context_size_high"), ): - value = search_costs.get(key) if value is None: continue try: diff --git a/litellm/main.py b/litellm/main.py index 327a7a1c4b5..c22cd8eb11c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -802,6 +802,11 @@ def _transcription_response_from_cached_dict(cached_response_dict: Mapping[str, return TranscriptionResponse(**cached_response_dict) +def _text_completion_response_from_cached_dict(cached_response_dict: Mapping[str, object]) -> TextCompletionResponse: + fields: Final = {key: value for key, value in cached_response_dict.items() if key != "stream"} + return TextCompletionResponse(stream=bool(cached_response_dict.get("stream", False)), **fields) + + async def _async_streaming(response, model, custom_llm_provider, args): try: print_verbose(f"received response in _async_streaming: {response}") @@ -3560,7 +3565,7 @@ def _complete_nadir(ctx: CompletionDispatchContext) -> _CompletionDispatchResult encoding=_get_encoding(), api_key=api_key, logging_obj=ctx.logging, - client=ctx.client, + client=_dispatch_client_http(ctx), ) ctx.logging.post_call(input=ctx.messages, api_key=api_key, original_response=response) @@ -4197,6 +4202,7 @@ _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) _BASE_URL_HINT: Final = TypeAdapter(BaseUrlRoutingHint) +_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _AWS_BEDROCK_CLIENT_CREDENTIAL_KEYS: Final = ( @@ -4741,7 +4747,7 @@ def _complete_petals(ctx: CompletionDispatchContext) -> _CompletionDispatchResul ) if stream is True: ## [BETA] # Fake streaming for petals - resp_string: Final = model_response["choices"][0]["message"]["content"] + resp_string: Final = _OPAQUE_VALUE_ADAPTER.validate_python(model_response["choices"][0]["message"]["content"]) return CustomStreamWrapper( resp_string, model, @@ -7719,7 +7725,7 @@ async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextComp init_response: Final = await loop.run_in_executor(None, func_with_context) if isinstance(init_response, dict) or isinstance(init_response, TextCompletionResponse): ## CACHING SCENARIO if isinstance(init_response, dict): - response = TextCompletionResponse(**init_response) + response = _text_completion_response_from_cached_dict(init_response) else: response = init_response elif asyncio.iscoroutine(init_response): @@ -7899,7 +7905,6 @@ def text_completion( tokenizer: Final = get_encoding("p50k_base") ## if it's a 2d list - each element in the list is a text_completion() request if len(prompt) > 0 and isinstance(prompt[0], list): - responses: Final = [None for x in prompt] # init responses def process_prompt(i, individual_prompt): decoded_prompt: Final = tokenizer.decode(individual_prompt) @@ -7916,16 +7921,16 @@ def text_completion( text_completion_response["object"] = "text_completion" text_completion_response["created"] = response.get("created", None) text_completion_response["model"] = response.get("model", None) - return response["choices"][0] + return response.choices[0] with concurrent.futures.ThreadPoolExecutor() as executor: completed_futures: Final = [ executor.submit(process_prompt, i, individual_prompt) for i, individual_prompt in enumerate(prompt) ] - for i, future in enumerate(concurrent.futures.as_completed(completed_futures)): - responses[i] = future.result() - text_completion_response.choices = responses + text_completion_response.choices = [ + future.result() for future in concurrent.futures.as_completed(completed_futures) + ] return text_completion_response # else: diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py index 6a0d22cb63b..f3abbc2e93d 100644 --- a/litellm/models/organization_membership.py +++ b/litellm/models/organization_membership.py @@ -6,6 +6,7 @@ Canonical definition for ``litellm_organizationmembership``. Re-exported from """ from datetime import datetime +from typing import Final from pydantic import ConfigDict, model_validator @@ -31,9 +32,10 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): @model_validator(mode="after") def populate_user_email(self) -> "LiteLLM_OrganizationMembershipTable": - if self.user_email is None and self.user is not None: - if isinstance(self.user, dict): - self.user_email = self.user.get("user_email") + user: Final[object] = self.user + if self.user_email is None and user is not None: + if isinstance(user, dict): + self.user_email = user.get("user_email") else: - self.user_email = getattr(self.user, "user_email", None) + self.user_email = getattr(user, "user_email", None) return self diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7f0d392ab1b..824633bdf5a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -99,6 +99,7 @@ if TYPE_CHECKING: _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({}) OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +_CACHED_VALUE: Final = TypeAdapter(object) def _normalize_caller_admission_credential(value: str) -> str: @@ -4258,7 +4259,7 @@ class MCPRequestHandler: cache_key: Final = f"agent_object_permission_id:{agent_id}" try: - cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final = _CACHED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 71eb2112bd5..2db548f5863 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -1173,12 +1173,9 @@ async def list_gateway_catalog( try: page: Final = await fetch_optional_catalog_page(context, request, server, allowed, cursor) + page_meta: Final[Mapping[str, object]] = page.meta or {} return page.model_copy( - update={ - "meta": { - key: value for key, value in (page.meta or {}).items() if key != SERVER_OUTCOMES_META_KEY - } - } + update={"meta": {key: value for key, value in page_meta.items() if key != SERVER_OUTCOMES_META_KEY}} ) except Exception as error: if cursor is not None: diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 7302b74f15a..0666cc7e6c0 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -165,6 +165,7 @@ _CLIENT_FORWARDED_AUTH_TYPES: Final["frozenset[str]"] = frozenset({"true_passthr _MINTED_TOKEN_CREDENTIAL_FIELDS: Final["frozenset[str]"] = frozenset({"access_token", "refresh_token", "expires_in"}) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_DECODED_JSON: Final = TypeAdapter(object) def _bind_submitted_oauth_client( @@ -339,7 +340,7 @@ def _decrypt_env_vars_on_returned_row(row: object) -> None: return if isinstance(env_vars, str): try: - env_vars = json.loads(env_vars) + env_vars = _DECODED_JSON.validate_python(json.loads(env_vars)) except (json.JSONDecodeError, TypeError): return if not isinstance(env_vars, list): @@ -476,8 +477,9 @@ def _prepare_mcp_server_data( # Serialize JSON fields from ``data_dict`` (not ``data``) so the # exclude_unset filter is respected. Reading back from ``data`` would # reintroduce defaults (e.g. ``env={}``) for fields the caller never set. - if data_dict.get("static_headers") is not None: - data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"]) + static_headers: Final[Mapping[str, str] | None] = data_dict.get("static_headers") + if static_headers is not None: + data_dict["static_headers"] = encrypt_secret_map(static_headers) # env_vars is read from ``data_dict`` (not ``data``) like every other JSON # column so the exclude_unset filter is respected: a partial update that @@ -489,11 +491,13 @@ def _prepare_mcp_server_data( _encrypt_global_env_var_values(serialized_env_vars) data_dict["env_vars"] = safe_dumps(serialized_env_vars) - if data_dict.get("mcp_info") is not None: - data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"]) + mcp_info: Final[object] = data_dict.get("mcp_info") + if mcp_info is not None: + data_dict["mcp_info"] = safe_dumps(mcp_info) - if data_dict.get("env") is not None: - data_dict["env"] = encrypt_secret_map(data_dict["env"]) + env: Final[Mapping[str, str] | None] = data_dict.get("env") + if env is not None: + data_dict["env"] = encrypt_secret_map(env) if "tool_name_to_display_name" in data_dict: data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {}) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3b2c2e09df5..f5c54451006 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -701,7 +701,7 @@ async def _store_per_user_token_server_side( expires_in = None refresh_token: Final[str | None] = token_response.get("refresh_token") or None - raw_scope: Final[object] = token_response.get("scope", requested_scope) + raw_scope: Final[object] = _DECODED_JSON.validate_python(token_response.get("scope", requested_scope)) scopes: Final[list[str] | None] = raw_scope.split() if isinstance(raw_scope, str) else None try: diff --git a/litellm/proxy/_experimental/mcp_server/faults/classify.py b/litellm/proxy/_experimental/mcp_server/faults/classify.py index bb41436b495..7826fc34eac 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/classify.py +++ b/litellm/proxy/_experimental/mcp_server/faults/classify.py @@ -10,6 +10,7 @@ from __future__ import annotations from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.faults.types import ( @@ -25,6 +26,8 @@ from litellm.proxy._experimental.mcp_server.faults.types import ( UpstreamReportedFault, ) +_DECODED_JSON: Final = TypeAdapter(object) + def _safe_text(response: httpx.Response) -> str: try: @@ -35,7 +38,7 @@ def _safe_text(response: httpx.Response) -> str: def _safe_json(response: httpx.Response) -> object: try: - return response.json() + return _DECODED_JSON.validate_python(response.json()) except Exception: return None diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index 84673c8c738..a3e33929f6d 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -226,7 +226,7 @@ def rate_limit_retry_delay( ) -> float: if retry_after_seconds is not None: return min(retry_after_seconds, max_delay) - return min(max_delay, base_delay * 2**attempt * (1 + jitter)) + return min(max_delay, base_delay * 2.0**attempt * (1 + jitter)) T: Final = TypeVar("T") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7eddf592262..db70c7d77de 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -488,6 +488,10 @@ class _AuthorizationServerMetadataPayload(TypedDict, total=False): token_endpoint_auth_methods_supported: Sequence[str] +def _mapping_value(source: Mapping[str, object], key: str, default: object) -> object: + return source.get(key, default) + + def _blank_to_none(value: str | None) -> str | None: """Collapse an absent, empty, or whitespace-only string to ``None``. @@ -2757,7 +2761,9 @@ class MCPServerManager: server_id=server_id, name=name_for_prefix, protocol_version=TypeAdapter(MCPUpstreamProtocol).validate_python( - server_config.get("protocol_version", mcp_info.get("protocol_version", "auto")) + _mapping_value( + server_config, "protocol_version", _mapping_value(mcp_info, "protocol_version", "auto") + ) ), alias=alias, server_name=server_name, diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 6d4fefcdb6c..cc92b8cce9b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -11,6 +11,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -41,6 +42,8 @@ from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +_DECODED_JSON: Final = TypeAdapter(object) + class MCPOAuth2TokenCache(InMemoryCache): """ @@ -160,7 +163,7 @@ class MCPOAuth2TokenCache(InMemoryCache): f"failed with status {exc.response.status_code}" ) from exc - body: Final = response.json() + body: Final = _DECODED_JSON.validate_python(response.json()) if not isinstance(body, dict): raise ValueError( diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index c31560a9f63..b3aa040c507 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -18,6 +18,7 @@ from typing import Final, Literal, Protocol, TypeAlias import jwt from fastapi import HTTPException from jwt.types import Options +from pydantic import TypeAdapter from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -38,6 +39,7 @@ _ALLOWED_ID_TOKEN_ALGORITHMS: Final = ( "PS512", ) _JWKS_CACHE_TTL_SECONDS: Final = 3600 +_DECODED_JSON: Final = TypeAdapter(object) _jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) JwksFetcher: TypeAlias = Callable[ @@ -103,7 +105,7 @@ async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> Sequence[Mappi client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response: Final = await client.get(jwks_url) response.raise_for_status() - document: Final = response.json() + document: Final = _DECODED_JSON.validate_python(response.json()) keys: Final = document.get("keys") if isinstance(document, dict) else None if not isinstance(keys, list): raise TypeError(f"JWKS document at {jwks_url} has no 'keys' array") @@ -116,7 +118,7 @@ async def _discover_jwks_url(issuer: str) -> str: client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response: Final = await client.get(discovery_url) response.raise_for_status() - metadata: Final = response.json() + metadata: Final = _DECODED_JSON.validate_python(response.json()) jwks_uri: Final = metadata.get("jwks_uri") if isinstance(metadata, dict) else None if not isinstance(jwks_uri, str) or not jwks_uri: raise ValueError(f"OIDC discovery at {discovery_url} returned no jwks_uri") diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py index 2501f751d8e..d84a65fee02 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py @@ -12,6 +12,7 @@ from __future__ import annotations from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -34,6 +35,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger _GATEWAY_FAULT_OAUTH_ERRORS: Final = frozenset( {"invalid_client", "unauthorized_client", "unsupported_grant_type", "invalid_target", "invalid_scope"} ) +_DECODED_JSON: Final = TypeAdapter(object) def _oauth_error_fields(response: httpx.Response) -> tuple[str | None, str | None]: @@ -46,7 +48,7 @@ def _oauth_error_fields(response: httpx.Response) -> tuple[str | None, str | Non it can carry IdP internals and must never reach the caller. """ try: - body: Final[object] = response.json() + body: Final = _DECODED_JSON.validate_python(response.json()) except Exception: # noqa: BLE001 return None, None if not isinstance(body, dict): diff --git a/litellm/proxy/_experimental/mcp_server/result_conversion.py b/litellm/proxy/_experimental/mcp_server/result_conversion.py index 61435f94c94..fdb9a01fea6 100644 --- a/litellm/proxy/_experimental/mcp_server/result_conversion.py +++ b/litellm/proxy/_experimental/mcp_server/result_conversion.py @@ -13,7 +13,7 @@ from __future__ import annotations import json import math from collections.abc import Sequence -from typing import Final, TypeAlias, TypeVar +from typing import Final, Protocol, TypeAlias, TypeVar from mcp.types import CacheableResult, CallToolResult, ContentBlock, InputRequiredResult, TextContent, Tool from typing_extensions import ReadOnly, TypedDict, assert_never @@ -105,8 +105,17 @@ def complete_call_tool_result(outcome: ToolOutcome, compat: WireCompat) -> CallT return converted +class _StructuredContentCarrier(Protocol): + @property + def structured_content(self) -> object: ... + + +def _structured_content(result: _StructuredContentCarrier) -> object: + return result.structured_content + + def _downgrade_structured_content(result: CallToolResult) -> CallToolResult: - structured: Final = result.structured_content + structured: Final = _structured_content(result) if structured is None or isinstance(structured, dict): return result fallback: Final = TextContent(type="text", text=json.dumps(structured)) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 2f88d8c361f..a7748b9ca88 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -6,7 +6,7 @@ Filters MCP tools semantically for /chat/completions and /responses endpoints. import asyncio from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, TypeVar from litellm._logging import verbose_logger from litellm.exceptions import ContextWindowExceededError @@ -20,6 +20,8 @@ if TYPE_CHECKING: from litellm.router import Router +_Tool: Final = TypeVar("_Tool") + class SemanticToolFilterContextWindowError(Exception): """Raised when the embedding model exceeds its context window, so semantic filtering cannot run.""" @@ -265,9 +267,9 @@ class SemanticMCPToolFilter: async def filter_tools( self, query: str, - available_tools: list[Any], + available_tools: list[_Tool], top_k: int | None = None, - ) -> list[Any]: + ) -> list[_Tool]: """ Filter tools semantically based on query. @@ -390,7 +392,7 @@ class SemanticMCPToolFilter: separator: Final = client_name[-len(canonical) - 1] return separator in ("_", "-") - def _get_tools_by_names(self, tool_names: Sequence[str], available_tools: Sequence[object]) -> list[object]: + def _get_tools_by_names(self, tool_names: Sequence[str], available_tools: Sequence[_Tool]) -> list[_Tool]: """ Get tools from available_tools by their names, preserving the semantic router's ordering. @@ -406,13 +408,13 @@ class SemanticMCPToolFilter: # Exact matches win over suffix matches when both are present, and # each incoming tool is returned at most once even if two canonical # names happen to be tail-compatible with the same incoming name. - available_by_name: Final[dict[str, object]] = {} + available_by_name: Final[dict[str, _Tool]] = {} for tool in available_tools: client_name, _ = self.extract_tool_info(tool) if client_name and client_name not in available_by_name: available_by_name[client_name] = tool - matched: Final[list[object]] = [] + matched: Final[list[_Tool]] = [] used_ids: Final[set[int]] = set() for canonical in tool_names: tool = available_by_name.get(canonical) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index bb7c81f9d60..7207c4150ad 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -14,7 +14,7 @@ import time import types from collections import Counter from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set -from typing import TYPE_CHECKING, Final, NoReturn, Protocol +from typing import TYPE_CHECKING, Final, Literal, NoReturn, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -150,6 +150,16 @@ def _scope_headers(scope: Scope) -> Sequence[tuple[bytes, bytes]]: return asgi_scope_fields(scope).get("headers", []) +class _RawHeaderPairsSource(Protocol): + def get( + self, key: Literal["headers"], default: Sequence[tuple[bytes | str, bytes | str]], / + ) -> Sequence[tuple[bytes | str, bytes | str]]: ... + + +def _raw_header_pairs(source: _RawHeaderPairsSource) -> Sequence[tuple[bytes | str, bytes | str]]: + return source.get("headers", []) + + def unsupported_protocol_version(scope: Scope) -> str | None: """Admit configured HTTP revisions while keeping SSE on the legacy protocol.""" from litellm.proxy._experimental.mcp_server.capabilities import configured_versions @@ -1319,7 +1329,7 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + scope_headers: Final = _raw_header_pairs(scope) for header_name, header_value in scope_headers: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": @@ -1669,7 +1679,7 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + _headers: Final = _raw_header_pairs(scope) def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): @@ -2869,7 +2879,7 @@ if MCP_AVAILABLE: mcp_servers, _sse_client_ip, scoped_server_endpoint=scoped_server_endpoint, - is_initialize=scope.get("method") == "GET", + is_initialize=asgi_scope_fields(scope).get("method", "") == "GET", ): async with ( sse.connect_sse(transport_scope, receive, send) as (read_stream, write_stream), @@ -3090,7 +3100,7 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) + response_headers: Final = _raw_header_pairs(message) for key, value in response_headers: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index f3a78d206d7..35906c793b4 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -131,6 +131,10 @@ async def get_mcp_toolset_by_name( return _toolset_from_row(row) +def _set_fields(data: UpdateMCPToolsetRequest) -> Mapping[str, object]: + return data.model_dump(exclude_unset=True) + + async def update_mcp_toolset( prisma_client: PrismaClient, data: UpdateMCPToolsetRequest, @@ -143,7 +147,7 @@ async def update_mcp_toolset( data_dict: Final = dict( ( (field, json.dumps(value) if field == "tools" else value) - for field, value in data.model_dump(exclude_unset=True).items() + for field, value in _set_fields(data).items() if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) ), updated_by=touched_by, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 1684b9c5c30..ee8f3ea3b30 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -13,6 +13,8 @@ from collections.abc import Set as AbstractSet from typing import Any, Final, Literal, Protocol, overload from urllib.parse import quote +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -21,6 +23,7 @@ if typing.TYPE_CHECKING: MCP_SERVERS_TARGET: Final = "mcp_servers" MCP_OAUTH_TOKENS_TARGET: Final = "mcp_oauth_tokens" +_DECODED_JSON: Final = TypeAdapter(object) class ASGIScopeFields(Protocol): @@ -243,7 +246,7 @@ def _parse_mcp_info_dict(mcp_info: object) -> Mapping[str, object] | None: return mcp_info if isinstance(mcp_info, str): try: - parsed: Final[object] = json.loads(mcp_info) + parsed: Final = _DECODED_JSON.validate_python(json.loads(mcp_info)) except (ValueError, TypeError): return None return parsed if isinstance(parsed, dict) else None diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 28c701e0fc2..aa67d36af4f 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -356,7 +356,7 @@ class LazyFeatureMiddleware: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # Short-circuit once every feature has loaded. if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features): - path: str = scope.get("path", "") + path: str = str(_scope_value(scope, "path")) # Strip the request's root_path so prefix matching works under a # server root path. Without this, requests like /api/v1/policies/... # never match the registered prefixes (/policies/...) and lazy @@ -584,6 +584,7 @@ def _make_warmup_router(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY from fastapi import APIRouter, Depends, HTTPException from fastapi.openapi.utils import get_openapi + from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router: Final = APIRouter() @@ -603,17 +604,22 @@ def _make_warmup_router(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY await _force_load(app, feat, features) feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))] - full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes) + full: Final[Mapping[str, object]] = MappingProxyType( + get_openapi(title=app.title, version=app.version, routes=feat_routes) + ) + paths: Final = as_str_mapping(full.get("paths", {})) or {} + components: Final = as_str_mapping(full.get("components", {})) or {} # Force all operations under one tag so they group under a single Swagger # section — many lazy modules tag routes inconsistently. - for path_ops in full.get("paths", {}).values(): - for op in path_ops.values(): - if isinstance(op, dict): - op["tags"] = [feat.name] + for path_ops in paths.values(): + for op in (as_str_mapping(path_ops) or {}).values(): + operation = as_str_mapping(op) + if isinstance(operation, dict): + operation["tags"] = [feat.name] return { "stub_path": feat.path_prefixes[0], "paths": full.get("paths", {}), - "components": {"schemas": full.get("components", {}).get("schemas", {})}, + "components": {"schemas": components.get("schemas", {})}, } return router diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 00a05307a89..c5a63b0f990 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -16,6 +16,7 @@ import sys from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import TYPE_CHECKING, Final from typing_extensions import ReadOnly, TypedDict @@ -61,7 +62,7 @@ def load_snapshot() -> dict[str, dict] | None: return None -def _normalize_operation_ids(paths: dict[str, dict]) -> None: +def _normalize_operation_ids(paths: Mapping[str, object]) -> None: """Make FastAPI-generated operation IDs stable for multi-method routes. FastAPI derives the default operation ID suffix from the first item in the @@ -69,26 +70,30 @@ def _normalize_operation_ids(paths: dict[str, dict]) -> None: set iteration order can vary between processes, which makes the snapshot drift even when no routes changed. """ + from litellm.litellm_core_utils.core_helpers import as_str_mapping + for path_ops in paths.values(): - if not isinstance(path_ops, dict): + operations = as_str_mapping(path_ops) + if not isinstance(operations, dict): continue - methods = {method for method in path_ops if method in HTTP_METHOD_SUFFIXES} + methods = {method for method in operations if method in HTTP_METHOD_SUFFIXES} if not methods: continue - for method, operation in path_ops.items(): - if method not in HTTP_METHOD_SUFFIXES or not isinstance(operation, dict): + for method, operation in operations.items(): + entry = as_str_mapping(operation) + if method not in HTTP_METHOD_SUFFIXES or not isinstance(entry, dict): continue - operation_id = operation.get("operationId") + operation_id = entry.get("operationId") if not isinstance(operation_id, str): continue for suffix in methods: suffix_token = f"_{suffix}" if operation_id.endswith(suffix_token): - operation["operationId"] = operation_id[: -len(suffix_token)] + f"_{method}" + entry["operationId"] = operation_id[: -len(suffix_token)] + f"_{method}" break @@ -117,6 +122,7 @@ def _register_feature(app: "FastAPI", feat: "LazyFeature") -> str | None: def _feature_fragment(app: "FastAPI", feat: "LazyFeature", used_operation_ids: set[str]) -> SnapshotFragment | None: from fastapi.openapi.utils import get_openapi + from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.proxy_server import ensure_unique_openapi_operation_ids feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))] @@ -124,22 +130,29 @@ def _feature_fragment(app: "FastAPI", feat: "LazyFeature", used_operation_ids: s return None _stabilize_multi_method_route_ids(feat_routes) full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes) - paths: Final = full.get("paths", {}) + openapi: Final[Mapping[str, object]] = MappingProxyType(full) + paths: Final = as_str_mapping(openapi.get("paths", {})) or {} _normalize_operation_ids(paths) for path_ops in paths.values(): - for method, op in path_ops.items(): - if isinstance(op, dict): - operation_id = op.get("operationId") + for method, op in (as_str_mapping(path_ops) or {}).items(): + operation = as_str_mapping(op) + if isinstance(operation, dict): + operation_id = operation.get("operationId") if isinstance(operation_id, str): for suffix in HTTP_METHOD_SUFFIXES: if operation_id.endswith(f"_{suffix}"): - op["operationId"] = operation_id[: -len(suffix)] + method + operation["operationId"] = operation_id[: -len(suffix)] + method break - op["tags"] = [feat.name] - unique: Final = ensure_unique_openapi_operation_ids(full, used_operation_ids) + operation["tags"] = [feat.name] + unique: Final[Mapping[str, object]] = MappingProxyType( + ensure_unique_openapi_operation_ids(full, used_operation_ids) + ) + components: Final = as_str_mapping(unique.get("components", {})) or {} return { - "paths": paths, - "components": {"schemas": unique.get("components", {}).get("schemas", {})}, + "paths": { + path: operations for path, entry in paths.items() if (operations := as_str_mapping(entry)) is not None + }, + "components": {"schemas": as_str_mapping(components.get("schemas", {})) or {}}, } diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index ff58c9c85ec..03d1be7d185 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -16,10 +16,13 @@ fetcher dispatches by ``discovery_mode``: from collections.abc import Mapping from enum import Enum -from typing import Any, Final +from typing import Final from urllib.parse import urlencode +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -46,6 +49,8 @@ AGENT_CARD_WELL_KNOWN_PATHS: Final[tuple[str, ...]] = ( DEFAULT_DISCOVERY_TIMEOUT_SECONDS: Final = 10.0 +_DECODED_JSON: Final = TypeAdapter(object) + class AgentCardDiscoveryError(Exception): """Raised when none of the well-known paths returned a usable agent card.""" @@ -87,7 +92,7 @@ async def fetch_well_known_card( params: Mapping[str, object] | None = None, timeout: float = DEFAULT_DISCOVERY_TIMEOUT_SECONDS, headers: dict[str, str] | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Fetch an agent card from ``base_url`` using the strategy chosen by ``discovery_mode``. Returns the parsed JSON from the first path that @@ -137,13 +142,14 @@ async def fetch_well_known_card( continue try: - card = response.json() + decoded = _DECODED_JSON.validate_python(response.json()) except Exception as exc: last_error = f"{url}: invalid JSON ({exc})" continue + card = as_str_mapping(decoded) if not isinstance(card, dict): - last_error = f"{url}: expected JSON object, got {type(card).__name__}" + last_error = f"{url}: expected JSON object, got {type(decoded).__name__}" continue verbose_proxy_logger.debug("A2A discovery succeeded at %s", url) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c11d296837e..4691c7b0bc5 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -23,6 +23,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth @@ -281,7 +282,7 @@ async def _a2a_sse_event_source( request_id: str | int | None = None, extra_headers: Mapping[str, str] | None = None, served_version: A2AVersion = "0.3", -) -> AsyncGenerator[Mapping[str, object], None]: +) -> AsyncGenerator[object, None]: """Stream an upstream A2A SSE response as parsed JSON-RPC event dicts. Upstream HTTP/JSON-RPC errors are surfaced as a single JSON-RPC error event @@ -327,12 +328,14 @@ async def _a2a_sse_event_source( if not payload: continue try: - event = json.loads(payload) + event = _DECODED_JSON.validate_python(json.loads(payload)) except Exception: continue - if isinstance(event, dict): - event = normalize_stream_event(event, served_version, request_id=request_id) - yield event + event_fields = as_str_mapping(event) + if isinstance(event_fields, dict): + yield normalize_stream_event(event_fields, served_version, request_id=request_id) + else: + yield event finally: await resp.aclose() @@ -741,7 +744,7 @@ async def invoke_agent_a2a( # Get backend URL and agent name agent_card_params: Final = agent.agent_card_params or {} agent_url: Final = agent_card_params.get("url") - agent_name: Final = agent_card_params.get("name", agent_id) + agent_name: Final = (as_str_mapping(agent_card_params) or {}).get("name", agent_id) # Get litellm_params (may include custom_llm_provider for completion bridge) litellm_params: dict[str, object] = agent.litellm_params or {} diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 9a6f674db2c..fb727367caf 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -18,6 +18,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.handler import ( bridge_model_name, ) from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.databricks_oauth import ( resolve_databricks_backend_auth, @@ -173,13 +174,14 @@ async def route_a2a_agent_request( return _call_route(route_type, _bridge_agent_request(data, agent, bridge_auth.litellm_params)) # Get API base URL from agent config - if not agent.agent_card_params or "url" not in agent.agent_card_params: + card_params: Final = as_str_mapping(agent.agent_card_params) + if not card_params or "url" not in card_params: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) # Inject API base and route to litellm - data["api_base"] = agent.agent_card_params["url"] + data["api_base"] = card_params["url"] verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"]) return _call_route(route_type, data) diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 7425c240d10..bebba009d82 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -48,6 +48,7 @@ _DEFAULT_SCOPE: Final = "all-apis" _TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60 _DEFAULT_TTL_SECONDS: Final = 3600 _OAUTH_BLOCK: Final = TypeAdapter(Mapping[str, object]) +_TOKEN_RESPONSE_BODY: Final = TypeAdapter(object) def _resolve_secret(value: object) -> str | None: @@ -275,7 +276,7 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): except httpx.HTTPError as exc: raise ValueError(f"Databricks App OAuth token request failed: {exc}") from exc - body: Final[object] = response.json() + body: Final = _TOKEN_RESPONSE_BODY.validate_python(response.json()) if not isinstance(body, dict): raise ValueError( f"Databricks App OAuth token response returned non-object JSON (got {type(body).__name__})" diff --git a/litellm/proxy/agent_endpoints/managed_identity.py b/litellm/proxy/agent_endpoints/managed_identity.py index abab21901ee..fab7ebbc61d 100644 --- a/litellm/proxy/agent_endpoints/managed_identity.py +++ b/litellm/proxy/agent_endpoints/managed_identity.py @@ -119,6 +119,10 @@ def managed_write_fields( return AgentIdentityFailure(message=f"Invalid agent identity configuration: {exc}") +def _identity_fields(identity: EntraIdentityConfig) -> Mapping[str, object]: + return identity.model_dump() + + def _identity_write(identity: EntraIdentityConfig | None, existing: AgentResponse | None) -> ManagedWriteFields: if identity is None: unbind: Final[ManagedWriteFields] = { @@ -134,7 +138,7 @@ def _identity_write(identity: EntraIdentityConfig | None, existing: AgentRespons existing and existing.identity and existing.identity.active - and all(getattr(existing.identity, name) == value for name, value in identity.model_dump().items()) + and all(getattr(existing.identity, name) == value for name, value in _identity_fields(identity).items()) ): unchanged: Final[ManagedWriteFields] = {} return unchanged diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 8e348423837..b0088d93112 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -21,6 +21,7 @@ import json import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, get_args from urllib.parse import urlsplit @@ -299,7 +300,7 @@ def _validate_installation_preference(spec: PluginSpec) -> None: def _build_plugin_manifest(name: str, spec: PluginSpec) -> Mapping[str, object]: """Build the stored manifest dict shared by plugin create and update.""" - dumped: Final[Mapping[str, object]] = spec.model_dump(exclude_none=True) + dumped: Final[Mapping[str, object]] = MappingProxyType(spec.model_dump(exclude_none=True)) return {"name": name, **{key: value for key, value in dumped.items() if value and key != "name"}} diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8887d86570a..2c6768b8423 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -393,7 +393,7 @@ def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[st class _RawCacheRead(Protocol): - async def async_get_cache(self, *, key: str) -> object: ... + async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: @@ -2712,7 +2712,7 @@ async def get_team_membership( use_writer=True, ) - l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) + l1_cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) if not isinstance(l1_membership, _TeamMembershipCacheMiss): return l1_membership diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7a186d84aa5..25f900a5f9f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -924,8 +924,9 @@ def get_request_route(request: Request) -> str: scope: Final = request.scope if not isinstance(scope, dict): return str(request.url.path) - raw_path: Final[str] = str(scope.get("path", request.url.path)) - root_path: Final[str] = str(scope.get("app_root_path", scope.get("root_path", ""))).rstrip("/") + scope_fields: Final[Mapping[str, object]] = MappingProxyType(scope) + raw_path: Final[str] = str(scope_fields.get("path", request.url.path)) + root_path: Final[str] = str(scope_fields.get("app_root_path", scope_fields.get("root_path", ""))).rstrip("/") if not isinstance(raw_path, str): return str(request.url.path) # Strip root_path only when it matches whole path segments — guarding @@ -1520,8 +1521,15 @@ def get_model_rate_limit_from_metadata( metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit", "model_itpm_limit", "model_otpm_limit"], ) -> dict[str, int] | None: - if getattr(user_api_key_dict, metadata_accessor_key): - return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key) + metadata: Final = ( + user_api_key_dict.team_metadata + if metadata_accessor_key == "team_metadata" + else user_api_key_dict.organization_metadata + if metadata_accessor_key == "organization_metadata" + else user_api_key_dict.project_metadata + ) + if metadata: + return metadata.get(rate_limit_key) return None diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index f1ca2c70f3b..0b349e6a733 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -283,7 +283,7 @@ class IPAddressUtils: @staticmethod def get_mcp_client_ip( request: Request, - general_settings: dict[str, Any] | None = None, + general_settings: dict[str, object] | None = None, ) -> str | None: """ Extract client IP from a FastAPI request for MCP access control. @@ -371,4 +371,4 @@ class IPAddressUtils: return client_ip case _HopCountUnset(): pass - return get_request_ip_address(request, use_x_forwarded_for=use_xff) + return get_request_ip_address(request, use_x_forwarded_for=use_xff is True) diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index a9f349d1750..e291c78294c 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -3,7 +3,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final +from typing import Final, Protocol import httpx @@ -223,11 +223,15 @@ def rds_region_from_hostname(db_host: str) -> str | None: return match.group("region").lower() +class _RdsAuthTokenClient(Protocol): + def generate_db_auth_token(self, *, DBHostname: str, Port: str, DBUsername: str, Region: str | None) -> str: ... + + def generate_iam_auth_token( db_host: str, db_port: str, db_user: str, - client: Any | None = None, + client: _RdsAuthTokenClient | None = None, *, region: str | None = None, ) -> str: diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 720bfabb95a..3ae671d40ac 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -160,7 +160,7 @@ class RouteChecks: if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( route=route, - allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, + allowed_routes=LiteLLMRoutes[allowed_route].value, ): if ( allowed_route in _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 27fa4731baf..b1e8bdbe787 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,6 +13,7 @@ import re import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast import fastapi @@ -177,6 +178,7 @@ from litellm.proxy.utils import ( normalize_route_for_root_path, ) from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository +from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.router_utils.fallback_event_handlers import ( @@ -1319,7 +1321,7 @@ async def fetch_global_spend_with_event_coordination( async def _load_global_spend() -> float | None: proxy_budget_row: Final = await bounded_db_lookup( - prisma_client.db.litellm_usertable.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), + UserRepository(prisma_client).table.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), name="proxy_budget", ) return float(proxy_budget_row.spend) if proxy_budget_row is not None else None @@ -1635,7 +1637,7 @@ async def _auto_register_jwt_mapping( if minted: try: async with db_span("delete_orphaned_jwt_key", "LiteLLM_VerificationToken"): - await prisma_client.db.litellm_verificationtoken.delete(where={"token": token_hash}) + await VerificationTokenRepository(prisma_client).table.delete(where={"token": token_hash}) except Exception as delete_err: # Don't fail the request if cleanup fails — the orphan is # unmapped and inert. Log so an operator can prune it later. @@ -2693,7 +2695,7 @@ async def user_api_key_auth_builder( and valid_token.last_refreshed_at is not None and team_obj.last_refreshed_at > valid_token.last_refreshed_at ): - team_obj_dict: Final = team_obj.__dict__ + team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) for k, v in team_obj_dict.items(): field_name = f"team_{k}" diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 597080be3b3..775d00ed53e 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request @@ -19,6 +20,10 @@ router: Final = APIRouter( ) +def _copy_cache_params(cache_params: Mapping[str, object]) -> Mapping[str, object]: + return dict(**cache_params) + + def _extract_cache_params() -> dict[str, object]: """ Safely extracts and cleans cache parameters. @@ -41,7 +46,7 @@ def _extract_cache_params() -> dict[str, object]: try: cache_params: Final = vars(litellm.cache.cache) cleaned_params: Final = ( - HealthCheckCacheParams.model_validate(dict(**cache_params)).model_dump() if cache_params else {} + HealthCheckCacheParams.model_validate(_copy_cache_params(cache_params)).model_dump() if cache_params else {} ) return masker.mask_dict(cleaned_params) except (AttributeError, TypeError) as e: diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index c027bf2fe98..ccdef44f9e8 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -59,7 +59,7 @@ class ChatClient: frequency_penalty: float | None = None, user: str | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Create a chat completion. @@ -184,7 +184,8 @@ class ChatClient: response.raise_for_status() # Parse SSE stream - for line in response.iter_lines(): + lines: Final[Iterator[object]] = response.iter_lines() + for line in lines: if line: line = _SSE_LINE.validate_python(line).decode("utf-8") if line.startswith("data: "): diff --git a/litellm/proxy/client/cli/commands/encryption.py b/litellm/proxy/client/cli/commands/encryption.py index ef73b4f0d2d..61c8db617af 100644 --- a/litellm/proxy/client/cli/commands/encryption.py +++ b/litellm/proxy/client/cli/commands/encryption.py @@ -4,9 +4,14 @@ from typing import Final import click import rich +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import as_str_mapping from ...http_client import HTTPClient -from .auth import cli_context_obj +from ._cli_context import cli_context_values + +_DECODED_JSON: Final = TypeAdapter(object) @click.group() @@ -40,23 +45,32 @@ def migrate(ctx: click.Context, check_only: bool, dry_run: bool): lite encryption migrate --check # attestation scan, no writes lite encryption migrate # perform the migration """ - client: Final = HTTPClient(cli_context_obj(ctx)["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = HTTPClient(context["base_url"], context["api_key"]) - if check_only: - response = client.request("GET", "/credentials/migrate-encryption/check") - else: - response = client.request( + response: Final = _DECODED_JSON.validate_python( + client.request("GET", "/credentials/migrate-encryption/check") + if check_only + else client.request( "POST", "/credentials/migrate-encryption", json={}, params={"dry_run": "true"} if dry_run else None, ) + ) rich.print_json(data=response) - report: Final = response.get("report", {}) if isinstance(response, dict) else {} + report_value: Final = (as_str_mapping(response) or {}).get("report", {}) + report: Final = as_str_mapping(report_value) + if report is None: + raise click.ClickException(f"Unexpected report in the proxy response: {report_value!r}") residual: Final = report.get("residual_legacy") - if residual is not None and residual > 0: + if residual is None: + return + if not isinstance(residual, int | float): + raise click.ClickException(f"Unexpected residual_legacy in the proxy response: {residual!r}") + if residual > 0: rich.print(f"[yellow]Residual legacy values remaining: {residual}[/yellow]") elif residual == 0: rich.print("[green]No legacy values remaining (residual_legacy == 0).[/green]") diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 5aeefa0c739..fef3e313ee5 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -89,6 +89,28 @@ class Fetch(Protocol): def __call__(self, credentials: Credentials, session_id: str) -> Fetched: ... +class _ResponseBody(Protocol): + def read(self) -> bytes: ... + + +class _OpenedResponse(Protocol): + def __enter__(self) -> _ResponseBody: ... + + def __exit__(self, *exc_info: object) -> None: ... + + +class _UrlOpener(Protocol): + def __call__(self, url: urllib.request.Request, /, *, timeout: float) -> _OpenedResponse: ... + + +def _url_opener() -> _UrlOpener: + return urllib.request.urlopen + + +def _json_decoder() -> Callable[[bytes | str], object]: + return json.loads + + def as_mapping(value: object) -> Mapping[str, object]: return value if isinstance(value, dict) else EMPTY @@ -106,7 +128,7 @@ def printable(value: object) -> str: def load_json(raw: bytes | str) -> object: try: - return json.loads(raw) + return _json_decoder()(raw) except ValueError: return None @@ -195,7 +217,7 @@ def fetch_session(credentials: Credentials, session_id: str) -> Fetched: }, ) try: - with urllib.request.urlopen(request, timeout=FETCH_TIMEOUT_SECONDS) as response: + with _url_opener()(request, timeout=FETCH_TIMEOUT_SECONDS) as response: raw: Final[bytes] = response.read() except urllib.error.HTTPError as error: return Fetched(session=None, definitive=400 <= error.code < 500) diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index f7c545ef48e..7c01732119e 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -189,7 +189,7 @@ def execute_command(path: ShellScope, ctx: click.Context, cli: click.Group) -> N try: if path[0].startswith("-"): cli.main( - ["--base-url", ctx.obj["base_url"], *path], + ["--base-url", cli_context_obj(ctx)["base_url"], *path], prog_name="lite", obj={"secret_vault": context_secret_vault(ctx)}, standalone_mode=False, diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 3f1e7f4259c..a1964d8a1bb 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -4,6 +4,7 @@ from collections.abc import Callable, Iterator, Mapping, MutableMapping from typing import Final, Literal, Protocol import click +from pydantic import TypeAdapter from litellm._version import version as litellm_version from litellm.litellm_core_utils.cli_token_utils import load_cli_token @@ -89,10 +90,13 @@ class LazyCommands(MutableMapping[str, click.Command]): return len(self._loaders) +_MODULE_ATTRIBUTE: Final = TypeAdapter(object) + + def _command_in(module: str, attribute: str) -> Callable[[], click.Command]: @functools.cache def load() -> click.Command: - match getattr(importlib.import_module(module, __package__), attribute): + match _MODULE_ATTRIBUTE.validate_python(getattr(importlib.import_module(module, __package__), attribute)): case click.Command() as command: return command case _: diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index ebcd79457a6..649fe529936 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -75,7 +75,7 @@ class CredentialsManagementClient: credential_values: Mapping[str, object], return_request: bool = False, display_name: str | None = None, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Create a new credential. @@ -121,7 +121,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Delete a credential by name. @@ -182,7 +182,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Get a credential by name. diff --git a/litellm/proxy/client/http_client.py b/litellm/proxy/client/http_client.py index aa0b986b1ad..96a170c75cb 100644 --- a/litellm/proxy/client/http_client.py +++ b/litellm/proxy/client/http_client.py @@ -4,6 +4,11 @@ from collections.abc import Mapping from typing import Any, Final import requests +from typing_extensions import TypedDict, Unpack + + +class _RequestOptions(TypedDict, total=False, extra_items=object): + pass class HTTPClient: @@ -29,7 +34,7 @@ class HTTPClient: data: Mapping[str, object] | list | bytes | None = None, json: Mapping[str, object] | list | None = None, headers: dict[str, str] | None = None, - **kwargs: Any, + **kwargs: Unpack[_RequestOptions], ) -> Any: """Make an HTTP request to the LiteLLM proxy server. diff --git a/litellm/proxy/client/model_groups.py b/litellm/proxy/client/model_groups.py index fef307600c4..1555d170212 100644 --- a/litellm/proxy/client/model_groups.py +++ b/litellm/proxy/client/model_groups.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final import requests @@ -31,7 +31,7 @@ class ModelGroupsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def info(self, return_request: bool = False) -> list[dict[str, Any]] | requests.Request: + def info(self, return_request: bool = False) -> list[dict[str, object]] | requests.Request: """ Get detailed information about all model groups from the server. diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 1b50f8c24f3..7342974c0cc 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -158,7 +158,7 @@ class ModelsManagementClient: model_id: str | None = None, model_name: str | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Get information about a specific model by its ID or name. diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 40cdb117476..b2640d2a765 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,11 +1,15 @@ import asyncio import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter + from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports ConfigSyncPubSub, _ConfigSyncPubSub, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -30,6 +34,7 @@ _pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong _in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES) _BACKOFF_INITIAL_SECONDS: Final = 5.0 _BACKOFF_MAX_SECONDS: Final = 60.0 +_DECODED_JSON: Final = TypeAdapter(object) def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @@ -46,7 +51,9 @@ class _CacheInvalidationMessage: def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: - message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + message: Final[Mapping[str, object]] = MappingProxyType( + asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + ) return json.dumps({field: value for field, value in message.items() if value is not None}) @@ -62,10 +69,10 @@ def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if not isinstance(data, str): return None try: - parsed: Final = json.loads(data) + parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(data))) except json.JSONDecodeError: return None - if not isinstance(parsed, dict): + if parsed is None: return None cache_key: Final = parsed.get("cache_key") if not isinstance(cache_key, str): diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 46434a997b7..f889748c504 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -39,6 +40,7 @@ from litellm.types.utils import ( ) _CALLBACK_VAR_MASKER: Final = SensitiveDataMasker() +_LOADED_CALLBACK: Final = TypeAdapter(object) # Compound names that are credential-bearing but don't contain any of the # default sensitive segments (so SensitiveDataMasker won't flag them). _EXTRA_SENSITIVE_CALLBACK_KEYS: Final = {"gcs_path_service_account"} @@ -389,9 +391,11 @@ def initialize_callbacks_on_proxy( imported_list.append( _loaded_callback_or_raise( entry=callback, - loaded=get_instance_fn( - value=callback, - config_file_path=config_file_path, + loaded=_LOADED_CALLBACK.validate_python( + get_instance_fn( + value=callback, + config_file_path=config_file_path, + ) ), ) ) @@ -408,9 +412,11 @@ def initialize_callbacks_on_proxy( litellm.callbacks = [ _loaded_callback_or_raise( entry=value, - loaded=get_instance_fn( - value=value, - config_file_path=config_file_path, + loaded=_LOADED_CALLBACK.validate_python( + get_instance_fn( + value=value, + config_file_path=config_file_path, + ) ), ) ] diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 33633150711..ec4baa76514 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -11,12 +11,14 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import TypeAdapter from typing_extensions import ReadOnly from litellm import get_secret_str from litellm._logging import verbose_proxy_logger from litellm.constants import PYTHON_GC_THRESHOLD from litellm.litellm_core_utils.bug_report import EnvironmentReport +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.bug_report_config import build_proxy_environment_report @@ -26,6 +28,7 @@ if TYPE_CHECKING: from litellm.router import Router router: Final = APIRouter() +_DECODED_JSON: Final = TypeAdapter(object) # Configure garbage collection thresholds from environment variables @@ -885,8 +888,8 @@ def init_verbose_loggers(): return if os.path.isfile(worker_config): return - _settings: Final = json.loads(worker_config) - if not isinstance(_settings, dict): + _settings: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(worker_config))) + if _settings is None: return debug: Final = _settings.get("debug", None) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 245a56678eb..e34d03a04c3 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -164,7 +164,7 @@ def coerce_numeric_form_fields( def _parse_binary_body(body: bytes) -> dict: """JSON sent under a binary content type still parses; real binary (protobuf) carries no params -> {}.""" try: - parsed: Final = orjson.loads(body) + parsed: Final = _DECODED_JSON.validate_python(orjson.loads(body)) if isinstance(parsed, dict): return parsed except orjson.JSONDecodeError: diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 4a082eb307b..640a947f9c5 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase _BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_YAML: Final = TypeAdapter(object) class BucketObjectFetcher(Protocol): @@ -31,7 +32,7 @@ class SyncBucketObjectReader(Protocol): def _parsed_config(object_key: str, file_contents: str) -> object | None: try: - parsed: Final = yaml.safe_load(file_contents) + parsed: Final = _DECODED_YAML.validate_python(yaml.safe_load(file_contents)) except yaml.YAMLError as e: verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e) return None diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 089d0b7a226..02da6f706d1 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -37,7 +37,7 @@ This module provides a single proxy-side error class that: import json from collections.abc import Mapping -from typing import Any, Final +from typing import Final from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError @@ -171,7 +171,7 @@ class ProxyRateLimitError(HTTPException, RateLimitError): def __init__( self, - detail: Any, + detail: object, headers: Mapping[str, object] | None = None, category: str | RateLimitErrorCategory = RateLimitErrorCategory.LITELLM_RATE_LIMIT, rate_limit_type: str | RateLimitType | None = None, diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 35bd69ef0bd..a2ffdd4d59d 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -13,6 +13,7 @@ from typing import Final from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -26,6 +27,8 @@ from litellm.proxy.container_endpoints.ownership import ( get_container_forwarding_params, ) +_UPLOADED_FILE_FIELD: Final = TypeAdapter(object) + def _load_endpoints_config() -> dict: """Load the endpoints configuration from JSON file.""" @@ -304,7 +307,7 @@ async def _process_multipart_upload_request( raise HTTPException(status_code=400, detail="Missing required 'file' field") # convert_upload_files_to_file_data returns list of tuples, extract single file - file_list: Final = data["file"] + file_list: Final = _UPLOADED_FILE_FIELD.validate_python(data["file"]) if isinstance(file_list, list) and len(file_list) > 0: data["file"] = file_list[0] diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index aeb0bb1e3d0..9e4ab3c3c4e 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,13 +1,14 @@ import json from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet -from typing import TYPE_CHECKING, Any, Final, TypeAlias +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias from fastapi import HTTPException from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils.core_helpers import has_attributes from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -111,12 +112,20 @@ def _get_response_id(response: object) -> str | None: return value if isinstance(value, str) else None -def _dump_response(response: Any) -> dict[str, object]: +class _SupportsModelDump(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _SupportsPydanticDict(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +def _dump_response(response: object) -> Mapping[str, object]: if isinstance(response, dict): return dict(response) - if hasattr(response, "model_dump"): + if has_attributes(response, _SupportsModelDump, "model_dump"): return response.model_dump() - if hasattr(response, "dict"): + if has_attributes(response, _SupportsPydanticDict, "dict"): return response.dict() return {"id": _get_response_id(response)} @@ -199,9 +208,11 @@ async def record_container_owner( original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) model_object_id: Final = _container_model_object_id(original_container_id, resolved_provider) - file_object: Final = _dump_response(response) - file_object["custom_llm_provider"] = resolved_provider - file_object["provider_container_id"] = original_container_id + file_object: Final = { + **_dump_response(response), + "custom_llm_provider": resolved_provider, + "provider_container_id": original_container_id, + } # Prisma Python requires Json fields to be serialized as a JSON string. file_object_json: Final[str] = json.dumps(file_object) diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 75882bc3fb1..4ddcf7e600b 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -526,7 +526,7 @@ async def _upsert_turn_with_retry( except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise - await asyncio.sleep(2**attempt + random.uniform(0, 1)) + await asyncio.sleep(2.0**attempt + random.uniform(0, 1)) else: return diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index be5b307b5ae..b290e0f85b4 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -16,6 +16,7 @@ from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Mappi from contextlib import asynccontextmanager from contextvars import ContextVar from datetime import datetime, timedelta, timezone +from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload from urllib.parse import quote, unquote @@ -104,6 +105,7 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) _SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_CALL_KWARG_VALUE: Final = TypeAdapter(object) _JsonPayload: TypeAlias = Mapping[str, object] | Sequence[object] | str | int | float | bool | None _JSON_VALUE_ADAPTER: Final[TypeAdapter[_JsonPayload]] = TypeAdapter( _JsonPayload, config=ConfigDict(strict=True, hide_input_in_errors=True) @@ -313,6 +315,20 @@ def _spend_tables_left_to_send( return cast(DBSpendUpdateTransactions, remaining) if any(remaining.values()) else None +class _SpendRestore(TypedDict, total=False): + db_spend_update_transactions: ReadOnly[DBSpendUpdateTransactions] + daily_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_team_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_org_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_end_user_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + daily_agent_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]] + window_spend_update_transactions: ReadOnly[Sequence[WindowSpendTransaction]] + + +def _merged_spend_restore(first: _SpendRestore, second: _SpendRestore) -> _SpendRestore: + return {**first, **second} + + def _report_spend_update_step_failure(step_name: str, *, payload: Mapping[str, object], exc: Exception) -> None: spend_log_error( "Spend tracking - %s failed, this request is missing from that spend rollup. request_id=%s api_key=%s", @@ -845,7 +861,7 @@ class DBSpendUpdateWriter: spend=payload["spend"], total_tokens=payload["total_tokens"], completion_response=completion_response, - realtime_tool_calls=(kwargs or {}).get("realtime_tool_calls"), + realtime_tool_calls=_CALL_KWARG_VALUE.validate_python((kwargs or {}).get("realtime_tool_calls")), ) if transaction is None: return @@ -1627,7 +1643,7 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for spend updates") - uncommitted: dict[str, Any] = {} # mutable-ok: tracks popped categories still needing commit + uncommitted: dict[str, _SpendRestore] = {} # mutable-ok: tracks popped categories still needing commit committed_spend_tables: Final[list[_SpendTableName]] = [] # mutable-ok: filled as each table lands try: @@ -1642,13 +1658,41 @@ class DBSpendUpdateWriter: ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() uncommitted = { - "db_spend_update_transactions": db_spend_update_transactions, - "daily_spend_update_transactions": daily_spend_update_transactions, - "daily_team_spend_update_transactions": daily_team_spend_update_transactions, - "daily_org_spend_update_transactions": daily_org_spend_update_transactions, - "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, - "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, - "window_spend_update_transactions": window_spend_update_transactions, + "db_spend_update_transactions": ( + {} + if db_spend_update_transactions is None + else {"db_spend_update_transactions": db_spend_update_transactions} + ), + "daily_spend_update_transactions": ( + {} + if daily_spend_update_transactions is None + else {"daily_spend_update_transactions": daily_spend_update_transactions} + ), + "daily_team_spend_update_transactions": ( + {} + if daily_team_spend_update_transactions is None + else {"daily_team_spend_update_transactions": daily_team_spend_update_transactions} + ), + "daily_org_spend_update_transactions": ( + {} + if daily_org_spend_update_transactions is None + else {"daily_org_spend_update_transactions": daily_org_spend_update_transactions} + ), + "daily_end_user_spend_update_transactions": ( + {} + if daily_end_user_spend_update_transactions is None + else {"daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions} + ), + "daily_agent_spend_update_transactions": ( + {} + if daily_agent_spend_update_transactions is None + else {"daily_agent_spend_update_transactions": daily_agent_spend_update_transactions} + ), + "window_spend_update_transactions": ( + {} + if window_spend_update_transactions is None + else {"window_spend_update_transactions": window_spend_update_transactions} + ), } if db_spend_update_transactions is not None: @@ -1677,9 +1721,12 @@ class DBSpendUpdateWriter: on_table_committed=committed_spend_tables.append, ) except Exception as e: - uncommitted["db_spend_update_transactions"] = _spend_tables_left_to_send( + left_to_send: Final = _spend_tables_left_to_send( db_spend_update_transactions, committed_spend_tables, e ) + uncommitted["db_spend_update_transactions"] = ( + {} if left_to_send is None else {"db_spend_update_transactions": left_to_send} + ) raise uncommitted.pop("db_spend_update_transactions", None) @@ -1753,7 +1800,7 @@ class DBSpendUpdateWriter: exc=e, ) finally: - to_restore = {name: txns for name, txns in uncommitted.items() if txns is not None} + to_restore: Final = reduce(_merged_spend_restore, uncommitted.values(), _SpendRestore()) if to_restore: await self.redis_update_buffer.restore_transactions_to_redis(**to_restore) await self.pod_lock_manager.release_lock( diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 627b30e83f7..1116cd604a3 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -1,7 +1,7 @@ import asyncio import json import logging -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -14,7 +14,9 @@ from litellm.types.services import ServiceTypes POD_LOCK_TARGET: Final = "pod_lock" if TYPE_CHECKING: - ProxyLogging = Any + from litellm.proxy.utils import ProxyLogging as _ProxyLogging + + ProxyLogging: TypeAlias = _ProxyLogging else: ProxyLogging = Any diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 54d20bd4507..66833fd688a 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -11,6 +11,7 @@ from datetime import datetime from functools import reduce from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast +from pydantic import TypeAdapter from redis.exceptions import RedisError from litellm._internal_context import with_service_target @@ -98,6 +99,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( ) _ValueT = TypeVar("_ValueT") +_DECODED_JSON: Final = TypeAdapter(object) def _spend_log_json_default(value: object) -> str: @@ -109,7 +111,7 @@ def _encode_spend_log_row(row: SpendLogRow) -> str: def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: - decoded: Final = json.loads(encoded) + decoded: Final = _DECODED_JSON.validate_python(json.loads(encoded)) return decoded if isinstance(decoded, dict) else None diff --git a/litellm/proxy/db/prisma_query_span.py b/litellm/proxy/db/prisma_query_span.py index 75625c6c0f2..2747a765f2a 100644 --- a/litellm/proxy/db/prisma_query_span.py +++ b/litellm/proxy/db/prisma_query_span.py @@ -20,6 +20,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from litellm.integrations.otel.model.spans import PG_CATALOG, PRISMA_RELATIONS @@ -34,6 +36,7 @@ class PrismaQuery: UNKNOWN_PRISMA_QUERY: Final = PrismaQuery("prisma_query", None, None) +_DECODED_JSON: Final = TypeAdapter(object) _ROOT_FIELD: Final = re.compile(r"result:\s*(\w+)") _RAW_SQL: Final = re.compile(r'query:\s*"((?:[^"\\]|\\.)*)"') @@ -111,7 +114,7 @@ def sql_operation(sql: str) -> tuple[str | None, str | None]: def _query_text(content: str) -> str: try: - payload: Final[object] = json.loads(content) + payload: Final = _DECODED_JSON.validate_python(json.loads(content)) except ValueError: return content query: Final = payload.get("query") if isinstance(payload, dict) else None diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 673b549b645..3214c0974b4 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -56,6 +56,7 @@ class _LiveWorkerCountRow(TypedDict): _COUNT_ROWS_ADAPTER: Final = TypeAdapter(tuple[_LiveWorkerCountRow, ...]) +_RAW_QUERY_RESULT: Final = TypeAdapter(object) class _HeartbeatExecute(Protocol): @@ -99,7 +100,9 @@ async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: db: Final = prisma_client.db primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db async with db_span("count_live_proxy_workers", "LiteLLM_ProxyWorkerHeartbeat"): - rows: Final = await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + rows: Final = _RAW_QUERY_RESULT.validate_python( + await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + ) return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 48ad57fd06e..52a0be3d49e 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -298,7 +298,7 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: - return getattr(self.read_target, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self.read_target, name)) writer_attr: Final = _DELEGATED_ATTRIBUTE.validate_python(getattr(self._writer, name)) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 7ac414bec00..591ad81a4e5 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -19,7 +19,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from itertools import chain, groupby -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Final, Literal, Protocol from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import LiteralString @@ -122,7 +122,7 @@ def build_tool_usage_transaction( spend: float, total_tokens: int, completion_response: object, - realtime_tool_calls: Any = None, + realtime_tool_calls: object = None, ) -> ToolUsageTransaction | None: """None when the request invoked no tools. Realtime sessions carry invoked tools in kwargs["realtime_tool_calls"] (OpenAI tool_calls shape) rather than diff --git a/litellm/proxy/discovery_endpoints/agent_skills_archive.py b/litellm/proxy/discovery_endpoints/agent_skills_archive.py index 1f2fca3992e..94a0c3e2cf1 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_archive.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_archive.py @@ -16,6 +16,7 @@ from types import MappingProxyType from typing import Final import yaml +from pydantic import TypeAdapter MAX_ARCHIVE_UNPACKED_BYTES: Final = 50 * 1024 * 1024 MAX_ARCHIVE_ENTRIES: Final = 1000 @@ -26,6 +27,7 @@ _ZIP_ENTRY_PERMISSIONS: Final = 0o644 << 16 _FRONTMATTER_PATTERN: Final = re.compile(r"^---\s*\n(.*?)\n---\s*(?:\n|$)", re.DOTALL) _WINDOWS_DRIVE_PATTERN: Final = re.compile(r"^[A-Za-z]:") _EMPTY_FRONTMATTER: Final[Mapping[str, object]] = MappingProxyType({}) +_DECODED_YAML: Final = TypeAdapter(object) @dataclass(frozen=True, slots=True) @@ -100,7 +102,7 @@ def _manifest_frontmatter(manifest: bytes) -> Mapping[str, object]: if match is None: return _EMPTY_FRONTMATTER try: - parsed: Final = yaml.safe_load(match.group(1)) + parsed: Final = _DECODED_YAML.validate_python(yaml.safe_load(match.group(1))) except yaml.YAMLError: return _EMPTY_FRONTMATTER if not isinstance(parsed, dict): diff --git a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py index 2050cc40e6d..830f7b0334e 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -10,7 +10,6 @@ import asyncio import re from collections.abc import Sequence from itertools import groupby -from operator import itemgetter from types import MappingProxyType from typing import Final @@ -173,11 +172,15 @@ def _slugify(raw: str | None) -> str | None: return _NON_SLUG_PATTERN.sub("-", raw.lower()).strip("-")[:MAX_SKILL_NAME_LENGTH].rstrip("-") or None +def _entry_name(entry: tuple[int, str]) -> str: + return entry[1] + + def _deduplicated(names: Sequence[str]) -> tuple[str, ...]: ordinals: Final = MappingProxyType( { position: ordinal - for _, duplicates in groupby(sorted(enumerate(names), key=itemgetter(1)), key=itemgetter(1)) + for _, duplicates in groupby(sorted(enumerate(names), key=_entry_name), key=_entry_name) for ordinal, (position, _) in enumerate(duplicates) } ) diff --git a/litellm/proxy/example_config_yaml/custom_auth.py b/litellm/proxy/example_config_yaml/custom_auth.py index b7646ce5e3d..c53074638c6 100644 --- a/litellm/proxy/example_config_yaml/custom_auth.py +++ b/litellm/proxy/example_config_yaml/custom_auth.py @@ -30,7 +30,7 @@ async def generate_key_fn(data: GenerateKeyRequest): data_json: Final = data.json() # Unpacking variables - team_id: Final = data_json.get("team_id") + team_id: Final = data.team_id data_json.get("duration") data_json.get("models") data_json.get("aliases") diff --git a/litellm/proxy/example_config_yaml/custom_callbacks.py b/litellm/proxy/example_config_yaml/custom_callbacks.py index 2c689811dc8..5ae55060930 100644 --- a/litellm/proxy/example_config_yaml/custom_callbacks.py +++ b/litellm/proxy/example_config_yaml/custom_callbacks.py @@ -7,10 +7,14 @@ import traceback sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import inspect +from typing import Final + +from pydantic import TypeAdapter import litellm from litellm.integrations.custom_logger import CustomLogger -from typing import Final + +_ATTRIBUTE_VALUE: Final = TypeAdapter(object) # This file includes the custom callbacks for LiteLLM Proxy @@ -27,7 +31,11 @@ class MyCustomHandler(CustomLogger): print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger") try: print_verbose("Logger Initialized with following methods:") - methods: Final = [method for method in dir(self) if inspect.ismethod(getattr(self, method))] + methods: Final = [ + method + for method in dir(self) + if inspect.ismethod(_ATTRIBUTE_VALUE.validate_python(getattr(self, method))) + ] # Pretty print_verbose the methods for method in methods: diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 17bea031794..8581dca603b 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -9,8 +9,9 @@ every text fragment. """ from collections.abc import Callable, Iterator, Mapping, Sequence -from typing import Any, Final +from typing import Final +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.responses.litellm_completion_transformation.compaction_items import ( compaction_summary, input_items, @@ -231,7 +232,7 @@ def is_string_batch_input(data: Mapping[str, object]) -> bool: return isinstance(input_value, list) and bool(input_value) and all(isinstance(item, str) for item in input_value) -def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequence[object]) -> bool: +def apply_redacted_messages_back(data: dict[str, object], redacted_messages: Sequence[object]) -> bool: """Write redacted messages back to whichever field(s) the caller used. Mask/anonymize paths take a synthesised messages list (from @@ -249,9 +250,9 @@ def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequen carry exactly one message per inspected element: a partial rewrite would forward the remaining originals unredacted. Callers must block on False. """ + input_value: Final = data.get("input") if is_string_batch_input(data): - batch: Final = data["input"] - inspected_indices: Final = tuple(idx for idx, item in enumerate(batch) if item) + inspected_indices: Final = tuple(idx for idx, item in enumerate(iter_items(input_value)) if item) if len(redacted_messages) != len(inspected_indices): return False if any(not isinstance(message, Mapping) or message.get("content") is None for message in redacted_messages): @@ -259,12 +260,12 @@ def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequen redacted_texts: Final = tuple( "\n".join(_iter_text_parts_in_content(message["content"])) for message in redacted_messages ) - for idx, text in zip(inspected_indices, redacted_texts): - batch[idx] = text + if isinstance(input_value, list): + for idx, text in zip(inspected_indices, redacted_texts): + input_value[idx] = text return True if "messages" in data: data["messages"] = redacted_messages - input_value: Final = data.get("input") if isinstance(input_value, str): text_parts: Final[list[str]] = [] for msg in redacted_messages: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 70e7aac12b5..a5822bfb0e3 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -90,6 +90,7 @@ if TYPE_CHECKING: router: Final = APIRouter() GUARDRAIL_REGISTRY: Final = GuardrailRegistry() _DECODED_JSON: Final = TypeAdapter(object) +_DECODED_YAML: Final = TypeAdapter(object) def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]: @@ -1682,7 +1683,7 @@ async def validate_blocked_words_file(request: dict[str, str]): if not file_content: return {"valid": False, "error": "No file content provided"} - data: Final = yaml.safe_load(file_content) + data: Final = _DECODED_YAML.validate_python(yaml.safe_load(file_content)) if not isinstance(data, dict) or "blocked_words" not in data: return { diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index ac0a6911184..aabac874ca2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -63,6 +63,7 @@ _INVALID_ASSERTION_AADSTS_PREFIX: Final = "50027" _AADSTS_CODES_ADAPTER: Final = TypeAdapter(tuple[int, ...]) _MCP_CALL_TYPES: Final[tuple[str, ...]] = ("mcp_call", "call_mcp_tool") _TOOL_INPUT_SCHEMA_ADAPTER: Final = TypeAdapter(dict[str, object]) +_DECODED_JSON: Final = TypeAdapter(object) _OBO_CACHE_MAX_ENTRIES: Final = 1000 _DEFAULT_TOKEN_TTL_SECONDS: Final = 3599.0 _TOKEN_EXPIRY_SLACK_SECONDS: Final = 60.0 @@ -343,7 +344,7 @@ class Agent365Guardrail(CustomGuardrail): latency_ms: float, ) -> dict: # mutable-ok: returns the request data dict per hook contract try: - parsed_verdict: Final = response.json() + parsed_verdict: Final = _DECODED_JSON.validate_python(response.json()) except ValueError: return self._handle_unavailable( data=data, @@ -486,7 +487,7 @@ class Agent365Guardrail(CustomGuardrail): response=response, ) try: - parsed_body: Final = response.json() + parsed_body: Final = _DECODED_JSON.validate_python(response.json()) except ValueError as exc: raise Agent365MalformedResponseError("the Entra token endpoint returned a non-JSON body") from exc if not isinstance(parsed_body, dict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index 993c1d30949..389706d7417 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -2,12 +2,13 @@ import os import uuid from collections.abc import Mapping from typing import ( - Any, # noqa: TID251 # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs Final, Literal, ) import httpx +from pydantic import TypeAdapter +from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -36,6 +37,12 @@ DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy. SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call) +_DECODED_JSON: Final = TypeAdapter(object) + + +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass + def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool: """Whether this is a streamed response on a route the framework can only block, not rewrite. @@ -97,7 +104,7 @@ class AiriaGuardrail(CustomGuardrail): timeout: float | None = None, supported_event_hooks: list[GuardrailEventHooks] | None = None, # mutable-ok: matches CustomGuardrail.__init__ async_handler: AsyncHTTPHandler | None = None, # dependency-injected in tests; defaults to a real client - **kwargs: Any, # kwargs-ok: passed straight through to CustomGuardrail.__init__ + **kwargs: Unpack[_CustomGuardrailOptions], # kwargs-ok: passed straight through to CustomGuardrail.__init__ ) -> None: resolved_timeout: Final = timeout or float(os.getenv("AIRIA_TIMEOUT", "10")) @@ -161,7 +168,7 @@ class AiriaGuardrail(CustomGuardrail): headers={"Authorization": f"Bearer {self.api_key}"}, ) response.raise_for_status() - raw_body: Final[object] = response.json() + raw_body: Final = _DECODED_JSON.validate_python(response.json()) body: Final = _as_object(raw_body) if body is None: raise TypeError(f"expected a JSON object in the response body, got {type(raw_body).__name__}") diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py index 64770fdf0f7..5bf471265c8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations from .block_code_execution import BlockCodeExecutionGuardrail @@ -26,7 +27,7 @@ def _get_param( value: Final = getattr(litellm_params, key, default) if value is not None: return value - raw: Final = guardrail.get("litellm_params") + raw: Final = as_str_mapping(guardrail.get("litellm_params")) if isinstance(raw, dict) and key in raw: return raw[key] return default diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 936f862b10b..19b9bccd8aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -23,6 +23,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_ssl_configuration, @@ -222,7 +223,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return [] @staticmethod - def _iter_schema_string_refs(data: Mapping[str, Any]): + def _iter_schema_string_refs(data: Mapping[str, object]): """Yield ``(container, key)`` for every non-empty schema string the proxy forwards to the model inside tool/function and structured-output schemas: each ``tools[].function`` and legacy ``functions[]`` entry plus the @@ -235,10 +236,10 @@ class CatoNetworksGuardrail(CustomGuardrail): list_keys: Final = ("enum", "examples") stack: Final[list] = [] - for tool in data.get("tools") or []: + for tool in iter_items(data.get("tools") or ()): if isinstance(tool, dict) and isinstance(tool.get("function"), dict): stack.append(tool["function"]) - for function in data.get("functions") or []: + for function in iter_items(data.get("functions") or ()): if isinstance(function, dict): stack.append(function) response_format: Final = data.get("response_format") diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index ab43d32e260..45e785fd53c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -35,6 +35,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -1743,7 +1744,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_masked_entity_count( - rules: list[dict[str, Any]], + rules: Sequence[Mapping[str, object]], ) -> dict[str, int] | None: """Count entity-type detections per Cisco rule for the logging payload.""" if not rules: @@ -1752,8 +1753,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): for rule in rules: if not isinstance(rule, dict): continue - entity_types = rule.get("entity_types") or [] - for entity_type in entity_types: + for entity_type in iter_items(rule.get("entity_types") or ()): if not isinstance(entity_type, str): continue counts[entity_type] = counts.get(entity_type, 0) + 1 diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py index 498f9bf4099..c5ca97ff760 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Final +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.guardrails import ( GuardrailEventHooks, Mode, @@ -26,7 +27,7 @@ def _coerce_event_hook( def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object: if optional_params is not None: - value: Final = getattr(optional_params, attribute_name, None) + value: Final = attribute_or_none(optional_params, attribute_name) if value is not None: return value return getattr(litellm_params, attribute_name, None) diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 029feba89fb..4d910cde603 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -28,6 +28,7 @@ from urllib.parse import urlparse import httpx from fastapi import HTTPException from httpx import Response as HttpxResponse +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -72,6 +73,7 @@ if TYPE_CHECKING: ) BYPASS_HEADER: Final = "x-compresr-bypass" +_AGENTIC_LOOP_MODEL: Final = TypeAdapter(object) COMPRESR_RETRIEVE_TOOL_NAME: Final = "compresr_retrieve" DEFAULT_API_BASE: Final = "https://api.compresr.ai" DEFAULT_COMPRESSION_MODEL: Final = "latte_v2" @@ -1118,7 +1120,7 @@ class CompresrGuardrail(CustomGuardrail): full_model_name = model if logging_obj is not None: agentic_params: Final = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) - candidate: Final = agentic_params.get("model", model) + candidate: Final = _AGENTIC_LOOP_MODEL.validate_python(agentic_params.get("model", model)) if isinstance(candidate, str) and candidate: full_model_name = candidate diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 34725803c60..ee3a957f607 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -5,7 +5,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Optional, cast from fastapi import HTTPException -from pydantic import ConfigDict, Field, ValidationError +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import override from litellm._logging import verbose_proxy_logger @@ -37,6 +37,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + class CrowdStrikeAIDRGuardrailMissingSecrets(Exception): """Custom exception for missing CrowdStrike AIDR secrets.""" @@ -362,7 +364,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): assert response is not None response.raise_for_status() - response_body: Final[object] = response.json() + response_body: Final = _DECODED_JSON.validate_python(response.json()) raw_result: Final[object] = response_body.get("result") if isinstance(response_body, dict) else None blocked_signal: Final[object] = raw_result.get("blocked") if isinstance(raw_result, dict) else None diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py index 13d7178dfe9..0f9ce15e65b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -39,7 +39,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" raise ValueError("Custom code guardrail requires a guardrail_name") # Get the custom code from litellm_params - custom_code: Final = getattr(litellm_params, "custom_code", None) + custom_code: Final = litellm_params.custom_code if not custom_code: raise ValueError("Custom code guardrail requires 'custom_code' in litellm_params") diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 199da7df30d..ecd23d86540 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, @@ -126,6 +127,13 @@ def _sanitize_inbound_headers( return sanitized or None +def _logged_inbound_headers(details: Mapping[str, object]) -> Mapping[object, object] | None: + litellm_params: Final = as_str_mapping(details.get("litellm_params", {})) + metadata: Final = as_str_mapping(litellm_params.get("metadata", {})) if litellm_params is not None else None + headers: Final = metadata.get("headers", None) if metadata is not None else None + return headers if isinstance(headers, dict) else None + + def _extract_inbound_headers( request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"], @@ -163,7 +171,7 @@ def _extract_inbound_headers( if logging_obj and getattr(logging_obj, "model_call_details", None): try: details: Final = logging_obj.model_call_details or {} - headers = details.get("litellm_params", {}).get("metadata", {}).get("headers", None) + headers = _logged_inbound_headers(details) if headers: return _sanitize_inbound_headers(headers, extra_allowlist=extra_allowlist) except Exception: @@ -283,7 +291,7 @@ class GenericGuardrailAPI(CustomGuardrail): ) self.num_retries: int = num_retries or 0 self.headers = headers or {} - self.extra_headers = extra_headers or [] + self.extra_headers: Sequence[object] = extra_headers or [] # If api_key is provided, add it as x-api-key header if api_key: diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index f32228a6204..7ff5aa2466b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj GRAYSWAN_BLOCK_ERROR_MSG: Final = "Blocked by Gray Swan Guardrail" +_DECODED_JSON: Final = TypeAdapter(object) class _GraySwanMonitorResponse(TypedDict): @@ -705,7 +706,9 @@ class GraySwanGuardrail(CustomGuardrail): {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers ) if cleaned_litellm_metadata: - sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + sanitized: Final = _DECODED_JSON.validate_python( + safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + ) if isinstance(sanitized, dict) and sanitized: payload["litellm_metadata"] = sanitized diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index d5edcc19a02..a1d0900ebed 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -2,6 +2,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Final from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -23,6 +24,8 @@ from litellm.types.utils import CallTypesLiteral, GuardrailStatus if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + class JavelinGuardrail(CustomGuardrail): @classmethod @@ -113,7 +116,9 @@ class JavelinGuardrail(CustomGuardrail): json=dict(request), timeout=self.timeout, ) - verbose_proxy_logger.debug("Javelin Guardrail: Javelin guard API response: %s", response.json()) + verbose_proxy_logger.debug( + "Javelin Guardrail: Javelin guard API response: %s", _DECODED_JSON.validate_python(response.json()) + ) response_data: Final = response.json() # Ensure the response has the required assessments field if "assessments" not in response_data: diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 8efd1dc79e3..7c0cdb0d36f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import Final, Literal from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -40,6 +41,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( ) from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponse +_DECODED_JSON: Final = TypeAdapter(object) +_RESPONSE_FIELDS: Final = TypeAdapter(Mapping[str, object]) + _DETECTOR_CATEGORY_PHRASES: Final[Mapping[str, str]] = MappingProxyType( { "prompt_injection": "a potential prompt injection attempt", @@ -402,7 +406,9 @@ class LakeraAIGuardrail(CustomGuardrail): json=request, timeout=self.timeout, ) - verbose_proxy_logger.debug("Lakera AI v2 guard response: %s", response.json()) + verbose_proxy_logger.debug( + "Lakera AI v2 guard response: %s", _DECODED_JSON.validate_python(response.json()) + ) lakera_response = LakeraAIResponse(**response.json()) return lakera_response, masked_entity_count except Exception as e: @@ -801,7 +807,7 @@ class LakeraAIGuardrail(CustomGuardrail): choice_idx = choice_indices[idx] response_dict["choices"][choice_idx]["message"]["content"] = msg.get("content", "") add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) - return ModelResponse(**response_dict) + return ModelResponse(**_RESPONSE_FIELDS.validate_python(response_dict)) # inject_system_message has nothing left to inject into once a response # already exists, so it is treated the same as monitor: log and allow. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index e05db5003fa..97b35769f63 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -9,7 +9,7 @@ import json import os import uuid from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict +from typing import TYPE_CHECKING, Final, Literal, TypedDict try: import ulid @@ -39,6 +39,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.integrations.custom_guardrail import dc as global_cache +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -397,13 +398,15 @@ class LassoGuardrail(CustomGuardrail): HTTPException: If blocking violations are detected """ raw_messages: Final[list[dict[str, object]]] = data.get("messages") or [] - messages: list[dict[str, Any]] = self._expand_messages_for_classification(raw_messages) if raw_messages else [] - messages_count: Final = len(messages) - if data.get("input") is not None: - # Responses-API payloads carry text in data["input"]. Inspect it - # alongside any "messages" array — otherwise a caller can attach - # benign messages and stash blocked content in input to bypass. - messages.extend(build_inspection_messages({"input": data["input"]})) + expanded_messages: Final = self._expand_messages_for_classification(raw_messages) if raw_messages else [] + messages_count: Final = len(expanded_messages) + # Responses-API payloads carry text in data["input"]. Inspect it + # alongside any "messages" array — otherwise a caller can attach + # benign messages and stash blocked content in input to bypass. + input_messages: Final = ( + build_inspection_messages({"input": data["input"]}) if data.get("input") is not None else [] + ) + messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages] if not messages: return data @@ -421,7 +424,7 @@ class LassoGuardrail(CustomGuardrail): data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], ) -> dict: """Handle classification without masking.""" try: @@ -439,7 +442,7 @@ class LassoGuardrail(CustomGuardrail): data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], messages_count: int, ) -> dict: """Handle masking with classifix endpoint. @@ -470,7 +473,7 @@ class LassoGuardrail(CustomGuardrail): text_parts = [msg["content"] for msg in masked_for_input if isinstance(msg.get("content"), str)] if text_parts: data["input"] = "\n".join(text_parts) - self._log_masking_applied(message_type, dict(response)) + self._log_masking_applied(message_type, response) return data except Exception as e: @@ -479,7 +482,7 @@ class LassoGuardrail(CustomGuardrail): def _map_masked_messages_back( self, - original_messages: list[dict[str, Any]], + original_messages: Sequence[Mapping[str, object]], masked_messages: Sequence[Mapping[str, object]], ) -> list[dict[str, object]]: """Map Lasso-format masked messages back onto the original OpenAI-format messages. @@ -492,7 +495,7 @@ class LassoGuardrail(CustomGuardrail): """ # Index masked content by type so we can look up by id without caring about order. masked_tool_use: Final[dict[object, dict[str, object]]] = {} - masked_tool_result: Final[dict[str, str]] = {} + masked_tool_result: Final[dict[object, str]] = {} masked_text: Final[list[str]] = [] for msg in masked_messages: @@ -567,12 +570,12 @@ class LassoGuardrail(CustomGuardrail): def _update_tool_calls_from_masked( self, - tool_calls: list[object], + tool_calls: object, masked_tool_use: Mapping[object, Mapping[str, object]], ) -> list[object]: """Replace tool_call arguments with masked values returned by Lasso.""" updated: Final = [] - for call in tool_calls: + for call in iter_items(tool_calls): call_id = self._get_field(call, "id") if call_id and call_id in masked_tool_use: masked_input = masked_tool_use[call_id].get("input") @@ -627,7 +630,7 @@ class LassoGuardrail(CustomGuardrail): def _log_masking_applied( self, message_type: Literal["PROMPT", "COMPLETION"], - response: dict[str, Any], + response: LassoResponse, ) -> None: """Log masking application with structured context.""" conversation_id: Final = getattr(self, "conversation_id", "unknown") @@ -637,12 +640,12 @@ class LassoGuardrail(CustomGuardrail): "guardrail_name": getattr(self, "guardrail_name", "unknown"), "message_type": message_type, "violations_count": len(response.get("findings", {})), - "masked_fields": len(response.get("messages", [])), + "masked_fields": len(response.get("messages") or []), "conversation_id": conversation_id, }, ) - def _expand_messages_for_classification(self, messages: list[dict[str, Any]]) -> list[dict[str, object]]: + def _expand_messages_for_classification(self, messages: Sequence[Mapping[str, object]]) -> list[dict[str, object]]: """ Convert raw OpenAI-format messages to Lasso API format with content blocks. @@ -701,7 +704,7 @@ class LassoGuardrail(CustomGuardrail): expanded.append({"role": role, "content": content}) if role == "assistant": - for call in msg.get("tool_calls") or []: + for call in iter_items(msg.get("tool_calls") or []): call_id, name, input_data = self._extract_tool_call_fields(call) if not call_id or not name: verbose_proxy_logger.warning( @@ -749,7 +752,7 @@ class LassoGuardrail(CustomGuardrail): def _prepare_payload( self, - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py index 5e2b85250d2..76c3a86d701 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py @@ -13,6 +13,9 @@ import json from pathlib import Path from typing import Any, Final +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base import ( # noqa: F401 # legacy module exports BaseCompetitorIntentChecker, _compile_marker, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -114,6 +117,7 @@ AIRLINE_EXPLICIT_OTHER_MEANING_MARKER: Final = ( ) _MAJOR_AIRLINES_PATH: Final = Path(__file__).resolve().parent / "major_airlines.json" +_DECODED_JSON: Final = TypeAdapter(object) def _load_competitors_excluding_brand(brand_self: list[str]) -> list[str]: @@ -127,11 +131,11 @@ def _load_competitors_excluding_brand(brand_self: list[str]) -> list[str]: return [] try: with open(_MAJOR_AIRLINES_PATH, encoding="utf-8") as f: - airlines: Final = json.load(f) + airlines: Final = _DECODED_JSON.validate_python(json.load(f)) except (json.JSONDecodeError, OSError): return [] result: Final[list[str]] = [] - for entry in airlines: + for entry in iter_items(airlines): if not isinstance(entry, dict): continue match_str = entry.get("match") or "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 40e44151e65..f58c8705815 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -78,6 +78,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_JSON_ENTRIES: Final = TypeAdapter(list[object]) GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b") SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+") @@ -699,9 +700,8 @@ class ContentFilterGuardrail(CustomGuardrail): Severity mapping: 4,3 -> high; 2 -> medium; 1 -> low. """ with open(file_path, "r") as f: - entries = json.load(f) - if not isinstance(entries, list): - entries = [entries] + decoded: Final = _DECODED_ARGUMENTS.validate_python(json.load(f)) + entries: Final = _JSON_ENTRIES.validate_python(decoded) if isinstance(decoded, list) else [decoded] # Derive category name from filename (e.g. harm_toxic_abuse.json -> harm_toxic_abuse) category_name: Final = os.path.splitext(os.path.basename(file_path))[0] severity_map: Final = {4: "high", 3: "high", 2: "medium", 1: "low"} @@ -806,7 +806,7 @@ class ContentFilterGuardrail(CustomGuardrail): """ try: with open(file_path, "r") as f: - data: Final = yaml.safe_load(f) + data: Final = _DECODED_ARGUMENTS.validate_python(yaml.safe_load(f)) if not isinstance(data, dict) or "blocked_words" not in data: raise ValueError("Invalid format: file must contain 'blocked_words' key with list of words") diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index 7fb87080cd1..d2dafb04fe0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator, AsyncIterable, Iterator, Mapping, Sequence from typing import ( TYPE_CHECKING, - Any, # noqa: TID251 # the streaming hook and the flush copy pass provider chunks through untyped + Any, # noqa: TID251 # the streaming hook passes provider chunks through untyped ClassVar, Final, Literal, @@ -23,7 +23,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.litellm_core_utils.core_helpers import attribute_items +from litellm.litellm_core_utils.core_helpers import attribute_items, attribute_or_none, iter_items, require_attributes from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -99,13 +99,8 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): pass -class _ContentDelta(Protocol): - content: object - - -class _FinishingDelta(Protocol): - content: object - tool_calls: object +class _DeepCopyable(Protocol): + def model_copy(self, *, deep: bool) -> object: ... class LLMShieldProxyGuardrail(CustomGuardrail): @@ -528,7 +523,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): one tool call's accumulating `arguments`. A shared window would splice the text held back for one stream onto another. """ - delta: Final = getattr(choice, "delta", None) + delta: Final = attribute_or_none(choice, "delta") index: Final = choice_index(choice) is_final: Final = bool(getattr(choice, "finish_reason", None)) if isinstance(choice, TextChoices): @@ -566,7 +561,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): async def _restore_content_window( self, - delta: _ContentDelta, + delta: object, key: CarryKey, carries: CarryWindows, session_id: str, @@ -574,19 +569,19 @@ class LLMShieldProxyGuardrail(CustomGuardrail): ) -> None: """Restores one delta's content through its own window.""" carry: Final = carries.get(key, "") - text: Final = getattr(delta, "content", None) + text: Final = attribute_or_none(delta, "content") if not isinstance(text, str) or not text: if is_final and carry: flushed, remaining = await self._stream_step("", carry, True, session_id) carries[key] = remaining # rebind-ok: this stream's window advances. if flushed: - delta.content = flushed + write_field(delta, "content", flushed) return emitted, remaining = await self._stream_step(text, carry, is_final, session_id) carries[key] = remaining # rebind-ok: this stream's window advances. - delta.content = emitted + write_field(delta, "content", emitted) async def _restore_tool_call_window( self, @@ -618,7 +613,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): async def _flush_finished_choice( self, - delta: _FinishingDelta, + delta: object, choice_index: int, carries: CarryWindows, session_id: str, @@ -643,12 +638,12 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if not text: continue if tool_index is None: - delta.content = text + write_field(delta, "content", text) else: continuations.extend(continuation_delta(tool_index, text)) if continuations: - existing: Final = tuple(getattr(delta, "tool_calls", None) or ()) - delta.tool_calls = [*existing, *continuations] + existing: Final = tuple(attribute_items(delta, "tool_calls")) + write_field(delta, "tool_calls", [*existing, *continuations]) async def _flush_trailing( self, last_chunk: object, carries: CarryWindows, session_id: str @@ -676,7 +671,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): chunk = self._chunk_for_choice(last_chunk, choice_index) if chunk is None: continue - choice: object = chunk.choices[0] + choice = read_list(chunk, "choices")[0] delta = read_field(choice, "delta") if isinstance(choice, TextChoices): choice.text = text @@ -688,26 +683,26 @@ class LLMShieldProxyGuardrail(CustomGuardrail): yield chunk @staticmethod - def _chunk_for_choice(last_chunk: Any, index: int) -> Any: + def _chunk_for_choice(last_chunk: object, index: int) -> object | None: """A single-choice copy of the last chunk, carrying only `index`. Emitting one choice per chunk keeps a flush from reading as content on a choice it does not belong to. """ - chunk: Final = last_chunk.model_copy(deep=True) - raw_choices: Final = getattr(chunk, "choices", None) + chunk: Final = require_attributes(last_chunk, _DeepCopyable, "model_copy").model_copy(deep=True) + raw_choices: Final = attribute_or_none(chunk, "choices") if not raw_choices: return None - choices: Final[tuple[object, ...]] = tuple(raw_choices) + choices: Final = tuple(iter_items(raw_choices)) position: Final = next((at for at, choice in enumerate(choices) if choice_index(choice) == index), 0) - kept: Final = raw_choices[position] - if getattr(kept, "delta", None) is None and not isinstance(kept, TextChoices): + kept: Final = choices[position] + if attribute_or_none(kept, "delta") is None and not isinstance(kept, TextChoices): return None - kept.index = index - kept.finish_reason = None - chunk.choices = [kept] + write_field(kept, "index", index) + write_field(kept, "finish_reason", None) + write_field(chunk, "choices", [kept]) if hasattr(chunk, "usage"): - del chunk.usage + delattr(chunk, "usage") return chunk async def _stream_step(self, text: str, carry: str, final: bool, session_id: str) -> tuple[str, str]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py index 4ef8bbdbb3e..ef007155c46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -7,9 +7,12 @@ from enum import Enum from types import MappingProxyType from typing import ( Final, + Protocol, TypeAlias, ) +from pydantic import TypeAdapter + from .payload import ( JsonBody, MutableRequest, @@ -24,6 +27,8 @@ from .payload import ( write_field, ) +_OPAQUE_VALUE: Final = TypeAdapter(object) + ANTHROPIC_DELTA_FIELDS: Final = MappingProxyType({"text_delta": "text", "input_json_delta": "partial_json"}) SSE_EVENT_BOUNDARY: Final = re.compile(rb"(\r?\n\r?\n)") @@ -80,6 +85,15 @@ def opens_like_sse(head: bytes) -> bool | None: return False +class _ValuedEnum(Protocol): + @property + def value(self) -> object: ... + + +def _enum_value(member: _ValuedEnum) -> object: + return member.value + + def responses_event_type(chunk: object) -> str | None: """The event type of a Responses API stream event, or None for any other chunk. @@ -90,7 +104,7 @@ def responses_event_type(chunk: object) -> str | None: if isinstance(chunk, (bytes, str)): return None kind: Final = read_field(chunk, "type") - value: Final = kind.value if isinstance(kind, Enum) else kind + value: Final = _enum_value(kind) if isinstance(kind, Enum) else kind return value if isinstance(value, str) and value.startswith("response.") else None @@ -173,7 +187,7 @@ class AnthropicSSERestorer: return block line: Final = lines[data_lines[0]] try: - parsed: Final[object] = json.loads(line[len("data:") :]) + parsed: Final = _OPAQUE_VALUE.validate_python(json.loads(line[len("data:") :])) except ValueError: return block event: Final = as_object(parsed) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py index 95db4bd4f77..fa64d306f22 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Any, Final, cast +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol, cast from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -8,13 +10,22 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams +class _ParamsLookup(Protocol): + def get(self, key: str, /) -> object: ... + + +_NO_PARAMS: Final = MappingProxyType[str, object]({}) + + def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm # Default to always-on. Only disable if the user explicitly sets default_on: false. # We check the raw guardrail dict because LitellmParams normalizes None → False, # making it impossible to distinguish "not set" from "explicitly false" via litellm_params. - _raw_default_on: Final = cast(dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") + _raw_default_on: Final = ( + cast(Mapping[str, _ParamsLookup], guardrail).get("litellm_params", _NO_PARAMS).get("default_on") + ) _default_on: Final = _raw_default_on is not False _callback: Final = MCPEndUserPermissionGuardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 221c4b3752b..2e3f7120726 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -74,14 +74,14 @@ import os import re import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional import httpx import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -90,7 +90,6 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -98,6 +97,10 @@ if TYPE_CHECKING: from jwt.types import Options +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass + + class _OIDCDiscoveryDocument(TypedDict, total=False): jwks_uri: str @@ -279,11 +282,13 @@ class MCPJWTSigner(CustomGuardrail): debug_headers: bool = False, # FR-10: Configurable scopes allowed_scopes: list[str] | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs - super().__init__(**base_kwargs) + options: Final[_CustomGuardrailOptions] = { + "supported_event_hooks": list(self.get_supported_event_hooks()), + **kwargs, + } + super().__init__(**options) # --- Signing key setup --- key_material: Final = os.environ.get(self.SIGNING_KEY_ENV) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 8482567b6da..3cb20a00e2f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -6,7 +6,7 @@ on the LiteLLM gateway. Blocks or alerts when unregistered servers are found. """ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException @@ -21,6 +21,9 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( ) from litellm.types.guardrails import GuardrailEventHooks +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + class MCPSecurityGuardrail(CustomGuardrail): def __init__( @@ -40,7 +43,7 @@ class MCPSecurityGuardrail(CustomGuardrail): async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: Any, + cache: "DualCache", data: dict, call_type: str, ) -> Exception | str | dict | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py index ed955ac829d..06de4ef0e93 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py @@ -1,5 +1,6 @@ from typing import TYPE_CHECKING, Final +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.types.guardrails import SupportedGuardrailIntegrations from .purview_dlp import MicrosoftPurviewDLPGuardrail @@ -11,8 +12,8 @@ if TYPE_CHECKING: def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm - tenant_id: Final = getattr(litellm_params, "tenant_id", None) - client_id: Final = getattr(litellm_params, "client_id", None) + tenant_id: Final = attribute_or_none(litellm_params, "tenant_id") + client_id: Final = attribute_or_none(litellm_params, "client_id") # client_secret can be passed via the standard api_key field or as # a dedicated client_secret parameter. diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 718401b6e68..43eb54d1464 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -10,9 +10,12 @@ import json import os from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeAlias, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional, TypeAlias, cast from urllib.parse import urlparse +from pydantic import TypeAdapter +from typing_extensions import TypedDict, Unpack + from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, @@ -25,7 +28,6 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage -from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.noma import NomaV2StreamingParams from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus @@ -39,6 +41,7 @@ _DEFAULT_API_BASE: Final = "https://api.noma.security/" _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail" _INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls") _DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname +_DECODED_JSON: Final = TypeAdapter(object) _GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object] @@ -50,6 +53,10 @@ _LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + ( ) +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + pass + + class _Action(str, enum.Enum): BLOCKED = "BLOCKED" NONE = "NONE" @@ -74,7 +81,7 @@ class NomaV2Guardrail(CustomGuardrail): streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, streaming_buffer_release_on_scan: bool | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -115,10 +122,11 @@ class NomaV2Guardrail(CustomGuardrail): if streaming.streaming_buffer_release_on_scan is not None: self.streaming_buffer_release_on_scan: bool = streaming.streaming_buffer_release_on_scan - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - - base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs - super().__init__(**base_kwargs) + options: Final[_CustomGuardrailOptions] = { + "supported_event_hooks": list(self.get_supported_event_hooks()), + **kwargs, + } + super().__init__(**options) @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: @@ -219,7 +227,7 @@ class NomaV2Guardrail(CustomGuardrail): except (ValueError, TypeError): json_str = safe_dumps(payload) - safe_payload: Final[object] = safe_json_loads(json_str, default={}) + safe_payload: Final = _DECODED_JSON.validate_python(safe_json_loads(json_str, default={})) if safe_payload == {} and payload: verbose_proxy_logger.warning( "Noma v2 guardrail: payload serialization failed, falling back to empty payload" diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 1a6feb47215..7c26c4c4c56 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional import httpx from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -27,6 +28,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + class OnyxGuardrail(CustomGuardrail): @classmethod @@ -115,7 +118,7 @@ class OnyxGuardrail(CustomGuardrail): try: response: Final = ModelResponse(**request_data) parsed: Final = response.json() - payload = parsed.get("response", {}) + payload = _DECODED_JSON.validate_python(parsed.get("response", {})) except Exception as e: verbose_proxy_logger.error( "Error in converting request_data to ModelResponse: %s", diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py index fe526b72ac3..4d120a2dd26 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING, Final import litellm +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) @@ -40,12 +41,12 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return openai_moderation_guardrail -def _get_config_value(litellm_params, optional_params, attribute_name): +def _get_config_value(litellm_params: object, optional_params: object, attribute_name: str) -> object: if optional_params is not None: - value: Final = getattr(optional_params, attribute_name, None) + value: Final = attribute_or_none(optional_params, attribute_name) if value is not None: return value - return getattr(litellm_params, attribute_name, None) + return attribute_or_none(litellm_params, attribute_name) guardrail_initializer_registry: Final = { diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 7af5eaf1da9..98be35079c3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -40,6 +41,8 @@ if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIModerationResponse from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_DECODED_JSON: Final = TypeAdapter(object) + def _config_secret(value: object) -> object: return get_secret_str(value) if isinstance(value, str) and value.startswith("os.environ/") else value @@ -151,7 +154,9 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): timeout=self.timeout, ) - verbose_proxy_logger.debug("OpenAI Moderation guard response: %s", response.json()) + verbose_proxy_logger.debug( + "OpenAI Moderation guard response: %s", _DECODED_JSON.validate_python(response.json()) + ) if response.status_code != 200: raise HTTPException( diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 8e03f7db4f5..6426bfb8788 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -18,6 +18,7 @@ from urllib.parse import urlparse import httpx from fastapi import HTTPException from pydantic import ConfigDict, TypeAdapter, ValidationError, field_validator +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -129,6 +130,21 @@ _RESPONSES_INPUT: Final[TypeAdapter[_ResponsesInput]] = TypeAdapter(_ResponsesIn _DECODED_ARGUMENTS: Final = TypeAdapter(object) +class _AirsToolEvent(TypedDict): + metadata: ReadOnly[Mapping[str, object]] + input: NotRequired[ReadOnly[str]] + + +def _mcp_tool_input(mcp_arguments: object) -> str | None: + if mcp_arguments is not None and mcp_arguments != "": + serialized_args: Final = ( + json.dumps(mcp_arguments) if isinstance(mcp_arguments, (dict, list)) else str(mcp_arguments) + ) + if serialized_args.strip(): + return serialized_args + return None + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -339,7 +355,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): is_response: bool = False, metadata: Mapping[str, object] | None = None, call_id: object = None, - tool_event: dict[str, Any] | None = None, + tool_event: _AirsToolEvent | None = None, ) -> dict[str, object]: """Call PANW Prisma AIRS API to scan content or a tool_event.""" @@ -894,7 +910,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): if not line.startswith("data: "): continue try: - data = json.loads(line[6:]) + data = _DECODED_ARGUMENTS.validate_python(json.loads(line[6:])) except (json.JSONDecodeError, ValueError): continue if not isinstance(data, dict): @@ -1954,24 +1970,21 @@ class PanwPrismaAirsHandler(CustomGuardrail): # "mcp_tool_name"/"mcp_arguments". Check canonical first, then fallback. mcp_tool_name: Final = request_data.get("mcp_tool_name") or self._mcp_name_fallback(request_data) if mcp_tool_name and input_type == "request": - mcp_tool_event: Final[dict[str, object]] = { - "metadata": { - "ecosystem": "mcp", - "method": "tools/call", - "server_name": self._get_mcp_server_name(request_data, mcp_tool_name), - "tool_invoked": mcp_tool_name, - }, + mcp_tool_metadata: Final = { + "ecosystem": "mcp", + "method": "tools/call", + "server_name": self._get_mcp_server_name(request_data, mcp_tool_name), + "tool_invoked": mcp_tool_name, } mcp_arguments = request_data.get("mcp_arguments") if mcp_arguments is None: mcp_arguments = request_data.get("arguments") - if mcp_arguments is not None and mcp_arguments != "": - if isinstance(mcp_arguments, (dict, list)): - serialized_args = json.dumps(mcp_arguments) - else: - serialized_args = str(mcp_arguments) - if serialized_args.strip(): - mcp_tool_event["input"] = serialized_args + mcp_tool_input: Final = _mcp_tool_input(mcp_arguments) + mcp_tool_event: Final[_AirsToolEvent] = ( + {"metadata": mcp_tool_metadata} + if mcp_tool_input is None + else {"metadata": mcp_tool_metadata, "input": mcp_tool_input} + ) mcp_scan_result: Final = await self._call_panw_api( tool_event=mcp_tool_event, diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index cadc9ca689b..7e2f6bfec33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -85,6 +85,7 @@ class _PresidioAnonymizeResponse(TypedDict): _PRESIDIO_ANONYMIZE_ADAPTER: Final[TypeAdapter[_PresidioAnonymizeResponse | None]] = TypeAdapter( _PresidioAnonymizeResponse | None ) +_DECODED_JSON: Final = TypeAdapter(object) class _JsonResponse(Protocol): @@ -1595,10 +1596,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if line.startswith("data: ") and line != "data: [DONE]": raw_json = line[6:] try: - event = json.loads(raw_json) + event = _DECODED_JSON.validate_python(json.loads(raw_json)) delta = event.get("delta") if isinstance(event, dict) else None if ( - isinstance(delta, dict) + isinstance(event, dict) + and isinstance(delta, dict) and event.get("type") == "content_block_delta" and delta.get("type") == "text_delta" and isinstance(delta.get("text"), str) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index a92e67046bd..6150fccac0c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -10,6 +10,7 @@ import os from typing import Any, Final, Literal from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -29,6 +30,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs GUARDRAIL_NAME: Final = "qualifire" DEFAULT_QUALIFIRE_API_BASE: Final = "https://proxy.qualifire.ai" +_DECODED_JSON: Final = TypeAdapter(object) class QualifireGuardrail(CustomGuardrail): @@ -188,7 +190,7 @@ class QualifireGuardrail(CustomGuardrail): args = function_info.get("arguments", {}) if isinstance(args, str): try: - args = json.loads(args) + args = _DECODED_JSON.validate_python(json.loads(args)) except json.JSONDecodeError: args = {} api_tool_calls.append( diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index 1a571047977..342e5b64714 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import yaml +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD @@ -22,6 +23,7 @@ if TYPE_CHECKING: ROUTE_TEMPLATES_DIR: Final = os.path.join(os.path.dirname(__file__), "route_templates") +_DECODED_YAML: Final = TypeAdapter(object) class SemanticGuardRouteLoader: @@ -55,7 +57,7 @@ class SemanticGuardRouteLoader: if not os.path.exists(file_path): raise ValueError(f"SemanticGuard: custom routes file not found: {file_path}") with open(file_path, "r") as f: - data: Final = yaml.safe_load(f) + data: Final = _DECODED_YAML.validate_python(yaml.safe_load(f)) if isinstance(data, list): return data if isinstance(data, dict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index 2fbd50b5863..37acee28b53 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -6,13 +6,14 @@ via embedding similarity. Smarter than regex (understands intent), lighter than an LLM call (~20-50ms per request for embedding). """ -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import attribute_items, attribute_or_none, require_attributes from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.route_loader import ( SemanticGuardRouteLoader, ) @@ -188,17 +189,21 @@ def _extract_user_text(messages: list) -> str: return "" -def _extract_response_text(response: Any) -> str: +class _MessageContent(Protocol): + @property + def content(self) -> object: ... + + +def _extract_response_text(response: object) -> str: """Extract text from every LLM response choice.""" - if hasattr(response, "choices") and response.choices: - text_parts: Final[list[str]] = [] - for choice in response.choices: - if hasattr(choice, "message") and choice.message: - text = _content_to_text(choice.message.content) - if text: - text_parts.append(text) - return "\n".join(text_parts) - return "" + text_parts: Final[list[str]] = [] + for choice in attribute_items(response, "choices"): + message = attribute_or_none(choice, "message") + if message: + text = _content_to_text(require_attributes(message, _MessageContent, "content").content) + if text: + text_parts.append(text) + return "\n".join(text_parts) def _content_to_text(content: object) -> str: diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index de0f8856522..fbf1f2576ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -105,7 +105,7 @@ def _resolve_tool_names_from_messages(messages: list[dict]) -> dict[str, str]: fn = (tc.get("function") or {}).get("name") else: tc_id = getattr(tc, "id", None) - fn_obj = getattr(tc, "function", None) + fn_obj: object = getattr(tc, "function", None) fn = getattr(fn_obj, "name", None) if fn_obj else None if tc_id and fn: mapping[tc_id] = fn diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 3ad9d1a2c42..d713f1078a2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime import anyio from fastapi import HTTPException +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -50,6 +51,8 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +_DECODED_JSON: Final = TypeAdapter(object) + _RequestData: TypeAlias = dict[str, object] @@ -183,7 +186,7 @@ def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) return item.get("id") if isinstance(item, str): try: - obj = json.loads(item.strip()) + obj = _DECODED_JSON.validate_python(json.loads(item.strip())) if isinstance(obj, dict) and "id" in obj: return obj.get("id") except (json.JSONDecodeError, TypeError): @@ -729,7 +732,7 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict: UserAPIKeyAuth, call_type: str, sampling_rate: int, - end_of_stream_only: bool, + end_of_stream_only: object, mappings: Mapping[CallTypes, type["BaseTranslation"]], ) -> AsyncGenerator[object, None]: """Emit guardrail text transformations as new deltas on the stream. @@ -1135,16 +1138,19 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply = request_data.pop("guardrail_to_apply", None) typed_request_data: Final[_RequestData] = request_data - def _streaming_flag(name: str, default: object) -> Any: + def _streaming_setting(name: str, default: object) -> object: return self.resolve_streaming_flag(guardrail_to_apply, name, default) + def _streaming_flag(name: str, default: object) -> Any: + return _streaming_setting(name, default) + sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). - end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False) + end_of_stream_only: object = _streaming_setting("streaming_end_of_stream_only", False) # "block_only" (default) drops guardrail text rewrites on the streaming # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). - streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only") + streaming_transform_mode: Final = _streaming_setting("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -1152,10 +1158,10 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated: bool = _streaming_flag( + buffer_until_moderated: object = _streaming_setting( "streaming_buffer_until_moderated", buffer_until_moderated_default ) - release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False) + release_on_scan: Final = _streaming_setting("streaming_buffer_release_on_scan", False) if ( buffer_until_moderated diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 7521c1e9f45..24b30e65e2b 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -102,6 +102,16 @@ def _as_json_array(value: object) -> list[object] | None: return _JSON_ARRAY.validate_python(value) if isinstance(value, list) else None +def _extra_param_items(litellm_params: LitellmParams) -> Iterable[tuple[str, object]]: + return ( + litellm_params.model_dump(exclude_none=True).items() + if hasattr(litellm_params, "model_dump") + else iter(litellm_params) + if litellm_params + else () + ) + + def contains_encrypted_marker(value: object, depth: int = 0) -> bool: """True if any string in value, at any JSON depth, starts with the encrypted-value prefix.""" if depth > DEFAULT_MAX_RECURSE_DEPTH: @@ -904,13 +914,7 @@ class InMemoryGuardrailHandler: # This matches the behavior of other guardrail initializers (e.g., initialize_lakera) # and aligns with the documented behavior for custom guardrails excluded_extra_param_keys: Final = frozenset(("guardrail", "mode", "default_on", "stream_scope")) - extra_params_items: Final = ( - litellm_params.model_dump(exclude_none=True).items() - if hasattr(litellm_params, "model_dump") - else iter(litellm_params) - if litellm_params - else () - ) + extra_params_items: Final = _extra_param_items(litellm_params) extra_params: Final = MappingProxyType( {key: value for key, value in extra_params_items if key not in excluded_extra_param_keys} ) diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py index 526f58dcd36..c49835905fb 100644 --- a/litellm/proxy/guardrails/tool_name_extraction.py +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -6,8 +6,8 @@ knowledge lives in one place. Uses guardrail translation handlers where availabl with standalone extractors for generate_content and MCP. """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Callable, Mapping +from typing import Final from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings @@ -15,7 +15,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import list_field from litellm.types.utils import CallTypes # Call types that have no guardrail translation handler; we use standalone extractors -STANDALONE_EXTRACTORS: Final[dict[str, Any]] = {} +STANDALONE_EXTRACTORS: Final[dict[str, Callable[[Mapping[str, object]], list[str]]]] = {} def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str]: @@ -28,7 +28,7 @@ def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str return names -def _extract_mcp_tool_names(data: dict) -> list[str]: +def _extract_mcp_tool_names(data: Mapping[str, object]) -> list[str]: """MCP call_tool: name or mcp_tool_name in body""" names: Final[list[str]] = [] name: Final = data.get("name") or data.get("mcp_tool_name") diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 2f466f22ef8..caa66bddf39 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -11,10 +11,11 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload from fastapi import APIRouter, Depends, Query -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.usage_tracking import guardrail_status_to_action @@ -838,28 +839,41 @@ def _snippet(text: object, max_len: int = 200) -> str | None: return result -def _input_snippet_for_log(sl: "prisma_models.LiteLLM_SpendLogs") -> str | None: - """Snippet for request input: prefer messages, fall back to proxy_server_request (same as drawer).""" - out = _snippet(sl.messages) +_DECODED_JSON: Final = TypeAdapter(object) + + +def _request_dict(value: object) -> Mapping[str, object] | None: + mapping: Final = as_str_mapping(value) + return mapping if isinstance(mapping, dict) else None + + +def _proxy_request_snippet(psr: object) -> str | None: + request: Final = _request_dict(psr) + if request is None: + return _snippet(psr) + msgs: Final = request.get("messages") + body: Final = _request_dict(request.get("body")) if msgs is None else None + out: Final = _snippet(body.get("messages") if body is not None else msgs) if out: return out - psr = getattr(sl, "proxy_server_request", None) + return _snippet(psr) + + +def _input_snippet_for_log(sl: "prisma_models.LiteLLM_SpendLogs") -> str | None: + """Snippet for request input: prefer messages, fall back to proxy_server_request (same as drawer).""" + out: Final = _snippet(sl.messages) + if out: + return out + psr: Final[object] = getattr(sl, "proxy_server_request", None) if not psr: return None - if isinstance(psr, str): - try: - psr = json.loads(psr) - except Exception: - return _snippet(psr) - if isinstance(psr, dict): - msgs = psr.get("messages") - if msgs is None and isinstance(psr.get("body"), dict): - msgs = psr["body"].get("messages") - out = _snippet(msgs) - if out: - return out + if not isinstance(psr, str): + return _proxy_request_snippet(psr) + try: + decoded: Final = _DECODED_JSON.validate_python(json.loads(psr)) + except Exception: return _snippet(psr) - return _snippet(psr) + return _proxy_request_snippet(decoded) @router.get( diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 47c3869fb64..8f9aebecf30 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -455,10 +455,10 @@ class KeyManagementEventHooks: try: from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.proxy_server import prisma_client except ImportError: return None - prisma_client: Final = getattr(proxy_server_module, "prisma_client", None) user_api_key_cache: Final = getattr(proxy_server_module, "user_api_key_cache", None) if prisma_client is None or user_api_key_cache is None: diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 47ea50debc9..05c64317a44 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -892,7 +892,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): {**_hidden_params, "additional_headers": _additional_headers}, ) - return await super().async_post_call_success_hook(data, user_api_key_dict, response) + await super().async_post_call_success_hook(data, user_api_key_dict, response) PROXY_MaxParallelRequestsHandler: Final = _PROXY_MaxParallelRequestsHandler diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a2bb691b401..567b8452ae2 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -101,7 +101,7 @@ if TYPE_CHECKING: from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation - Span = _Span | Any + Span = _Span InternalUsageCache = _InternalUsageCache else: Span = Any diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index bbe38068638..c4949834896 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -7,6 +7,7 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.types.router import ModelGroupInfo from litellm.types.utils import PriorityReservationDict @@ -85,10 +86,10 @@ def _resolve_provider_from_router_alias( if llm_router is None: return None try: - model_list: Final = getattr(llm_router, "model_list", None) + model_list: Final[object] = getattr(llm_router, "model_list", None) if not model_list: return None - for deployment in model_list: + for deployment in iter_items(model_list): if not isinstance(deployment, dict): continue if deployment.get("model_name") != model: diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 5f31aa788e2..a87455e5a60 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -9,6 +9,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import iter_items from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.proxy._types import ( CommonProxyErrors, @@ -166,6 +167,16 @@ class _AccessGroupTx(Protocol): def litellm_agentstable(self) -> _AgentTable: ... +class _AccessGroupTransaction(Protocol): + async def __aenter__(self) -> _AccessGroupTx: ... + + async def __aexit__(self, *exc_info: object) -> None: ... + + +def _access_group_transaction(prisma_client: PrismaClient) -> _AccessGroupTransaction: + return prisma_client.db.tx() + + def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( @@ -191,6 +202,10 @@ def _tool_permissions_json(value: Mapping[str, Sequence[str]] | None) -> "Json": return Json(None) if value is None else Json({server: list(tools) for server, tools in value.items()}) +def _explicitly_set_fields(data: AccessGroupUpdateRequest) -> Mapping[str, object]: + return data.model_dump(exclude_unset=True) + + @dataclass(frozen=True, slots=True) class _ResourceNames: mcp_servers: Mapping[str, str] @@ -539,7 +554,7 @@ async def create_access_group( try: tx: _AccessGroupTx - async with prisma_client.db.tx() as tx: + async with _access_group_transaction(prisma_client) as tx: existing: Final = await tx.litellm_accessgrouptable.find_unique( where={"access_group_name": data.access_group_name} ) @@ -647,7 +662,7 @@ async def update_access_group( _require_proxy_admin(user_api_key_dict) prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - update_fields: Final = data.model_dump(exclude_unset=True) + update_fields: Final = _explicitly_set_fields(data) update_data: Final[dict] = {"updated_by": user_api_key_dict.user_id} for field, value in update_fields.items(): if field in ( @@ -658,7 +673,7 @@ async def update_access_group( "access_agent_ids", "access_skill_names", ): - value = list(value or ()) + value = list(iter_items(value or ())) update_data[field] = ( _tool_permissions_json(data.mcp_tool_permissions) if field == "mcp_tool_permissions" else value ) @@ -672,7 +687,7 @@ async def update_access_group( try: tx: _AccessGroupTx - async with prisma_client.db.tx() as tx: + async with _access_group_transaction(prisma_client) as tx: # Read inside the transaction so delta computation is consistent with the write, # avoiding a TOCTOU race where a concurrent update could make deltas stale. existing: Final = await tx.litellm_accessgrouptable.find_unique(where={"access_group_id": access_group_id}) @@ -687,10 +702,10 @@ async def update_access_group( old_team_ids: Final[set[str]] = set(attached[access_group_id]) old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or []) new_team_ids: Final[set[str]] = ( - set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids + set(data.assigned_team_ids or []) if "assigned_team_ids" in update_fields else old_team_ids ) new_key_ids: Final[set[str]] = ( - set(update_fields["assigned_key_ids"] or []) if "assigned_key_ids" in update_fields else old_key_ids + set(data.assigned_key_ids or []) if "assigned_key_ids" in update_fields else old_key_ids ) teams_to_add = list(new_team_ids - old_team_ids) @@ -747,7 +762,7 @@ async def delete_access_group( affected_key_tokens: list[str] = [] tx: _AccessGroupTx - async with prisma_client.db.tx() as tx: + async with _access_group_transaction(prisma_client) as tx: existing: Final = await tx.litellm_accessgrouptable.find_unique(where={"access_group_id": access_group_id}) if existing is None: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index d40f64f435f..9c46e711cf6 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy.management_helpers.auto_router_permissions import ( from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.daily_activity_sql import build_where_clause +from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository @@ -179,7 +180,7 @@ class _ShadowEvalAttemptRow(Protocol): class _ShadowEvalFunnelTable(Protocol): - async def create_many(self, data: Sequence[Mapping[str, object]], skip_duplicates: bool) -> int: ... + async def create_many(self, data: Sequence[Mapping[str, object]], *, skip_duplicates: bool) -> int: ... class _ShadowEvalAttemptTable(Protocol): @@ -204,16 +205,28 @@ def _user_rows(prisma_client: "PrismaClient") -> _UserRowsTable: return UserRepository(prisma_client).table +class _ShadowEvalJobRepository(PrismaTableRepository[_ShadowEvalJobRow]): + table_name = "litellm_shadowevaljob" + + +class _ShadowEvalFunnelRepository(PrismaTableRepository[object]): + table_name = "litellm_shadowevalfunnel" + + +class _ShadowEvalAttemptRepository(PrismaTableRepository[_ShadowEvalAttemptRow]): + table_name = "litellm_shadowevalattempt" + + def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable: - return prisma_client.db.litellm_shadowevaljob + return _ShadowEvalJobRepository(prisma_client).table def _shadow_eval_funnel(prisma_client: "PrismaClient") -> _ShadowEvalFunnelTable: - return prisma_client.db.litellm_shadowevalfunnel # pyright: ignore[reportAttributeAccessIssue] # generated client + return _ShadowEvalFunnelRepository(prisma_client).table def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTable: - return prisma_client.db.litellm_shadowevalattempt + return _ShadowEvalAttemptRepository(prisma_client).table async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: @@ -221,6 +234,10 @@ async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) - return await prisma_client.db.query_raw(query, *args) +async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: + return await prisma_client.db.execute_raw(query, *args) + + async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> LiteLLM_TeamTable | None: from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, @@ -1788,7 +1805,7 @@ async def start_shadow_eval( # deliberate. Sweep and claim filter on exact (target_type, id) pairs so a team id # that happens to equal a key hash never matches the other kind's slot. for target_type, ids in requested_by_type: - await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, list(ids), target_type) + await _execute_raw(prisma_client, _SWEEP_FINISHED_JOBS_SQL, list(ids), target_type) claimed: Final = await _shadow_eval_jobs(prisma_client).find_many( where={ "OR": [ @@ -2005,8 +2022,8 @@ async def stop_shadow_eval_job( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) stamp: Final = datetime.now(timezone.utc) operator: Final = user_api_key_dict.user_id or "operator" - claimed: Final = await prisma_client.db.execute_raw( - _STOP_JOB_SQL, job_id, operator, stamp.replace(tzinfo=None).isoformat() + claimed: Final = await _execute_raw( + prisma_client, _STOP_JOB_SQL, job_id, operator, stamp.replace(tzinfo=None).isoformat() ) legs: Final = _LEG_ROWS.validate_python( await _shadow_eval_jobs(prisma_client).find_many(where={"group_id": job_id}) or () diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index b1f0002e143..14dd9ce6e8a 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -42,6 +42,7 @@ from litellm.proxy.db.db_span import db_span if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import PrismaClient + from litellm.repositories.prisma_protocols import TableActions from litellm.proxy.common_utils.encrypt_decrypt_utils import ( # noqa: F401 # legacy module exports _ALGO_AES_GCM, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _ENCRYPTION_ALGORITHM_SETTING, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -523,7 +524,7 @@ async def _scan_one_table( scalar_columns: tuple, ) -> LocationReport: report: Final = LocationReport(location=location) - table: Final = getattr(prisma_client.db, db_attr, None) + table: Final[TableActions[object] | None] = getattr(prisma_client.db, db_attr, None) if table is None: return report try: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index eebc537509b..0ad4258fb04 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -211,6 +211,7 @@ if TYPE_CHECKING: from prisma import models as prisma_models _RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) +_RawRowT = TypeVar("_RawRowT") class _UserRowLike(Protocol): @@ -234,6 +235,12 @@ class _TxTables(Protocol): litellm_proxymodeltable: TableActions[object] +class _ModelReencryptTransaction(Protocol): + async def __aenter__(self) -> _TxTables: ... + + async def __aexit__(self, *exc_info: object) -> None: ... + + class _ModelParamsUpdate(TypedDict): litellm_params: ReadOnly["prisma.Json"] @@ -334,6 +341,14 @@ def _credentials_table(prisma_client: PrismaClient) -> TableActions[CredentialIt ) +async def _query_raw_rows(prisma_client: PrismaClient, query: str, *args: object) -> Sequence[_RawRowT]: + return await prisma_client.db.query_raw(query, *args) + + +def _model_reencrypt_transaction(prisma_client: PrismaClient) -> _ModelReencryptTransaction: + return prisma_client.db.tx(timeout=timedelta(minutes=2)) + + def _config_table(prisma_client: PrismaClient, *, use_writer: bool = False) -> _ConfigTableActions: return cast( # cast-ok: ConfigRepository.table hides the write actions this module needs on that same object "_ConfigTableActions", ConfigRepository(prisma_client, use_writer=use_writer).table @@ -1313,11 +1328,6 @@ async def validate_team_id_used_in_service_account_request( return True -_BUDGET_NUMERIC_KEYS = frozenset( - ["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit", "tpd_limit"] -) - - def _enforce_upperbound_key_params( data: GenerateKeyRequest | UpdateKeyRequest, fill_defaults: bool = True, @@ -1332,13 +1342,20 @@ def _enforce_upperbound_key_params( # (GHSA-2rv4-xv66-fpjg): float('nan') passes every `< 0` check because # nan < 0 is False, and spend >= nan is always False, permanently disabling # budget enforcement for any key that carries it. - for elem in data: - key, value = elem - if key in _BUDGET_NUMERIC_KEYS and value is not None: - if not math.isfinite(value): + budget_numeric_fields: Final = ( + ("max_budget", data.max_budget), + ("max_parallel_requests", data.max_parallel_requests), + ("tpm_limit", data.tpm_limit), + ("rpm_limit", data.rpm_limit), + ("tpd_limit", data.tpd_limit), + ("soft_budget", data.soft_budget), + ) + for field_name, field_value in budget_numeric_fields: + if field_value is not None: + if not math.isfinite(field_value): raise HTTPException( status_code=400, - detail={"error": f"{key} must be a finite number. Received: {value}"}, + detail={"error": f"{field_name} must be a finite number. Received: {field_value}"}, ) if litellm.upperbound_key_generate_params is None: @@ -5746,8 +5763,7 @@ async def rotate_master_key( ] ) verbose_proxy_logger.debug("Re-encrypting litellm_params on %s model rows", len(reencrypted_models)) - async with prisma_client.db.tx(timeout=timedelta(minutes=2)) as tx_ctx: - tx: Final[_TxTables] = tx_ctx + async with _model_reencrypt_transaction(prisma_client) as tx: for reencrypted_model in reencrypted_models: await tx.litellm_proxymodeltable.update_many( data=_ModelParamsUpdate(litellm_params=prisma.Json(reencrypted_model.litellm_params)), @@ -7344,7 +7360,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final[Sequence[Mapping[str, int]]] = await _query_raw_rows(prisma_client, count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -7357,7 +7373,9 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final[Sequence[Mapping[str, str]]] = await _query_raw_rows( + prisma_client, aliases_sql, *aliases_params + ) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index b93c05c555f..da84c2d0013 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,10 +1,12 @@ """`/management/v1/spend_logs` facets.""" +from collections.abc import Sequence from datetime import datetime, timezone from functools import partial -from typing import Annotated, Final, Literal +from typing import Annotated, Final, Literal, Protocol from fastapi import APIRouter, Depends, Query, Request +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth @@ -19,6 +21,7 @@ from litellm.proxy.list_api.common import ( reject_unknown_query_params, ) from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX +from litellm.proxy.utils import PrismaClient from litellm.types.proxy.management_endpoints.management_v1 import ( FacetListResponse, PageMeta, @@ -33,6 +36,19 @@ router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) SPEND_LOGS_FACET_SCAN_CAP: Final = 10000 +class _FacetRow(TypedDict, total=False): + end_user: ReadOnly[str | None] + user: ReadOnly[str | None] + + +class _FacetQuery(Protocol): + async def __call__(self, query: str, *args: object) -> Sequence[_FacetRow]: ... + + +def _facet_query(prisma_client: PrismaClient) -> _FacetQuery: + return prisma_client.db.query_raw + + def _as_utc(value: datetime) -> datetime: return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) @@ -125,8 +141,8 @@ async def _list_spend_log_facet( f" ORDER BY {column_sql} ASC" f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" ) - rows: Final = await prisma_client.db.query_raw(facet_sql, *params) - values: Final[list[str]] = [row[column] for row in rows if row.get(column)] + rows: Final = await _facet_query(prisma_client)(facet_sql, *params) + values: Final[list[str]] = [value for row in rows if (value := row.get(column))] has_more: Final = len(values) > page_size return FacetListResponse( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a014b9cf8a5..776072c69be 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -104,6 +104,14 @@ class _HasServerId(Protocol): server_id: str +class _TemporaryServerCacheWriter(Protocol): + async def async_set_cache(self, *, key: str, value: str, ttl: int) -> object: ... + + +class _TemporaryServerCacheReader(Protocol): + async def async_get_cache(self, *, key: str) -> object: ... + + def _session_cookie_claims(claims: Mapping[str, object]) -> Mapping[str, object]: return claims @@ -527,7 +535,7 @@ if MCP_AVAILABLE: """ if litellm.cache is None or not hasattr(litellm.cache, "cache"): return - cache_backend: Final = getattr(litellm.cache, "cache", None) + cache_backend: Final[_TemporaryServerCacheWriter | None] = getattr(litellm.cache, "cache", None) if cache_backend is None or not hasattr(cache_backend, "async_set_cache"): return @@ -566,7 +574,7 @@ if MCP_AVAILABLE: """ if litellm.cache is None or not hasattr(litellm.cache, "cache"): return None - cache_backend: Final = getattr(litellm.cache, "cache", None) + cache_backend: Final[_TemporaryServerCacheReader | None] = getattr(litellm.cache, "cache", None) if cache_backend is None or not hasattr(cache_backend, "async_get_cache"): return None diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index fc0868a5fc9..11f616a5e18 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -3217,12 +3217,13 @@ async def update_useful_links( ) -def _validated_labeled_tiers( - tier_labels: dict[ComplexityTier, str], # mutable-ok: Pydantic materializes JSON object fields as dicts -) -> tuple[tuple[ComplexityTier, str], ...]: +_DECODED_JSON: Final = TypeAdapter(object) + + +def _validated_labeled_tiers(tier_labels: object) -> tuple[tuple[ComplexityTier, str], ...]: """Validate tier labels once for both prompt-preview transports.""" try: - return ComplexityRouterConfig(tier_labels=tier_labels).labeled_tiers() + return ComplexityRouterConfig.model_validate({"tier_labels": tier_labels}).labeled_tiers() except (TypeError, ValidationError) as e: raise ProxyException( message=f"tier_labels must be a JSON object of tier name to display name: {e}", @@ -3237,7 +3238,7 @@ def _labeled_tiers_from_query(tier_labels: str | None) -> tuple[tuple[Complexity if not tier_labels: return None try: - parsed: Final = json.loads(tier_labels) + parsed: Final = _DECODED_JSON.validate_python(json.loads(tier_labels)) except JSONDecodeError as e: raise ProxyException( message=f"tier_labels must be a JSON object of tier name to display name: {e}", diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 64566221d4d..1f4ca442d20 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -71,10 +71,11 @@ def _get_routing_strategies_from_router_class() -> list[str]: # Get the routing_strategy parameter routing_strategy_param: Final = sig.parameters.get("routing_strategy") + routing_strategy_annotation: Final[object] = routing_strategy_param.annotation if routing_strategy_param else None - if routing_strategy_param and routing_strategy_param.annotation: + if routing_strategy_param and routing_strategy_annotation: # Extract Literal values using get_args - literal_values: Final = get_args(routing_strategy_param.annotation) + literal_values: Final = get_args(routing_strategy_annotation) if literal_values: return list(literal_values) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c7adae2ce8d..361778ab2ea 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -479,6 +479,14 @@ def _team_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models. return TeamRepository(prisma_client).table +class _ExecuteRaw(Protocol): + async def __call__(self, query: str, *args: object) -> int: ... + + +def _execute_raw(prisma_client: PrismaClient) -> _ExecuteRaw: + return prisma_client.db.execute_raw + + def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]": return cast( # cast-ok: generated actions type Json columns as str; TableActions widens inputs to Mapping "TableActions[prisma_models.LiteLLM_TeamTable]", tx.litellm_teamtable @@ -3092,7 +3100,7 @@ async def _process_team_members( # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models member_allowed_models = data.allowed_models - team_default_member_models: Final = getattr(complete_team_data, "default_team_member_models", None) + team_default_member_models: Final = complete_team_data.default_team_member_models if member_allowed_models is None and team_default_member_models: member_allowed_models = team_default_member_models @@ -6572,7 +6580,7 @@ async def append_team_models( # When the team currently has models=[] (unrestricted access), the # CASE expression inserts the 'all-proxy-models' sentinel first. models_to_add: Final = list(data.models) - await prisma_client.db.execute_raw( + await _execute_raw(prisma_client)( 'UPDATE "LiteLLM_TeamTable" ' "SET models = (" " SELECT ARRAY(SELECT DISTINCT unnest(" diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index fd5eb672811..02494de8c87 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -278,6 +278,12 @@ class _CustomSsoCall(Protocol): async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... +def _configured_custom_sso() -> _CustomSsoCall | None: + from litellm.proxy.proxy_server import user_custom_sso + + return user_custom_sso + + class _ServicePrincipalAssignment(Protocol): def get(self, key: str) -> str: ... @@ -1051,6 +1057,11 @@ class _UserinfoFallbackCapableSSO(Protocol): ) -> CustomOpenID | None: ... +class _GenericSSOClient(_UserinfoFallbackCapableSSO, Protocol): + @property + def refresh_token(self) -> str | None: ... + + async def _verify_generic_sso_or_fall_back_to_id_token( generic_sso: _UserinfoFallbackCapableSSO, request: Request, @@ -1578,7 +1589,7 @@ async def get_generic_sso_response( team_mappings: Final = await _setup_team_mappings() generic_include_token_claims: Final = os.getenv("GENERIC_INCLUDE_TOKEN_CLAIMS", "false").lower() == "true" - def response_convertor(response: Mapping[str, object], httpx_session: object): + def response_convertor(response: Mapping[str, object], httpx_session: object) -> CustomOpenID: nonlocal received_response # return for user debugging response_id_token: Final = response.get("id_token") response_access_token: Final = response.get("access_token") @@ -1613,7 +1624,7 @@ async def get_generic_sso_response( discovery_document=discovery, response_convertor=response_convertor, ) - generic_sso: Final = SSOProvider( + generic_sso: Final[_GenericSSOClient] = SSOProvider( client_id=generic_client_id, client_secret=generic_client_secret, redirect_uri=redirect_url, @@ -1712,9 +1723,7 @@ async def get_generic_sso_response( access_token_str: str | None = combined_response.get("access_token") else: result = await _verify_generic_sso_or_fall_back_to_id_token( - generic_sso=cast( # cast-ok: fastapi-sso builds the provider class dynamically via create_provider - _UserinfoFallbackCapableSSO, generic_sso - ), + generic_sso=generic_sso, request=request, token_exchange_params=token_exchange_params, additional_headers=additional_generic_sso_headers_dict, @@ -2350,9 +2359,7 @@ async def _build_cli_sso_user_defined_values( result: OpenID | dict, parsed_openid_result: ParsedOpenIDResult, ) -> SSOUserDefinedValues | None: - from litellm.proxy.proxy_server import user_custom_sso - - custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + custom_sso_handler: Final = _configured_custom_sso() user_id: Final = parsed_openid_result.get("user_id") if custom_sso_handler is not None: if inspect.iscoroutinefunction(custom_sso_handler): @@ -2438,10 +2445,8 @@ async def _complete_cli_sso_callback_session( ): from fastapi.responses import HTMLResponse - from litellm.proxy.proxy_server import user_custom_sso - user_id: Final = parsed_openid_result.get("user_id") - if user_custom_sso is None: + if _configured_custom_sso() is None: _require_sso_user_id(user_id) user_email: Final = parsed_openid_result.get("user_email") user_info: Final = await get_user_info_from_db( @@ -3807,7 +3812,6 @@ class SSOAuthenticationHandler: proxy_logging_obj, redis_usage_cache, user_api_key_cache, - user_custom_sso, ) from litellm.proxy.utils import get_prisma_client_or_throw from litellm.types.proxy.ui_sso import ReturnedUITokenObject @@ -3832,7 +3836,7 @@ class SSOAuthenticationHandler: # But if it is, we want their models preferences user_defined_values: SSOUserDefinedValues | None = None - custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + custom_sso_handler: Final = _configured_custom_sso() if custom_sso_handler is not None: if inspect.iscoroutinefunction(custom_sso_handler): user_defined_values = await custom_sso_handler(result) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 3468396c6f6..920970427ba 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -310,7 +310,7 @@ async def _resolve_mcp_server_identifiers_to_ids( identifiers=identifiers, prisma_client=prisma_client, ): - server_id = getattr(server, "server_id", None) + server_id = server.server_id if not server_id: continue for identifier in identifiers: diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 34c80bcdce3..4457df460d4 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -181,8 +181,8 @@ async def _assert_write_access( """ if _is_admin(user_api_key_dict): return - row_user_id: Final = getattr(row, "user_id", None) - row_team_id: Final = getattr(row, "team_id", None) + row_user_id: Final = row.user_id + row_team_id: Final = row.team_id # Personal ownership. if row_user_id and row_user_id == user_api_key_dict.user_id: diff --git a/litellm/proxy/moyai_endpoints.py b/litellm/proxy/moyai_endpoints.py index c4d8ad139f7..45f4502ec4e 100644 --- a/litellm/proxy/moyai_endpoints.py +++ b/litellm/proxy/moyai_endpoints.py @@ -20,9 +20,10 @@ from typing import TYPE_CHECKING, Annotated, Final from urllib.parse import urlencode, urlparse from fastapi import APIRouter, Depends, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._internal_context import with_service_target +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( @@ -114,6 +115,9 @@ def _sign_connect_code(master_key: str, moyai_url: str, user_id: str | None) -> return f"{_b64url(payload)}.{_b64url(signature)}" +_DECODED_CONNECT_PAYLOAD: Final = TypeAdapter(object) + + def _decode_connect_code(master_key: str, code: str) -> _MoyaiConnectCode: try: payload_b64, signature_b64 = code.split(".", 1) @@ -125,18 +129,23 @@ def _decode_connect_code(master_key: str, code: str) -> _MoyaiConnectCode: if not hmac.compare_digest(signature, expected): raise HTTPException(status_code=400, detail="Invalid Moyai connect code") try: - payload: Final = json.loads(payload_raw) + payload: Final = as_str_mapping(_DECODED_CONNECT_PAYLOAD.validate_python(json.loads(payload_raw))) except (ValueError, TypeError): raise HTTPException(status_code=400, detail="Invalid Moyai connect code") - if not isinstance(payload, dict) or not isinstance(payload.get("exp"), int) or payload["exp"] < int(time.time()): + if payload is None: raise HTTPException(status_code=400, detail="Invalid Moyai connect code") - if not isinstance(payload.get("moyai_origin"), str) or not isinstance(payload.get("nonce"), str): + exp: Final = payload.get("exp") + if not isinstance(exp, int) or exp < int(time.time()): + raise HTTPException(status_code=400, detail="Invalid Moyai connect code") + moyai_origin: Final = payload.get("moyai_origin") + nonce: Final = payload.get("nonce") + if not isinstance(moyai_origin, str) or not isinstance(nonce, str): raise HTTPException(status_code=400, detail="Invalid Moyai connect code") user_id: Final = payload.get("user_id") return _MoyaiConnectCode( - moyai_origin=payload["moyai_origin"], - nonce=payload["nonce"], - exp=payload["exp"], + moyai_origin=moyai_origin, + nonce=nonce, + exp=exp, user_id=user_id if isinstance(user_id, str) else None, ) diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 1aa24eed4cd..5d596aece32 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -4,12 +4,15 @@ from dataclasses import dataclass from itertools import chain from typing import BinaryIO, Final, NoReturn +from pydantic import TypeAdapter from typing_extensions import assert_never +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException from litellm.proxy.openai_files_endpoints.file_usage_caps import FileUsageLimit, describe_limit_source _MB: Final = 1024 * 1024 +_DECODED_LINE: Final = TypeAdapter(object) @dataclass(frozen=True, slots=True) @@ -95,10 +98,10 @@ def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: try: - parsed: Final = json.loads(raw_line) + parsed: Final = as_str_mapping(_DECODED_LINE.validate_python(json.loads(raw_line))) except (ValueError, RecursionError): return BatchFileInvalidJsonLine(line_number=line_number) - if not isinstance(parsed, dict): + if parsed is None: return BatchFileLineNotObject(line_number=line_number) missing: Final = next((key for key in line_shape.required_keys if key not in parsed), None) if missing is None: diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index afc1e1c96fa..1f507982e65 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -8,10 +8,11 @@ storage backends (e.g., Azure Blob Storage) and managing associated metadata. import base64 import time from collections.abc import Sequence -from typing import Any, Final, cast +from typing import Final from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid as uuid_module +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend @@ -19,6 +20,7 @@ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.llms.openai import OpenAIFileObject, OpenAIFilesPurpose +from litellm.types.passthrough_endpoints.managed_id_rewriter import ManagedFileIdWriter from litellm.types.utils import ExtractedFileData, SpecialEnums @@ -232,11 +234,10 @@ class StorageBackendFileService: proxy_logging_obj: Proxy logging object user_api_key_dict: User API key authentication data """ - managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if not managed_files_obj or not isinstance(managed_files_obj, BaseFileEndpoints): verbose_proxy_logger.warning("Managed files hook not available, skipping managed files storage") return - managed_files_obj = cast(Any, managed_files_obj) # Create model mappings using storage URL model_mappings: Final = {model_name: storage_url for model_name in target_model_names} @@ -260,7 +261,10 @@ class StorageBackendFileService: ) # Store in managed files - await managed_files_obj.store_unified_file_id( + managed_files_writer: Final = require_attributes( + managed_files_obj, ManagedFileIdWriter, "store_unified_file_id" + ) + await managed_files_writer.store_unified_file_id( file_id=base64_unified_file_id, file_object=file_object, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..9eef4f0ef9d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1305,7 +1305,7 @@ async def handle_bedrock_count_tokens( verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers - provider_headers: Final = getattr(getattr(e, "response", None), "headers", None) + provider_headers: Final = e.response.headers raise HTTPException( status_code=e.status_code, detail={"error": e.message}, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py index 7527c504c59..3fb4fbd555d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py @@ -27,11 +27,15 @@ COMPREHEND_MEDICAL_COST_PER_UNIT_USD: Final[Mapping[str, float]] = MappingProxyT COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS: Final = frozenset(COMPREHEND_MEDICAL_COST_PER_UNIT_USD) +def _operation_from_headers(headers: Mapping[str, str]) -> str: + target: Final = headers.get("x-amz-target", "") + return target.split(".")[-1] + + class ComprehendMedicalPassthroughLoggingHandler: @staticmethod def _operation_from_response(httpx_response: httpx.Response) -> str: - target: Final[str] = httpx_response.request.headers.get("x-amz-target", "") - return target.split(".")[-1] + return _operation_from_headers(httpx_response.request.headers) @staticmethod def get_cost_for_operation(operation: str, text: str) -> float: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py index e277655f1e1..2afc1d17bc2 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -36,6 +36,7 @@ from litellm.types.passthrough_endpoints.tinyfish import ( from litellm.types.utils import StandardPassThroughResponseObject _RUN_ADAPTER: Final = TypeAdapter(TinyfishRun) +_DECODED_JSON: Final = TypeAdapter(object) _EMPTY_KWARGS: Final[Mapping[str, object]] = MappingProxyType({}) @@ -323,7 +324,7 @@ class TinyFishPassthroughLoggingHandler: "TinyFish passthrough: GET /v1/runs/%s returned %s", safe_run_id, response.status_code ) return None - payload: Final[object] = response.json() # any-ok: httpx Response.json() -> Any + payload: Final = _DECODED_JSON.validate_python(response.json()) return _parse_run(payload) except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request verbose_proxy_logger.warning("[Non blocking logging error] TinyFish run fetch failed: %s", e) @@ -413,7 +414,7 @@ def _run_id_from_sse_chunks(all_chunks: Sequence[str]) -> str | None: if not line.startswith("data:"): continue try: - event_payload: object = json.loads(line[5:].strip()) # any-ok: json.loads -> Any + event_payload: object = _DECODED_JSON.validate_python(json.loads(line[5:].strip())) except json.JSONDecodeError: continue event = _parse_run(event_payload) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py index 79f6dee90f7..5b45316a44c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py @@ -492,6 +492,10 @@ def s3_media_url(media_uri: str, aws_region_name: str) -> str | None: return f"https://{bucket}.s3.{aws_region_name}.{dns_suffix}/{quote(key)}" +def _header_or_default(headers: Mapping[str, str], name: str, default: str) -> str: + return headers.get(name, default) + + def media_predates_job(headers: Mapping[str, str], job_created_at: float) -> bool: try: modified_at: Final = parsedate_to_datetime(headers["last-modified"]).timestamp() @@ -501,7 +505,7 @@ def media_predates_job(headers: Mapping[str, str], job_created_at: float) -> boo async def write_media_within_limit(response: httpx.Response, media_file: IO[bytes], max_bytes: int) -> bool: - if int(response.headers.get("content-length", "0")) > max_bytes: + if int(_header_or_default(response.headers, "content-length", "0")) > max_bytes: return False async for chunk in response.aiter_bytes(): _ = media_file.write(chunk) @@ -591,8 +595,7 @@ class TranscribePassthroughLoggingHandler: @staticmethod def _operation_from_response(httpx_response: httpx.Response) -> str: - headers: Final[Mapping[str, str]] = httpx_response.request.headers - target: Final[str] = headers.get("x-amz-target", "") + target: Final = _header_or_default(httpx_response.request.headers, "x-amz-target", "") return target.split(".")[-1] @staticmethod diff --git a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py index 59b8d40577a..e70d583a027 100644 --- a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py +++ b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py @@ -7,6 +7,7 @@ records the reported values without recomputing them. """ import math +from collections.abc import Mapping from dataclasses import dataclass from typing import Final @@ -89,7 +90,7 @@ def _parse_total_tokens(raw_value: str | None) -> int | None: return total_tokens -def parse_upstream_reported_usage(headers: httpx.Headers) -> UpstreamReportedUsage | None: +def parse_upstream_reported_usage(headers: Mapping[str, str]) -> UpstreamReportedUsage | None: """Read the reported totals off an upstream pass-through response. Returns ``None`` when neither header is present, which is the normal case diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index c751091ff9b..15e7b7e8482 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -5,10 +5,10 @@ Attachments define WHERE policies apply, separate from the policy definitions. This allows the same policy to be attached to multiple scopes. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict +from typing import TYPE_CHECKING, Final, TypedDict from litellm._logging import verbose_proxy_logger from litellm.repositories.table_repositories import PolicyAttachmentRepository @@ -102,7 +102,7 @@ class AttachmentRegistry: self._initialized = True verbose_proxy_logger.info("Loaded %s policy attachments", len(self._attachments)) - def _parse_attachment(self, attachment_data: dict[str, Any]) -> PolicyAttachment: + def _parse_attachment(self, attachment_data: Mapping[str, object]) -> PolicyAttachment: """ Parse an attachment from raw configuration data. @@ -112,15 +112,17 @@ class AttachmentRegistry: Returns: Parsed PolicyAttachment object """ - return PolicyAttachment( - policy=attachment_data.get("policy", ""), - scope=attachment_data.get("scope"), - teams=attachment_data.get("teams"), - keys=attachment_data.get("keys"), - models=attachment_data.get("models"), - tags=attachment_data.get("tags"), - priority=attachment_data.get("priority"), - default=attachment_data.get("default", False), + return PolicyAttachment.model_validate( + { + "policy": attachment_data.get("policy", ""), + "scope": attachment_data.get("scope"), + "teams": attachment_data.get("teams"), + "keys": attachment_data.get("keys"), + "models": attachment_data.get("models"), + "tags": attachment_data.get("tags"), + "priority": attachment_data.get("priority"), + "default": attachment_data.get("default", False), + } ) def get_attached_policies( diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 1b1b2302236..4ad29bd3478 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -52,6 +52,8 @@ try: except ImportError: HTTPException = None +_LEGACY_HOOK_REPLACEMENT: Final = TypeAdapter(object) + class UndeliverableStreamRewrite(Exception): def __init__(self, guardrail_name: str, reason: str) -> None: @@ -232,10 +234,12 @@ class _LegacyHookStreamAdapter(CustomGuardrail): logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response")) - replacement: Final = await self.inner.async_post_call_success_hook( - data=request_data, - user_api_key_dict=self.user_api_key_dict, - response=hooked, + replacement: Final = _LEGACY_HOOK_REPLACEMENT.validate_python( + await self.inner.async_post_call_success_hook( + data=request_data, + user_api_key_dict=self.user_api_key_dict, + response=hooked, + ) ) rewrite: Final = hooked if replacement is None else replacement if rewrite is None: diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index f4b38bea14e..68ac71c567b 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -623,7 +623,7 @@ async def test_pipeline( ``` """ try: - validated_pipeline: Final = GuardrailPipeline(**request.pipeline) + validated_pipeline: Final = GuardrailPipeline.model_validate(request.pipeline) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid pipeline: {e}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ade719ae0ed..44f43808dc8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -24,6 +24,7 @@ from collections.abc import ( Awaitable, Callable, Collection, + Container, Iterable, Mapping, MutableMapping, @@ -360,6 +361,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_parent_otel_span_from_kwargs, iter_items, normalize_drop_params, + require_attributes, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -474,6 +476,7 @@ from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.credential_hydration import decrypted_or_stored +from litellm.proxy.common_utils.custom_openapi_spec import JsonObject from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.discoverable_model_filter import ( @@ -1140,7 +1143,7 @@ def _redact_worker_config_for_logging(worker_config: str | dict[str, JsonValue] return None if isinstance(worker_config, dict): return _redact_secret_values_in_obj(worker_config) - parsed: Final = safe_json_loads(worker_config, default=None) + parsed: Final[JsonValue] = safe_json_loads(worker_config, default=None) if isinstance(parsed, dict): return safe_dumps(_redact_secret_values_in_obj(parsed)) return worker_config @@ -1466,11 +1469,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ) _module_path, _func_name = _hook_spec.rsplit(":", 1) _module = importlib.import_module(_module_path) - _hook_fn = getattr(_module, _func_name) + _hook_fn = _UNTYPED_VALUE.validate_python(getattr(_module, _func_name)) if inspect.iscoroutinefunction(_hook_fn): await _hook_fn() - else: + elif callable(_hook_fn): _hook_fn() + else: + raise TypeError(f"'{type(_hook_fn).__name__}' object is not callable") verbose_proxy_logger.info("Worker startup hook '%s' executed successfully", _hook_spec) except Exception as e: verbose_proxy_logger.error("Worker startup hook '%s' failed: %s", _hook_spec, e) @@ -1918,15 +1923,22 @@ def _strip_operation_id_method_suffix(operation_id: str) -> str: return operation_id +class _SupportsValues(Protocol): + def values(self) -> Iterable[JsonValue]: ... + + +def _openapi_path_items(openapi_schema: JsonObject) -> Iterable[JsonObject]: + paths: Final = require_attributes(openapi_schema.get("paths", {}), _SupportsValues, "values") + return (path_item for path_item in paths.values() if isinstance(path_item, dict)) + + def ensure_unique_openapi_operation_ids( - openapi_schema: dict[str, Any], + openapi_schema: JsonObject, reserved_operation_ids: set[str] | None = None, -) -> dict[str, Any]: +) -> JsonObject: operation_entries: Final = [] operation_id_counts: Final[dict[str, int]] = {} - for path_item in openapi_schema.get("paths", {}).values(): - if not isinstance(path_item, dict): - continue + for path_item in _openapi_path_items(openapi_schema): for method, operation in path_item.items(): if method not in _OPENAPI_HTTP_METHODS or not isinstance(operation, dict): continue @@ -2003,6 +2015,19 @@ from starlette.routing import BaseRoute from litellm.proxy.common_utils.openapi_route import OpenAPISchema, serve_openapi_off_the_event_loop +class _OpenAPIQueryParam(Protocol): + @property + def name(self) -> object: ... + + @property + def required(self) -> object: ... + + +def _query_param_name_and_required(param: object) -> tuple[object, object]: + fields: Final = require_attributes(param, _OpenAPIQueryParam, "name", "required") + return fields.name, fields.required + + def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_routes: list) -> dict: """ Add a synthetic GET stub for each WebSocket route so it appears in Swagger UI. @@ -2019,14 +2044,15 @@ def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_ try: if hasattr(route, "dependant") and route.dependant is not None: # Handle both FastAPI <0.120 and >=0.120 - query_params = getattr(route.dependant, "query_params", []) + query_params = _UNTYPED_VALUE.validate_python(getattr(route.dependant, "query_params", [])) if query_params: - for param in query_params: + for param in iter_items(query_params): + name, required = _query_param_name_and_required(param) parameters.append( { - "name": param.name, + "name": name, "in": "query", - "required": param.required, + "required": required, "schema": {"type": "string"}, } ) @@ -5576,6 +5602,13 @@ def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None: _DB_CONFIG_CALLBACK_EVENT_TYPES: Final[Mapping[str, tuple[Literal["success", "failure"], ...]]] = MappingProxyType( {"success_callback": ("success",), "failure_callback": ("failure",), "callbacks": ("success", "failure")} ) +_DB_CONFIG_CALLBACK_LISTS: Final[Mapping[str, Callable[[], Container[object]]]] = MappingProxyType( + { + "success_callback": lambda: litellm.success_callback, + "failure_callback": lambda: litellm.failure_callback, + "callbacks": lambda: litellm.callbacks, + } +) _CALLBACK_LIST_NAMES: Final = ( "callbacks", "success_callback", @@ -5586,7 +5619,9 @@ _CALLBACK_LIST_NAMES: Final = ( def _callback_list_entries(list_name: str) -> tuple[tuple[str, object], ...]: - return tuple((list_name, entry) for entry in getattr(litellm, list_name)) + return tuple( + (list_name, entry) for entry in iter_items(_UNTYPED_VALUE.validate_python(getattr(litellm, list_name))) + ) def _registered_callback_entries() -> tuple[tuple[str, object], ...]: @@ -5615,7 +5650,7 @@ def _restore_displaced_callback(list_name: str, entry: object) -> None: def _unregister_db_callback(registration: _DbCallbackRegistration) -> None: for list_name, entry in registration.added: litellm.logging_callback_manager.remove_callback_from_list_by_object( - getattr(litellm, list_name), entry, require_self=False + _UNTYPED_VALUE.validate_python(getattr(litellm, list_name)), entry, require_self=False ) for list_name, entry in registration.displaced: _restore_displaced_callback(list_name, entry) @@ -7824,7 +7859,7 @@ class ProxyConfig: self, callback: str, event_types: list[Literal["success", "failure"]], - existing_callbacks: list, + existing_callbacks: Container[object], ) -> None: """ Helper method to add a single callback to litellm for specified event types. @@ -7867,7 +7902,7 @@ class ProxyConfig: self._add_callback_from_db_to_in_memory_litellm_callbacks( callback=callback, event_types=list(_DB_CONFIG_CALLBACK_EVENT_TYPES[setting_key]), - existing_callbacks=getattr(litellm, setting_key), + existing_callbacks=_DB_CONFIG_CALLBACK_LISTS[setting_key](), ) after: Final = _registered_callback_entries() return _DbCallbackRegistration( @@ -8362,8 +8397,12 @@ class ProxyConfig: environment_read: Final = next(self._environment_read_numbers) sections: Final = tuple(self._settings_stores) responses: Final = await asyncio.gather(*(get_config_param(prisma_client, section) for section in sections)) + response: _ConfigParamValueRow | None for section, response in zip(sections, responses): - if response is None or (param_value := getattr(response, "param_value", None)) is None: + if response is None: + continue + param_value: JsonValue = getattr(response, "param_value", None) + if param_value is None: continue verbose_proxy_logger.debug( @@ -8390,7 +8429,10 @@ class ProxyConfig: self, prisma_client: PrismaClient, row: object, environment_before_db: Mapping[str, str], read: int ) -> None: current: Final = await self._environment_row_not_behind_the_applied_one(prisma_client, row) - if current is None or (param_value := getattr(current, "param_value", None)) is None: + if current is None: + return + param_value: Final[object] = getattr(current, "param_value", None) + if param_value is None: return self._apply_db_environment_variables(param_value, _config_row_revision(current), environment_before_db, read) @@ -15997,7 +16039,23 @@ async def _get_caller_byok_team_scope( return key_team_scope | set(user_row.teams or []) -def _byok_row_outside_caller_teams(model_info_dict: dict[str, JsonValue], allowed_team_ids: set[str] | None) -> bool: +class _SupportsGet(Protocol): + def get(self, key: str, default: object = None, /) -> object: ... + + +def _with_get(value: object) -> _SupportsGet: + return require_attributes(value, _SupportsGet, "get") + + +class _SupportsLower(Protocol): + def lower(self) -> object: ... + + +def _lowered(value: object) -> object: + return require_attributes(value, _SupportsLower, "lower").lower() + + +def _byok_row_outside_caller_teams(model_info_dict: object, allowed_team_ids: set[str] | None) -> bool: """Whether a team BYOK row belongs to a team the caller is not a member of. `team_id` is only set on team BYOK rows; non-team rows fall through @@ -16005,7 +16063,7 @@ def _byok_row_outside_caller_teams(model_info_dict: dict[str, JsonValue], allowe """ if allowed_team_ids is None: return False - team_id: Final = model_info_dict.get("team_id") + team_id: Final = _with_get(model_info_dict).get("team_id") if team_id is None: return False return team_id not in allowed_team_ids @@ -16106,7 +16164,7 @@ async def _fetch_db_models_for_search( async def _apply_search_filter_to_models( - all_models: list[dict[str, Any]], + all_models: list[dict[str, object]], search: str, prisma_client: PrismaClient | None, proxy_config: ProxyConfig, @@ -16158,7 +16216,7 @@ async def _apply_search_filter_to_models( prisma_client=prisma_client, ) - def _is_byok_outside_caller_teams(model_info_dict: dict[str, JsonValue]) -> bool: + def _is_byok_outside_caller_teams(model_info_dict: object) -> bool: return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) def _model_matches_search(m: dict[str, Any]) -> bool: @@ -16195,7 +16253,7 @@ async def _apply_search_filter_to_models( db_model_ids_in_router: Final = set() for m in filtered_router_models: - model_info = m.get("model_info", {}) + model_info = _with_get(m.get("model_info", {})) is_db_model = model_info.get("db_model", False) model_id = model_info.get("id") @@ -16306,7 +16364,7 @@ def _sort_models( reverse: Final = sort_order.lower() == "desc" def get_sort_key(model: dict[str, Any]) -> Any: - model_info: Final = model.get("model_info", {}) + model_info: Final = _with_get(_UNTYPED_VALUE.validate_python(model.get("model_info", {}))) if sort_by == "model_name": # Team BYOK models persist an internal `model_name` (e.g. @@ -16318,7 +16376,7 @@ def _sort_models( team_public_model_name: Final = model_info.get("team_public_model_name") if team_public_model_name: return str(team_public_model_name).lower() - return model.get("model_name", "").lower() + return _lowered(_UNTYPED_VALUE.validate_python(model.get("model_name", ""))) elif sort_by == "created_at": created_at: Final = model_info.get("created_at") @@ -16344,8 +16402,8 @@ def _sort_models( return normalized_dt elif sort_by == "costs": - input_cost: Final = model_info.get("input_cost_per_token", 0) or 0 - output_cost: Final = model_info.get("output_cost_per_token", 0) or 0 + input_cost: Final = model.get("model_info", {}).get("input_cost_per_token", 0) or 0 + output_cost: Final = model.get("model_info", {}).get("output_cost_per_token", 0) or 0 total_cost: Final = input_cost + output_cost # Put 0 or None costs at the end for asc, at the start for desc if total_cost == 0: @@ -16580,7 +16638,7 @@ async def _authorize_team_id_query( async def _filter_models_by_team_id( - all_models: list[dict[str, Any]], + all_models: list[dict[str, object]], team_id: str, prisma_client: PrismaClient, llm_router: Router, @@ -16624,7 +16682,7 @@ async def _filter_models_by_team_id( # every public model the admin can call. filtered_models: Final = [] for _model in all_models: - model_info = _model.get("model_info", {}) + model_info = _with_get(_model.get("model_info", {})) model_id = model_info.get("id", None) # BYOK rows owned by this team are always accessible to it, even if diff --git a/litellm/proxy/public_endpoints/public_v1/model_hub.py b/litellm/proxy/public_endpoints/public_v1/model_hub.py index cd634e490b6..e58149fba0b 100644 --- a/litellm/proxy/public_endpoints/public_v1/model_hub.py +++ b/litellm/proxy/public_endpoints/public_v1/model_hub.py @@ -107,10 +107,14 @@ def _features(row: ModelGroupInfoProxy) -> tuple[str, ...]: The hub's feature control has always been a multi-select over the `supports_*` flags. One boolean filter per flag would AND them, which is the opposite of what it does. """ + return _enabled_features(row.model_dump()) + + +def _enabled_features(fields: Mapping[str, object]) -> tuple[str, ...]: return tuple( sorted( name.removeprefix(FEATURE_PREFIX) - for name, value in row.model_dump().items() + for name, value in fields.items() if name.startswith(FEATURE_PREFIX) and value is True ) ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 803df9f83ca..0b91e57f01e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -728,7 +728,9 @@ async def _get_team_member_budget_counter( membership_cache_key: Final = team_membership_reservation_cache_key( user_id=valid_token.user_id, team_id=team_object.team_id ) - cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key) + cached_team_membership: Final = _CACHED_VALUE.validate_python( + await user_api_key_cache.async_get_cache(key=membership_cache_key) + ) team_membership: LiteLLM_TeamMembership | None = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): team_membership = cached_team_membership diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index ea543aeae6e..4abc79febe1 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -21,6 +21,8 @@ from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol +from pydantic import TypeAdapter + from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -135,6 +137,9 @@ def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: return str(getattr(row, "model_name", "") or "") +_DECODED_JSON: Final = TypeAdapter(object) + + def _decode_model_info(raw: object) -> "Mapping[str, object] | None": """A deployment's model_info as a mapping, decoding a JSON string, else None. @@ -143,7 +148,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None": """ if isinstance(raw, str): try: - decoded: Final[object] = json.loads(raw) + decoded: Final = _DECODED_JSON.validate_python(json.loads(raw)) except (TypeError, ValueError): return None return decoded if isinstance(decoded, dict) else None diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 84278f4619d..3c65fc0bf06 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3045,7 +3045,7 @@ async def ui_view_spend_logs( ) sql_params.extend([page_size, skip]) - data: Final = await prisma_client.db.query_raw(sql_query, *sql_params) + data: Final[Sequence[Mapping[str, object]]] = await _query_raw(prisma_client, sql_query, *sql_params) if request_id is not None and not is_v2 and not is_admin_view: await _assert_user_owns_fetched_spend_rows( @@ -3064,7 +3064,7 @@ async def ui_view_spend_logs( return await _build_ui_spend_logs_response( prisma_client, - data, + list(data), total_records, page, page_size, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index efc83132568..6cf274884c7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,6 +19,7 @@ from collections.abc import ( Awaitable, Callable, Coroutine, + Iterable, Iterator, Mapping, Sequence, @@ -51,7 +52,7 @@ from typing import ( import anyio import httpx -from typing_extensions import Never, ReadOnly, TypedDict +from typing_extensions import Never, NotRequired, ReadOnly, TypedDict from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -323,10 +324,22 @@ class _RelTuplesRow(TypedDict): reltuples: ReadOnly[int] +class _TeamMemberRow(TypedDict): + role: ReadOnly[Literal["admin", "user"]] + user_id: ReadOnly[NotRequired[str | None]] + user_email: ReadOnly[NotRequired[str | None]] + + +class _TeamMembershipRow(TypedDict): + user_id: ReadOnly[str | None] + team_members_with_roles: ReadOnly[Iterable[_TeamMemberRow] | None] + + _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) +_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -364,6 +377,47 @@ class _UpdateManyBatch(Protocol): async def commit(self) -> None: ... +def _update_many_batch(db: PrismaWrapper | RoutingPrismaWrapper) -> _UpdateManyBatch: + return db.batch_() + + +def _query_raw_rows( + db: PrismaWrapper | RoutingPrismaWrapper, query: str, *args: object +) -> Awaitable[Sequence[Mapping[str, object]]]: + return db.query_raw(query, *args) + + +def _query_first_row( + db: PrismaWrapper | RoutingPrismaWrapper, query: str, *args: object +) -> Awaitable[Mapping[str, object] | None]: + return db.query_first(query, *args) + + +def _team_member_for_user(row: _TeamMembershipRow) -> Member | None: + members_with_roles: Final = row["team_members_with_roles"] + if members_with_roles is None: + return None + user_id: Final = row["user_id"] + if user_id is None: + return None + matches: Final = [ + Member(**member) + for member in members_with_roles + if member.get("user_id") is not None and user_id == member.get("user_id") + ] + return matches[-1] if matches else None + + +class _EndUserSpendTransaction(Protocol): + def batch_(self) -> contextlib.AbstractAsyncContextManager[_EndUserSpendBatch]: ... + + +def _end_user_spend_transaction( + db: PrismaWrapper | RoutingPrismaWrapper, timeout: timedelta +) -> contextlib.AbstractAsyncContextManager[_EndUserSpendTransaction]: + return db.tx(timeout=timeout) + + unified_guardrail: Final = UnifiedLLMGuardrails() NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages}) @@ -543,7 +597,7 @@ class _UpstreamStreamBoundary(Generic[_T]): self.failure: BaseException | None = None def __getattr__(self, name: str) -> object: - return getattr(self._source, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self._source, name)) def __aiter__(self) -> "_UpstreamStreamBoundary[_T]": return self @@ -5494,7 +5548,9 @@ class PrismaClient: raise e - async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> dict | None: + async def _query_first_with_cached_plan_fallback( + self, sql_query: str, *args: object + ) -> Mapping[str, object] | None: """ Execute a query, recovering once from PostgreSQL's "cached plan must not change result type" error. @@ -5745,7 +5801,7 @@ class PrismaClient: LIMIT $1 OFFSET $2 """ - response = await self.db.query_raw(sql_query, limit, offset) + response = await _query_raw_rows(self.db, sql_query, limit, offset) return response elif table_name == "spend": verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'") @@ -5889,7 +5945,7 @@ class PrismaClient: """ response = ( - await self.writer_db.query_first(sql_query, hashed_token) + await _query_first_row(self.writer_db, sql_query, hashed_token) if use_writer else await self._query_first_with_cached_plan_fallback(sql_query, hashed_token) ) @@ -5925,31 +5981,14 @@ class PrismaClient: if response["team_blocked"] is None: response["team_blocked"] = False - team_member: Member | None = None - if response["team_members_with_roles"] is not None and response["user_id"] is not None: - ## find the team member corresponding to user id - """ - [ - { - "role": "admin", - "user_id": "default_user_id", - "user_email": null - }, - { - "role": "user", - "user_id": null, - "user_email": "test@email.com" - } - ] - """ - for tm in response["team_members_with_roles"]: - if tm.get("user_id") is not None and response["user_id"] == tm.get("user_id"): - team_member = Member(**tm) + ## find the team member corresponding to user id + team_member: Final = _team_member_for_user(response) response["team_member"] = team_member - response = LiteLLM_VerificationTokenView(**response, last_refreshed_at=time.time()) + token_view: Final = LiteLLM_VerificationTokenView(**response, last_refreshed_at=time.time()) # for prisma we need to cast the expires time to str - if response.expires is not None and isinstance(response.expires, datetime): - response.expires = response.expires.isoformat() + if token_view.expires is not None and isinstance(token_view.expires, datetime): + token_view.expires = token_view.expires.isoformat() + response = token_view return response except Exception as e: import traceback @@ -6239,7 +6278,7 @@ class PrismaClient: """ Batch write update queries """ - batcher: _UpdateManyBatch = self.db.batch_() + batcher: _UpdateManyBatch = _update_many_batch(self.db) for idx, t in enumerate(data_list): # check if plain text or hash if t.token.startswith("sk-"): @@ -6264,7 +6303,7 @@ class PrismaClient: """ Batch write update queries """ - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for user in sorted(data_list, key=attrgetter("user_id")): try: data_json = self.jsonify_object(data=user.model_dump(exclude_none=True)) @@ -6289,7 +6328,7 @@ class PrismaClient: """ Batch write update queries """ - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for enduser in data_list: try: data_json = self.jsonify_object(data=enduser.model_dump(exclude_none=True)) @@ -6314,7 +6353,7 @@ class PrismaClient: """ Batch write update queries """ - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for budget in data_list: try: data_json = self.jsonify_object(data=budget.model_dump(exclude_none=True)) @@ -6337,7 +6376,7 @@ class PrismaClient: and isinstance(data_list, list) ): # Batch write update queries - batcher = self.db.batch_() + batcher = _update_many_batch(self.db) for idx, team in enumerate(data_list): try: data_json = self.jsonify_team_object(db_data=team.model_dump(exclude_none=True)) @@ -7002,7 +7041,7 @@ class PrismaClient: # caller's own retried query, a stronger check than SELECT 1, # and a reader that fails to come back sets `_reader_unavailable` # so reads fall through to the writer just recreated here. - await self.writer_db.query_raw("SELECT 1") + await _query_raw_rows(self.writer_db, "SELECT 1") # A recreate can decline: the optimistic-lock guard no-ops when # the writer generation moved since cycle entry, and the routing # wrapper then leaves the reader untouched as well. Callers that @@ -7350,7 +7389,7 @@ class PrismaClient: async def _view_setup_gate_table_present(self) -> bool: rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( - await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) + await _query_raw_rows(self.db, "SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) ) return rows[0]["present"] @@ -7360,7 +7399,7 @@ class PrismaClient: await asyncio.sleep(self._db_health_watchdog_interval_seconds) async with db_span("db_health_watchdog", None): await asyncio.wait_for( - self.db.query_raw("SELECT 1"), + _query_raw_rows(self.db, "SELECT 1"), timeout=self._db_health_watchdog_probe_timeout_seconds, ) if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: @@ -7420,7 +7459,7 @@ class PrismaClient: async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool: """True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006).""" rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python( - await writer.query_raw(_WRITER_WRITABILITY_PROBE_SQL) + await _query_raw_rows(writer, _WRITER_WRITABILITY_PROBE_SQL) ) return any(row.get("transaction_read_only") == "on" for row in rows) @@ -7447,7 +7486,7 @@ class PrismaClient: """ sql_query: Final = "SELECT 1" async with db_span("health_check", None): - response: Final[object] = await wrapper.query_raw(sql_query) + response: Final[object] = await _query_raw_rows(wrapper, sql_query) return response async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: @@ -8000,7 +8039,7 @@ class ProxyUpdateSpend: try: async with ( db_span("update_end_user_spend", "LiteLLM_EndUserTable"), - prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction, + _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) as transaction, ): batcher: _EndUserSpendBatch async with transaction.batch_() as batcher: diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index d4643994043..6d9249d5572 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -4,6 +4,7 @@ from typing import Final from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from litellm.proxy._types import * @@ -36,6 +37,7 @@ from litellm.types.videos.utils import ( ) router: Final = APIRouter() +_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) @router.post( @@ -582,7 +584,9 @@ async def video_create_character( if video_file: data["video"] = video_file[0] - target_model_name: Final = extract_model_from_target_model_names(data.get("target_model_names")) + target_model_name: Final = extract_model_from_target_model_names( + _OPAQUE_VALUE_ADAPTER.validate_python(data.get("target_model_names")) + ) if target_model_name and not data.get("model"): data["model"] = target_model_name diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index a2be42f1b76..56123e4bee3 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: from litellm.router import Router -def extract_model_from_target_model_names(target_model_names: Any) -> str | None: +def extract_model_from_target_model_names(target_model_names: object) -> str | None: if isinstance(target_model_names, str): model_names: Final = [m.strip() for m in target_model_names.split(",") if m.strip()] return model_names[0] if model_names else None diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 1e7aa6ac1c6..5520912dd41 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -106,11 +106,11 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) self.embedding_model = self.vector_store_config.get("embedding_model") or "amazon.titan-embed-text-v2:0" - self.wait_for_ingestion = self.vector_store_config.get("wait_for_ingestion", False) + self.wait_for_ingestion: object = self.vector_store_config.get("wait_for_ingestion", False) self.ingestion_timeout: int = _get_int(self.vector_store_config.get("ingestion_timeout"), 300) # Get AWS region using BaseAWSLLM method - _aws_region: Final = self.vector_store_config.get("aws_region_name") + _aws_region: Final[object] = self.vector_store_config.get("aws_region_name") self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( aws_region_name=str(_aws_region) if _aws_region else None ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 088dfb54155..7b9068ff95e 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -33,7 +33,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager, resolve_api_key_secret_reference -from ..litellm_core_utils.core_helpers import normalize_drop_params +from ..litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -83,6 +83,13 @@ def _with_resolved_session_model(session: dict[str, object], model_name: str) -> return {**session, "model": model_name} +def _resolved_request_session(request_data: Mapping[str, object], model_name: str) -> Mapping[str, object] | None: + session: Final = as_str_mapping(request_data.get("session")) + if not isinstance(session, dict): + return None + return _with_resolved_session_model(session, model_name) + + def _build_litellm_metadata(kwargs: dict) -> dict[str, object]: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" metadata: Final[dict[str, object]] = {**(kwargs.get("litellm_metadata") or {})} @@ -189,8 +196,9 @@ async def acreate_realtime_client_secret( custom_llm_provider=custom_llm_provider, ) request_data: Final = req.model_dump(exclude_none=True, exclude={"model"}) - if isinstance(request_data.get("session"), dict): - request_data["session"] = _with_resolved_session_model(request_data["session"], model_name) + resolved_session: Final = _resolved_request_session(request_data, model_name) + if resolved_session is not None: + request_data["session"] = resolved_session return await base_llm_http_handler.async_realtime_client_secret_handler( api_base=resolved_api_base, api_key=await _realtime_http_bearer_token( diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 81fba770b70..92f8ba984f4 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -8,6 +8,7 @@ from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable from pydantic import BaseModel +from litellm.litellm_core_utils.core_helpers import attribute_or_none from litellm.repositories.chunked_in import find_many_in from litellm.repositories.prisma_protocols import TableActions @@ -131,4 +132,4 @@ def is_unique_violation(exc: BaseException) -> bool: return "P2002" in str(exc) or "unique constraint" in str(exc).lower() if isinstance(exc, UniqueViolationError): return True - return getattr(exc, "code", None) == "P2002" + return attribute_or_none(exc, "code") == "P2002" diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 57f6fd33c11..b89d87aa354 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -75,6 +75,7 @@ class _PrismaClientView(Protocol): _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) +_RAW_COLUMN: Final = TypeAdapter(object) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", "model_spend", @@ -142,7 +143,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) if not rows: return None - raw_value: Final[object] = rows[0]["members_with_roles"] + raw_value: Final = _RAW_COLUMN.validate_python(rows[0]["members_with_roles"]) parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index f7e2616324f..65d9b89cf79 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -1,6 +1,6 @@ import asyncio import contextvars -from collections.abc import Awaitable +from collections.abc import Awaitable, Coroutine from functools import partial from typing import Final, Literal @@ -76,7 +76,7 @@ async def arerank( func_with_context: Final = partial(ctx.run, func) init_response: Final = await loop.run_in_executor(None, func_with_context) - if asyncio.iscoroutine(init_response): + if isinstance(init_response, Coroutine): response = await init_response else: response = init_response diff --git a/litellm/rerank_api/rerank_utils.py b/litellm/rerank_api/rerank_utils.py index f365ae749d1..4924625b011 100644 --- a/litellm/rerank_api/rerank_utils.py +++ b/litellm/rerank_api/rerank_utils.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -8,7 +8,7 @@ def get_optional_rerank_params( model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7121908517b..f29f6c83ef0 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -15,7 +15,7 @@ import json import time import uuid from collections.abc import Iterable, Sequence -from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast # noqa: TID251 # see kwargs-ok / cast-ok markers +from typing import TYPE_CHECKING, Final, TypeAlias, cast # noqa: TID251 # see cast-ok markers from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -517,7 +517,7 @@ async def aresponses_with_emulated_file_search( model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call - **kwargs: Any, # kwargs-ok: `object` would surface the caller's partially-unknown dict at its call site + **kwargs: object, # kwargs-ok: pass-through params forwarded as-is to the underlying aresponses call ) -> ResponsesAPIResponse: """ Emulated file_search for providers that don't support it natively. diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index fcda6b31b00..98e1984550c 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -23,6 +23,7 @@ from typing import Final, Literal, TypeAlias from pydantic import TypeAdapter, ValidationError from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( ChatCompletionToolParam, @@ -30,6 +31,7 @@ from litellm.types.llms.openai import ( ) _MAX_ARGUMENTS_LEN: Final = 1_000_000 +_DECODED_JSON: Final = TypeAdapter(object) TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE: Final = MappingProxyType( {"function_call": "fc", "custom_tool_call": "ctc", "tool_search_call": "tsc"} @@ -111,8 +113,8 @@ def unwrap_custom_tool_arguments(arguments: str) -> str: if len(arguments) > _MAX_ARGUMENTS_LEN: return arguments try: - parsed: Final = json.loads(arguments) - if isinstance(parsed, dict) and "content" in parsed: + parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(arguments))) + if parsed is not None and "content" in parsed: return str(parsed["content"]) except (json.JSONDecodeError, TypeError, ValueError): pass diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index c25dd4a405e..b5ef9e95741 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -5,6 +5,7 @@ from itertools import filterfalse from typing import Any, Final, cast import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.main import stream_chunk_builder @@ -84,14 +85,18 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | ) +def _is_signed_thinking_block(block: object) -> bool: + fields: Final = as_str_mapping(block) + return fields is not None and isinstance(block, dict) and bool(fields.get("signature") or fields.get("data")) + + def _streams_compaction(chunk: ModelResponseStream) -> bool: delta_fields: Final = chunk.choices[0].delta.provider_specific_fields if chunk.choices else None return delta_fields is not None and "compaction_blocks" in delta_fields def _delta_has_signed_thinking_block(delta: object) -> bool: - blocks: Final = getattr(delta, "thinking_blocks", None) or () - return any(isinstance(b, dict) and (b.get("signature") or b.get("data")) for b in blocks) + return any(_is_signed_thinking_block(block) for block in attribute_items(delta, "thinking_blocks")) def _delta_carries_output(delta: ChatCompletionDelta) -> bool: @@ -1427,7 +1432,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _is_unstreamed_empty_message(self, item: object) -> bool: if getattr(item, "type", None) != "message" or self._message_output_index is not None: return False - return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ()) + return not any(getattr(part, "text", None) for part in attribute_items(item, "content")) def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: if not isinstance(self.litellm_model_response, ModelResponse): diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e7380c26529..1ff80b0d71d 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -191,6 +191,7 @@ def mock_responses_api_response( _STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_FORWARDED_VALUE: Final = TypeAdapter(object) def _client_trace_id(request_kwargs: Mapping[str, object]) -> str | None: @@ -866,7 +867,7 @@ async def aresponses( stream: bool | None = None, temperature: float | None = None, text: Optional["ResponseText"] = None, - text_format: type["BaseModel"] | dict | None = None, + text_format: type["BaseModel"] | dict[str, object] | None = None, tool_choice: ToolChoice | None = None, tools: Iterable[ToolParam] | None = None, top_p: float | None = None, @@ -1355,7 +1356,7 @@ def _responses_try_dispatch_emulated_file_search( user: str | None, service_tier: str | None, safety_identifier: str | None, - text_format: type[BaseModel] | dict | None, + text_format: type[BaseModel] | dict[str, object] | None, allowed_openai_params: list[str] | None, extra_headers: dict[str, object] | None, extra_query: dict[str, object] | None, @@ -1390,7 +1391,7 @@ def _responses_try_dispatch_emulated_file_search( "background": background, "stream": stream, "temperature": temperature, - "text": text, + "text": _FORWARDED_VALUE.validate_python(text), "tool_choice": tool_choice, "top_p": top_p, "truncation": truncation, @@ -1435,7 +1436,7 @@ def responses( stream: bool | None = None, temperature: float | None = None, text: Optional["ResponseText"] = None, - text_format: type["BaseModel"] | dict | None = None, + text_format: type["BaseModel"] | dict[str, object] | None = None, tool_choice: ToolChoice | None = None, tools: Iterable[ToolParam] | None = None, top_p: float | None = None, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index b100858b306..cd072825ef5 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -4,6 +4,7 @@ import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call @@ -24,6 +25,8 @@ from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) + class _MCPCompletionKwargs(TypedDict, total=False, extra_items=object): """Extra keywords forwarded verbatim to ``litellm.acompletion``, which owns their contract.""" @@ -566,7 +569,7 @@ async def acompletion_with_mcp( def __getattr__(self, name: str) -> object: # Delegate all other attributes to original wrapper - return getattr(self._original_wrapper, name) + return _DELEGATED_ATTRIBUTE.validate_python(getattr(self._original_wrapper, name)) # Helper class to wrap async iterator for sync iteration class _SyncIteratorWrapper: diff --git a/litellm/responses/sse_output_recovery.py b/litellm/responses/sse_output_recovery.py index cdabcb504d3..e7d0a540022 100644 --- a/litellm/responses/sse_output_recovery.py +++ b/litellm/responses/sse_output_recovery.py @@ -11,9 +11,12 @@ import json from collections.abc import Mapping from typing import Final, SupportsInt, TypeAlias, cast # noqa: TID251 # int() re-checks the cast below at runtime +from pydantic import TypeAdapter + from litellm.constants import STREAM_SSE_DONE_STRING _MAX_CONTENT_INDEX: Final = 1024 +_DECODED_JSON: Final = TypeAdapter(object) _ConvertibleToInt: TypeAlias = SupportsInt | str @@ -33,7 +36,7 @@ def parse_sse_json_chunk(chunk: str) -> dict[str, object] | None: if not stripped_chunk or stripped_chunk == STREAM_SSE_DONE_STRING or stripped_chunk.startswith("event:"): return None try: - parsed_chunk: Final[object] = json.loads(stripped_chunk) + parsed_chunk: Final = _DECODED_JSON.validate_python(json.loads(stripped_chunk)) except json.JSONDecodeError: return None if not isinstance(parsed_chunk, dict): diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 0d59c93786c..50ab3e4f75a 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -28,7 +28,7 @@ from typing import ( import httpx from openai._streaming import SSEDecoder -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import TypeIs import litellm @@ -166,6 +166,7 @@ def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: _SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" _UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" +_DECODED_JSON: Final = TypeAdapter(object) def _load_json_object(payload: str | bytes) -> dict[str, object]: @@ -175,7 +176,7 @@ def _load_json_object(payload: str | bytes) -> dict[str, object]: def _load_json_value(payload: str | bytes) -> object: """Parse a JSON payload whose top-level shape the caller narrows itself.""" - return json.loads(payload) + return _DECODED_JSON.validate_python(json.loads(payload)) def _model_id_from_metadata(litellm_metadata: Mapping[str, object] | None) -> str | None: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index f256480d952..dc75e77d685 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -5,10 +5,12 @@ from functools import reduce from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload from pydantic import BaseModel, TypeAdapter +from typing_extensions import TypedDict as ExtraItemsTypedDict from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire payloads without a runtime conversion import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel from litellm.types.llms.openai import ( @@ -33,6 +35,10 @@ from litellm.types.utils import ( ) +class _UsageExtraFields(ExtraItemsTypedDict, total=False, extra_items=object): + pass + + def _apply_nested_drop_params(params: dict[str, object], additional_drop_params: list[str] | None) -> dict[str, object]: nested_paths: Final = tuple(path for path in additional_drop_params or () if is_nested_path(path)) return reduce(lambda acc, path: delete_nested_value(acc, path), nested_paths, params) @@ -64,6 +70,28 @@ def _is_object_dict( return isinstance(value, dict) +def _prompt_tokens_details_wrapper(details: object) -> PromptTokensDetailsWrapper: + if _is_object_dict(details): + input_details: Final = InputTokensDetails.model_validate(details) + return PromptTokensDetailsWrapper( + **input_details.model_dump(exclude={"cache_write_tokens_details"}), + cache_creation_token_details=_cache_write_lifetimes(input_details.cache_write_tokens_details), + ) + return PromptTokensDetailsWrapper( + cached_tokens=getattr(details, "cached_tokens", None), + audio_tokens=getattr(details, "audio_tokens", None), + text_tokens=getattr(details, "text_tokens", None), + image_tokens=getattr(details, "image_tokens", None), + cached_tokens_details=getattr(details, "cached_tokens_details", None), + video_tokens=getattr(details, "video_tokens", None), + cache_write_tokens=getattr(details, "cache_write_tokens", None), + cache_creation_token_details=_cache_write_lifetimes(getattr(details, "cache_write_tokens_details", None)), + cache_type=getattr(details, "cache_type", None), + web_search_requests=getattr(details, "web_search_requests", None), + google_maps_grounding_requests=getattr(details, "google_maps_grounding_requests", None), + ) + + def normalize_responses_api_stream_options( stream_options: object, ) -> ResponsesAPIStreamOptions | None: @@ -1060,7 +1088,7 @@ class ResponsesAPIRequestUtils: @staticmethod def _update_container_ids_in_response( - responses_api_response: ResponsesAPIResponse | dict[str, Any], + responses_api_response: ResponsesAPIResponse | dict[str, object], custom_llm_provider: str | None, litellm_metadata: dict[str, Any] | None = None, ) -> ResponsesAPIResponse | dict[str, object]: @@ -1074,15 +1102,16 @@ class ResponsesAPIRequestUtils: model_id: Final = model_info.get("id") # Get the output list - if isinstance(responses_api_response, dict): - output = responses_api_response.get("output", []) - else: - output = getattr(responses_api_response, "output", []) + output: Final = ( + responses_api_response.get("output", []) + if isinstance(responses_api_response, dict) + else attribute_or_none(responses_api_response, "output") + ) if not output: return responses_api_response - for item in output: + for item in iter_items(output): ResponsesAPIRequestUtils.encode_container_id_on_output_item( item=item, custom_llm_provider=custom_llm_provider, @@ -1257,34 +1286,11 @@ class ResponseAPILoggingUtils: response_api_usage = usage_input prompt_tokens: Final[int] = response_api_usage.input_tokens or 0 completion_tokens: Final[int] = response_api_usage.output_tokens or 0 - prompt_tokens_details: PromptTokensDetailsWrapper | None = None - if response_api_usage.input_tokens_details: - if isinstance(response_api_usage.input_tokens_details, dict): - input_details: Final = InputTokensDetails.model_validate(response_api_usage.input_tokens_details) - prompt_tokens_details = PromptTokensDetailsWrapper( - **input_details.model_dump(exclude={"cache_write_tokens_details"}), - cache_creation_token_details=_cache_write_lifetimes(input_details.cache_write_tokens_details), - ) - else: - prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None), - audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), - text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), - image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), - cached_tokens_details=getattr( - response_api_usage.input_tokens_details, "cached_tokens_details", None - ), - video_tokens=getattr(response_api_usage.input_tokens_details, "video_tokens", None), - cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), - cache_creation_token_details=_cache_write_lifetimes( - getattr(response_api_usage.input_tokens_details, "cache_write_tokens_details", None) - ), - cache_type=getattr(response_api_usage.input_tokens_details, "cache_type", None), - web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), - google_maps_grounding_requests=getattr( - response_api_usage.input_tokens_details, "google_maps_grounding_requests", None - ), - ) + prompt_tokens_details: Final = ( + _prompt_tokens_details_wrapper(response_api_usage.input_tokens_details) + if response_api_usage.input_tokens_details + else None + ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None output_tokens_details: Final[OutputTokensDetails | None] = getattr( response_api_usage, "output_tokens_details", None @@ -1308,20 +1314,22 @@ class ResponseAPILoggingUtils: audio_tokens=audio_tokens, ) - extra_usage_fields: Final = { - key: value - for key, value in (response_api_usage.model_extra or {}).items() - if key - not in ( - "input_token_details", - "output_token_details", - "prompt_tokens", - "completion_tokens", - "total_tokens", - "prompt_tokens_details", - "completion_tokens_details", - ) - } + extra_usage_fields: Final = _UsageExtraFields( + **{ + key: value + for key, value in (as_str_mapping(response_api_usage.model_extra) or {}).items() + if key + not in ( + "input_token_details", + "output_token_details", + "prompt_tokens", + "completion_tokens", + "total_tokens", + "prompt_tokens_details", + "completion_tokens_details", + ) + } + ) chat_usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, diff --git a/litellm/router.py b/litellm/router.py index 705accbcf8d..40c5c8afa6c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -934,7 +934,7 @@ MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 def _retry_policy_ceiling(policy: RetryPolicy) -> int: """The most retries any error class under this policy can be granted.""" - retry_counts: Final[Iterable[object]] = policy.model_dump().values() + retry_counts: Final[tuple[object, ...]] = tuple(policy.model_dump().values()) return max((retries for retries in retry_counts if isinstance(retries, int)), default=0) @@ -6504,7 +6504,7 @@ class Router: has_generated_content: bool, buffered_lifecycle_chunks: tuple[bytes, ...], model: str, - initial_kwargs: dict[str, Any], # mutable-ok: handed to _aanthropic_messages_fallback_attempt, which mutates it + initial_kwargs: dict[str, object], # mutable-ok: _aanthropic_messages_fallback_attempt mutates it wrapper: "FallbackAwareAnthropicMessagesStream", close_superseded: Callable[[], Awaitable[None]], ) -> AsyncGenerator[bytes, None]: diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index 1d119375778..d8c52b4c8b6 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -47,13 +47,13 @@ class QualityRouter(CustomLogger): model_name: str, litellm_router_instance: "Router", default_model: str | None = None, - quality_router_config: dict[str, Any] | None = None, + quality_router_config: Mapping[str, object] | None = None, ): self.model_name = model_name self.litellm_router_instance = litellm_router_instance if quality_router_config: - self.config = QualityRouterConfig(**quality_router_config) + self.config = QualityRouterConfig.model_validate(dict(**quality_router_config)) else: self.config = QualityRouterConfig() @@ -185,7 +185,7 @@ class QualityRouter(CustomLogger): if isinstance(raw_prefs, RoutingPreferences): prefs = raw_prefs elif isinstance(raw_prefs, dict): - prefs = RoutingPreferences(**raw_prefs) + prefs = RoutingPreferences.model_validate(dict(**raw_prefs)) else: # A Pydantic object of some other shape — coerce via its dict. prefs = RoutingPreferences( diff --git a/litellm/sandbox/__init__.py b/litellm/sandbox/__init__.py index 1974859fb1e..bcbb7bc5743 100644 --- a/litellm/sandbox/__init__.py +++ b/litellm/sandbox/__init__.py @@ -7,7 +7,10 @@ They resolve lazily so `import litellm` does not pull in litellm.harness. import importlib from typing import Final +from pydantic import TypeAdapter + _HARNESS_SANDBOX_MODULE: Final = "litellm.harness.sandbox" +_HARNESS_EXPORT: Final = TypeAdapter(object) _HARNESS_EXPORTS: Final = frozenset( { "local", @@ -23,5 +26,5 @@ _HARNESS_EXPORTS: Final = frozenset( def __getattr__(name: str) -> object: if name in _HARNESS_EXPORTS: - return getattr(importlib.import_module(_HARNESS_SANDBOX_MODULE), name) + return _HARNESS_EXPORT.validate_python(getattr(importlib.import_module(_HARNESS_SANDBOX_MODULE), name)) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/secret_managers/base_secret_manager.py b/litellm/secret_managers/base_secret_manager.py index e537dd1d998..9c25086c451 100644 --- a/litellm/secret_managers/base_secret_manager.py +++ b/litellm/secret_managers/base_secret_manager.py @@ -1,6 +1,6 @@ import re from abc import ABC, abstractmethod -from typing import Any, Final +from typing import Final import httpx @@ -73,7 +73,7 @@ class BaseSecretManager(ABC): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Asynchronously write a secret to the secret manager. diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index a1ec7c0a386..abfa9649e91 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -1,7 +1,7 @@ import asyncio import base64 import os -from typing import Any, Final +from typing import Final from urllib.parse import quote import httpx @@ -278,7 +278,7 @@ class CyberArkSecretManager(BaseSecretManager): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Writes a secret to CyberArk Conjur using an async HTTPX client. diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index f4234e26c60..b33cf0e6b7e 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -24,6 +24,7 @@ from litellm.types.llms.base import LiteLLMBaseModel oidc_cache: Final = DualCache() _PARSED_LITERAL: Final = TypeAdapter(object) +_SECRET_MANAGER_CLIENT: Final = TypeAdapter(object) @lru_cache(maxsize=128) @@ -383,7 +384,7 @@ def get_secret( try: if _should_read_secret_from_secret_manager() and litellm.secret_manager_client is not None: try: - client: Final = litellm.secret_manager_client + client: Final = _SECRET_MANAGER_CLIENT.validate_python(litellm.secret_manager_client) key_manager = "local" if key_management_system is not None: key_manager = key_management_system.value diff --git a/litellm/utils.py b/litellm/utils.py index bb94ca6d40a..1db5f7129df 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1641,7 +1641,7 @@ def _deployment_provider(kwargs: Mapping[str, object]) -> str: return "" -async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) -> dict[str, Any]: +async def async_pre_call_deployment_hook(kwargs: Mapping[str, object], call_type: str) -> dict[str, Any]: """ Allow modifying the request just before it's sent to the deployment. @@ -4287,7 +4287,7 @@ def get_optional_params_embeddings( # retrieve all parameters passed to the function passed_params: Final = locals() - custom_llm_provider = passed_params.pop("custom_llm_provider", None) + passed_params.pop("custom_llm_provider", None) special_params: Final = passed_params.pop("kwargs") passed_params.pop("drop_params", None) @@ -4325,7 +4325,7 @@ def get_optional_params_embeddings( provider_config: BaseEmbeddingConfig | None = None optional_params = {} - if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): + if custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_embedding_config( model=model, provider=LlmProviders(custom_llm_provider), @@ -9060,8 +9060,9 @@ _get_bundled_model_cost_map = get_bundled_model_cost_map def _get_model_cost_entry_for_provider_config( model: str, provider: LlmProviders, -) -> dict[str, object]: +) -> Mapping[str, object]: candidate_keys: Final = (model, f"{provider.value}/{model}") + model_info: Mapping[str, object] | None for model_key in candidate_keys: model_info = litellm.model_cost.get(model_key) if model_info is not None: @@ -10886,18 +10887,12 @@ def extract_duration_from_srt_or_vtt(srt_or_vtt_content: str) -> float | None: # Regular expression to match timestamps in the format "hh:mm:ss,ms" or "hh:mm:ss.ms" timestamp_pattern: Final = r"(\d{2}):(\d{2}):(\d{2})[.,](\d{3})" - timestamps: Final[Sequence[tuple[str, str, str, str]]] = re.findall(timestamp_pattern, srt_or_vtt_content) - - if not timestamps: - return None + timestamps: Final = re.finditer(timestamp_pattern, srt_or_vtt_content) # Convert timestamps to seconds and find the max (end time) - durations: Final = [] - match: tuple[str, str, str, str] - for match in timestamps: - hours, minutes, seconds, milliseconds = map(int, match) - total_seconds = hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0 - durations.append(total_seconds) + durations: Final = [ + int(match[1]) * 3600 + int(match[2]) * 60 + int(match[3]) + int(match[4]) / 1000.0 for match in timestamps + ] return max(durations) if durations else None diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 8b33b1ce12c..9f53accfa0b 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -5,7 +5,7 @@ LiteLLM SDK Functions for Creating and Searching Vector Stores import asyncio import builtins import contextvars -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Iterator, Mapping from functools import partial from typing import TYPE_CHECKING, Final @@ -53,6 +53,10 @@ def _direct_vector_store_embedding_executor( ) +def _items_without_key(values: Mapping[str, object], excluded_key: str) -> Iterator[tuple[str, object]]: + return ((key, value) for key, value in values.items() if key != excluded_key) + + def mock_vector_store_search_response( mock_results: builtins.list[VectorStoreSearchResult] | None = None, ): @@ -307,7 +311,7 @@ async def asearch( embedding_executor: Final = _direct_vector_store_embedding_executor( kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) - local_vars: Final = {key: value for key, value in locals().items() if key != "embedding_executor"} + local_vars: Final = dict(_items_without_key(locals(), "embedding_executor")) try: loop: Final = asyncio.get_event_loop() @@ -391,7 +395,7 @@ def search( embedding_executor: Final = _direct_vector_store_embedding_executor( kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) - local_vars: Final = {key: value for key, value in locals().items() if key != "embedding_executor"} + local_vars: Final = dict(_items_without_key(locals(), "embedding_executor")) try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) @@ -694,7 +698,7 @@ def list( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -): +) -> Mapping[str, object] | Coroutine[object, object, Mapping[str, object]]: """ List vector stores. @@ -992,7 +996,7 @@ def delete( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -): +) -> Mapping[str, object] | Coroutine[object, object, Mapping[str, object]]: """ Delete a vector store. diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 7610d602fff..a5bb5e0b340 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -42,7 +42,7 @@ IGNORE_FUNCTIONS = [ "contains_encrypted_marker", # max depth set. "_rotate_guardrail_row", # bounded by attempts_left. "_delete_nested_value_custom", # max depth set (bounded by number of path segments). - "filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. + "_filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. "__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion. "_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation. "_basic_json_schema_validate", # max depth set. diff --git a/tests/unit/a2a_protocol/test_utils.py b/tests/unit/a2a_protocol/test_utils.py index 8a219211c62..dfa8449f6e6 100644 --- a/tests/unit/a2a_protocol/test_utils.py +++ b/tests/unit/a2a_protocol/test_utils.py @@ -1,5 +1,8 @@ """Tests for litellm/a2a_protocol/utils.py token/usage extraction.""" +from types import SimpleNamespace +from typing import Final + import pytest pytest.importorskip("a2a.compat.v0_3.types") @@ -7,6 +10,7 @@ pytest.importorskip("a2a.compat.v0_3.types") from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest from litellm.a2a_protocol.utils import A2ARequestUtils +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder def _request(user_text: str) -> SendMessageRequest: @@ -39,3 +43,13 @@ def test_calculate_usage_counts_input_tokens_from_request_object(): assert prompt_tokens > 0 assert completion_tokens > 0 assert total_tokens == prompt_tokens + completion_tokens + + +def test_calculate_usage_counts_input_tokens_from_a_message_that_forwards_its_fields(): + request: Final = _request("count these input tokens please") + forwarded_request: Final = SimpleNamespace(params=SimpleNamespace(message=AttributeForwarder(request.params.message))) + response_dict: Final = {"result": {"kind": "message", "parts": [{"kind": "text", "text": "ok"}]}} + + assert A2ARequestUtils.calculate_usage_from_request_response( + request=forwarded_request, response_dict=response_dict + ) == A2ARequestUtils.calculate_usage_from_request_response(request=request, response_dict=response_dict) diff --git a/tests/unit/containers/test_container_proxy_ownership.py b/tests/unit/containers/test_container_proxy_ownership.py index 780ae9bf6d1..ac5adeee8c1 100644 --- a/tests/unit/containers/test_container_proxy_ownership.py +++ b/tests/unit/containers/test_container_proxy_ownership.py @@ -1,6 +1,7 @@ import json import sys from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -10,6 +11,7 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.container_endpoints import ownership from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.containers.main import ContainerListResponse, ContainerObject +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder @pytest.fixture(autouse=True) @@ -97,6 +99,23 @@ async def test_should_not_mutate_dict_container_response_when_recording_owner( assert file_obj["provider_container_id"] == "cntr_provider" +@pytest.mark.asyncio +async def test_should_store_the_full_container_of_a_response_that_forwards_its_fields(monkeypatch): + table: Final = AsyncMock() + table.find_unique.return_value = None + prisma_client: Final = SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table)) + monkeypatch.setattr(ownership, "_get_prisma_client", AsyncMock(return_value=prisma_client)) + + await ownership.record_container_owner( + response=AttributeForwarder(_container("cntr_provider")), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + custom_llm_provider="openai", + ) + + file_obj: Final = json.loads(table.create.await_args.kwargs["data"]["file_object"]) + assert file_obj["status"] == "active" + + @pytest.mark.asyncio async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 20ed8ac0c10..7f6e955e831 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -31,6 +31,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( encode_file_id_with_model, ) from litellm.types.llms.openai import OpenAIFileObject +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder class _InMemoryManagedFileTable: @@ -451,6 +452,20 @@ async def test_async_pre_call_deployment_hook_prefers_top_level_model_info(): ), "Should prefer top-level model_info over litellm_metadata" +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_reads_model_info_that_forwards_its_fields(): + proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()) + kwargs: Final = { + "input_file_id": "managed-file-abc", + "model_file_id_mapping": {"managed-file-abc": {"deployment-xyz": "file-provider-123"}}, + "model_info": AttributeForwarder({"id": "deployment-xyz"}), + } + + result: Final = await proxy_managed_files.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acreate_batch) + + assert result["input_file_id"] == "file-provider-123" + + @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_model_info_leaves_file_id_unchanged(): """ diff --git a/tests/unit/integrations/dotprompt/test_prompt_manager.py b/tests/unit/integrations/dotprompt/test_prompt_manager.py index dbd4e4a4c4c..cbf19fe7985 100644 --- a/tests/unit/integrations/dotprompt/test_prompt_manager.py +++ b/tests/unit/integrations/dotprompt/test_prompt_manager.py @@ -7,6 +7,7 @@ from fastapi.testclient import TestClient +from typing import Final from unittest.mock import MagicMock, Mock, patch import httpx @@ -362,6 +363,22 @@ def test_load_prompts_from_json_data(): assert rendered == "Dynamic prompt: test" +def test_load_prompts_from_json_data_skips_malformed_entries(): + manager: Final = PromptManager(prompt_directory=str(Path(__file__).parent)) + initial_count: Final = len(manager.prompts) + + manager.load_prompts_from_json_data( + { + "valid_prompt": {"content": "Valid: {{value}}", "metadata": {"model": "gpt-4"}}, + "string_entry": "not a prompt", + "null_entry": None, + } + ) + + assert len(manager.prompts) == initial_count + 1 + assert manager.render("valid_prompt", {"value": "ok"}) == "Valid: ok" + + def test_prompt_file_to_json_conversion(): """Test converting .prompt files to JSON format.""" # Create a temporary prompt file with frontmatter @@ -661,6 +678,27 @@ def test_prompt_initializer_registers_flat_db_prompt_under_base_id(): assert template.content == "AHOY {{name}}" +def test_prompt_initializer_accepts_params_without_dotprompt_content(): + from types import SimpleNamespace + + from litellm.integrations.dotprompt import prompt_initializer + from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec + + litellm_params: Final = SimpleNamespace( + prompt_data={"content": "AHOY {{name}}", "metadata": {"model": "gpt-4o-mini"}}, + ) + prompt_spec: Final = PromptSpec( + prompt_id="agent-prompt", + litellm_params=PromptLiteLLMParams(prompt_integration="dotprompt"), + ) + + dotprompt_manager: Final = prompt_initializer(litellm_params, prompt_spec) + + template: Final = dotprompt_manager.prompt_manager.get_prompt("agent-prompt") + assert template is not None + assert template.content == "AHOY {{name}}" + + def _swap_prompt_manager_and_spec(ignore_prompt_manager_model: bool) -> tuple[DotpromptManager, PromptSpec]: manager = DotpromptManager( prompt_data={"content": "You are a pirate assistant.", "metadata": {"model": "gpt-4o-mini"}}, diff --git a/tests/unit/integrations/newrelic/test_newrelic.py b/tests/unit/integrations/newrelic/test_newrelic.py index 00a115dd8f2..d09732c6eb2 100644 --- a/tests/unit/integrations/newrelic/test_newrelic.py +++ b/tests/unit/integrations/newrelic/test_newrelic.py @@ -1,6 +1,7 @@ import os import sys from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -491,6 +492,18 @@ class TestExtractAllMessagesTimestamps: input_msg = next(m for m in messages if not m.get("is_response")) assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0) + def test_input_messages_read_a_start_time_that_only_exposes_timestamp(self): + class EpochSeconds: + def timestamp(self) -> float: + return 1_000_000.0 + + kwargs: Final = {**make_kwargs(messages=[{"role": "user", "content": "Hi"}]), "start_time": EpochSeconds()} + + messages: Final = self.logger._extract_all_messages(kwargs, make_response(), response_model="gpt-4", vendor="openai") + + input_msg: Final = next(m for m in messages if not m.get("is_response")) + assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0) + def test_output_messages_get_end_time_timestamp(self): kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) response = make_response() diff --git a/tests/unit/litellm_core_utils/attribute_forwarder.py b/tests/unit/litellm_core_utils/attribute_forwarder.py new file mode 100644 index 00000000000..3056057fc97 --- /dev/null +++ b/tests/unit/litellm_core_utils/attribute_forwarder.py @@ -0,0 +1,6 @@ +class AttributeForwarder: + def __init__(self, target: object) -> None: + self._target = target + + def __getattr__(self, name: str) -> object: + return getattr(self._target, name) diff --git a/tests/unit/llms/compactifai/test_compactifai.py b/tests/unit/llms/compactifai/test_compactifai.py index 1367c703fda..56d5374f383 100644 --- a/tests/unit/llms/compactifai/test_compactifai.py +++ b/tests/unit/llms/compactifai/test_compactifai.py @@ -2,7 +2,7 @@ import json import os import sys from unittest.mock import AsyncMock, patch -from typing import Optional +from typing import Final, Optional import httpx import pytest @@ -271,3 +271,38 @@ async def test_compactifai_async_completion(respx_mock): assert response.choices[0].message.content == "Async response from CompactifAI" assert response.usage.total_tokens == 23 + + +@pytest.mark.parametrize( + ("body", "json_mode", "error", "message"), + [ + ([], None, TypeError, r"argument after \*\* must be a mapping, not list"), + ({"choices": None}, True, TypeError, "'NoneType' object is not iterable"), + ({"choices": [{"message": "hi"}]}, True, AttributeError, "'str' object has no attribute 'get'"), + ({"choices": [{"message": {"tool_calls": [5]}}]}, True, TypeError, "'int' object is not subscriptable"), + ], +) +def test_compactifai_malformed_response_body_raises_the_error_of_the_failing_lookup(body, json_mode, error, message): + from unittest.mock import MagicMock + + from litellm.llms.compactifai.chat.transformation import CompactifAIChatConfig + + raw_response: Final = httpx.Response( + 200, json=body, request=httpx.Request("POST", "https://api.compactif.ai/v1/chat/completions") + ) + + with pytest.raises(error, match=message) as raised: + CompactifAIChatConfig().transform_response( + model="cai-llama-3-1-8b-slim", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + json_mode=json_mode, + ) + + assert raised.type is error diff --git a/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py index aff90549a7b..0951ba6ba92 100644 --- a/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py +++ b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py @@ -1,6 +1,7 @@ import io import os import pathlib +from typing import Final from unittest.mock import MagicMock import httpx @@ -594,3 +595,24 @@ def test_transform_response_wraps_malformed_words_with_the_raw_body(words: objec DeepgramAudioTranscriptionConfig().transform_audio_transcription_response(raw_response) assert str(exc_info.value).endswith(f"\nResponse: {raw_response.text}") + + +@pytest.mark.parametrize( + ("payload", "error"), + [ + ({"results": None}, "'NoneType' object is not subscriptable"), + (_deepgram_payload({"transcript": "hello", "words": 5}, {}), "object of type 'int' has no len()"), + ( + {**_deepgram_payload({"transcript": "hello"}, {}), "metadata": None}, + "'NoneType' object is not subscriptable", + ), + ], +) +def test_transform_response_wraps_the_error_of_the_failing_lookup_in_a_malformed_body(payload: object, error: str): + raw_response: Final = httpx.Response(200, json=payload) + + with pytest.raises(ValueError, match="Error transforming Deepgram response: ") as exc_info: + DeepgramAudioTranscriptionConfig().transform_audio_transcription_response(raw_response) + + assert exc_info.type is ValueError + assert str(exc_info.value) == f"Error transforming Deepgram response: {error}\nResponse: {raw_response.text}" diff --git a/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py b/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py index efdb428db32..ee6ee187d64 100644 --- a/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py +++ b/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py @@ -98,6 +98,12 @@ class TestEmbedding: class TestErrors: + def test_a_non_object_body_fails_as_an_invalid_response_object(self, eden_key, respx_mock): + respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(200, json=["unexpected"])) + + with pytest.raises(litellm.APIConnectionError, match="Invalid response object"): + litellm.embedding(model=MODEL, input="hello") + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): respx_mock.post(EDEN_EMBEDDINGS_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) diff --git a/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py b/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py index d7b32affc70..4e0673b8c27 100644 --- a/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py +++ b/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py @@ -117,6 +117,12 @@ class TestImageGeneration: class TestErrors: + def test_a_non_object_body_fails_as_an_invalid_response_object(self, eden_key, respx_mock): + respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(200, json=["unexpected"])) + + with pytest.raises(litellm.APIConnectionError, match="Invalid response object"): + litellm.image_generation(model=MODEL, prompt="a red square") + def test_middleware_401_maps_to_authentication_error(self, eden_key, respx_mock): respx_mock.post(EDEN_IMAGES_URL).mock(return_value=httpx.Response(401, json={"detail": "Invalid token."})) diff --git a/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py b/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py index da8ab138416..175ce040d9f 100644 --- a/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py +++ b/tests/unit/llms/gemini/google_genai/test_google_genai_transformation.py @@ -106,6 +106,26 @@ def test_generate_content_response_parses_function_call() -> None: } +@pytest.mark.parametrize( + ("body", "message"), + [ + ([], r"argument after \*\* must be a mapping, not list"), + (["candidates"], "list indices must be integers or slices, not str"), + ], +) +def test_generate_content_response_non_object_body_raises_the_error_of_the_failing_lookup( + body: JsonValue, message: str +) -> None: + with pytest.raises(TypeError, match=message) as raised: + GoogleGenAIConfig().transform_generate_content_response( + model="gemini-2.5-flash", + raw_response=httpx.Response(200, json=body), + logging_obj=SimpleNamespace(model_call_details={}), + ) + + assert raised.type is TypeError + + def test_generate_content_request_preserves_tool_declarations() -> None: tools: Final = [ { diff --git a/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index bd9b7006e58..c21dec9b6bb 100644 --- a/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -241,6 +241,26 @@ class TestGeminiImageEditTransformation: ) assert logging_usage["completion_tokens_details"]["image_tokens"] == 1716 + @pytest.mark.parametrize( + ("body", "error", "message"), + [ + (["not", "an", "object"], AttributeError, "'list' object has no attribute 'get'"), + ({"candidates": None}, TypeError, "'NoneType' object is not iterable"), + ({"candidates": [{"content": {"parts": None}}]}, TypeError, "'NoneType' object is not iterable"), + ], + ) + def test_transform_image_edit_response_malformed_body_raises_the_error_of_the_failing_lookup( + self, body: object, error: type[Exception], message: str + ) -> None: + with pytest.raises(error, match=message) as raised: + self.config.transform_image_edit_response( + model=self.model, + raw_response=httpx.Response(200, json=body), + logging_obj=self.logging_obj, + ) + + assert raised.type is error + def test_transform_image_edit_request_without_image_raises(self) -> None: optional_params = {} diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 196b912752b..20c36878fef 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -355,6 +355,25 @@ class TestGeminiModelInfo: assert result == expected +@pytest.mark.parametrize( + ("body", "error", "message"), + [ + ([], TypeError, "list indices must be integers or slices, not str"), + ({"models": None}, TypeError, "'NoneType' object is not iterable"), + ({"models": [{"name": 1}]}, AttributeError, "'int' object has no attribute 'replace'"), + ], +) +def test_get_models_malformed_listing_raises_the_error_of_the_failing_lookup( + body: JsonValue, error: type[Exception], message: str, respx_mock: respx.MockRouter +) -> None: + respx_mock.get("https://models.test/v1beta/models").respond(200, json=body) + + with pytest.raises(error, match=message) as raised: + GeminiModelInfo().get_models(api_key="test-key", api_base="https://models.test") + + assert raised.type is error + + class TestGoogleAIStudioTokenCounter: """Test suite for GoogleAIStudioTokenCounter class""" diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index 0a5039791ee..8995b3cb068 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -8,6 +8,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.transformation impor from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder def test_transform_basic_request(): @@ -207,6 +208,23 @@ def test_messages_to_responses_input_with_tool(): } +def test_messages_to_responses_input_reads_a_tool_call_that_forwards_its_fields(): + tool_call: Final = {"id": "call_123", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + messages: Final = [ + {"role": "user", "content": "What's the weather?"}, + {"role": "assistant", "content": None, "tool_calls": [AttributeForwarder(tool_call)]}, + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items[1] == { + "type": "function_call", + "call_id": "call_123", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + } + + def test_messages_to_responses_input_preserves_images(): """An image block must survive the round trip, or OpenAI counts only the text. diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 2f10fd6fd0d..0d1b1aca9bc 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -47,6 +47,7 @@ from litellm.types.responses.main import ( OutputToolSearchCall, ) from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder class MockGuardrail(CustomGuardrail): @@ -1783,6 +1784,22 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_events_that_forward_their_fields(self): + handler: Final = OpenAIResponsesHandler() + plain_events: Final = self._ended_stream_events() + forwarded_delta: Final = AttributeForwarder(plain_events[0]) + events: Final = [forwarded_delta, *plain_events[1:]] + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert forwarded_delta.delta == "hello [MASKED]" + @staticmethod def _ended_function_call_stream_events() -> List[dict]: def item(arguments: str, status: str) -> dict: diff --git a/tests/unit/llms/runwayml/image_generation/__init__.py b/tests/unit/llms/runwayml/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py b/tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py new file mode 100644 index 00000000000..7eec9541b98 --- /dev/null +++ b/tests/unit/llms/runwayml/image_generation/test_runway_image_generation_transformation.py @@ -0,0 +1,87 @@ +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.runwayml.image_generation.transformation import RunwayMLImageGenerationConfig +from litellm.types.utils import ImageResponse + +TASK_URL = "https://api.dev.runwayml.com/v1/tasks/task_1" + + +@pytest.mark.respx() +def test_output_dict_with_non_string_values_still_yields_the_image(respx_mock): + respx_mock.get(TASK_URL).respond( + json={ + "id": "task_1", + "status": "SUCCEEDED", + "output": [{"url": "https://cdn.example/image.png", "b64_json": None, "width": 1024}], + }, + status_code=200, + ) + created: Final = httpx.Response( + 200, + json={"id": "task_1", "status": "PENDING"}, + request=httpx.Request( + "POST", "https://api.dev.runwayml.com/v1/text_to_image", headers={"Authorization": "Bearer k"} + ), + ) + + image_response: Final = RunwayMLImageGenerationConfig().transform_image_generation_response( + model="gen4_image", + raw_response=created, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert [image.url for image in image_response.data] == ["https://cdn.example/image.png"] + + +def test_non_object_task_body_raises_the_base_error_class(): + created: Final = httpx.Response( + 200, + json=["unexpected"], + request=httpx.Request( + "POST", "https://api.dev.runwayml.com/v1/text_to_image", headers={"Authorization": "Bearer k"} + ), + ) + + with pytest.raises(AttributeError): + RunwayMLImageGenerationConfig().transform_image_generation_response( + model="gen4_image", + raw_response=created, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.asyncio +async def test_async_non_object_task_body_raises_the_base_error_class(): + created: Final = httpx.Response( + 200, + json=["unexpected"], + request=httpx.Request( + "POST", "https://api.dev.runwayml.com/v1/text_to_image", headers={"Authorization": "Bearer k"} + ), + ) + + with pytest.raises(AttributeError): + await RunwayMLImageGenerationConfig().async_transform_image_generation_response( + model="gen4_image", + raw_response=created, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) diff --git a/tests/unit/proxy/auth/test_mcp_ip_filtering.py b/tests/unit/proxy/auth/test_mcp_ip_filtering.py index 4c219760762..c403daf733f 100644 --- a/tests/unit/proxy/auth/test_mcp_ip_filtering.py +++ b/tests/unit/proxy/auth/test_mcp_ip_filtering.py @@ -6,6 +6,7 @@ external callers only see servers with available_on_public_internet=True. """ import logging +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -176,6 +177,19 @@ class TestMCPClientIPExtraction: assert result == "203.0.113.5" + def test_non_boolean_use_x_forwarded_for_never_trusts_xff(self): + request: Final = _make_request("10.0.0.5", "203.0.113.9, 10.0.0.5") + + result: Final = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": "false", + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + + assert result == "10.0.0.5" + def _make_request(client_host, xff): request = MagicMock(spec=Request) diff --git a/tests/unit/proxy/client/cli/test_encryption_commands.py b/tests/unit/proxy/client/cli/test_encryption_commands.py index 3a86eb82593..6685d29fb7e 100644 --- a/tests/unit/proxy/client/cli/test_encryption_commands.py +++ b/tests/unit/proxy/client/cli/test_encryption_commands.py @@ -5,6 +5,8 @@ check vs POST migrate, dry-run param) and its residual-state messaging without a live proxy. """ +from typing import Final + import pytest from click.testing import CliRunner @@ -69,3 +71,20 @@ def test_migrate_reports_residual_legacy(runner): result = runner.invoke(cli_main.cli, ["encryption", "migrate", "--check"]) assert result.exit_code == 0, result.output assert "Residual legacy values remaining: 3" in result.output + + +@pytest.mark.parametrize("residual", [-1, float("nan")]) +def test_migrate_prints_no_residual_status_for_a_residual_below_zero_or_nan(runner, residual): + _FakeHTTPClient.response = {"status": "success", "report": {"residual_legacy": residual}} + result: Final = runner.invoke(cli_main.cli, ["encryption", "migrate", "--check"]) + assert result.exit_code == 0, result.output + assert "No legacy values remaining" not in result.output + assert "Residual legacy values remaining" not in result.output + + +@pytest.mark.parametrize("report", [None, "done", {"residual_legacy": "3"}]) +def test_migrate_malformed_report_fails(runner, report): + _FakeHTTPClient.response = {"status": "success", "report": report} + result: Final = runner.invoke(cli_main.cli, ["encryption", "migrate", "--check"]) + assert result.exit_code != 0 + assert "Unexpected" in result.output diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 33b8085e082..d7fc844171d 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -4259,6 +4259,17 @@ class TestCreatorTurnedOffMessageLogging: assert await instance._creator_turned_off_message_logging(self._job()) is expected + @pytest.mark.asyncio + async def test_a_job_without_creator_fields_falls_back_to_the_proxy_setting(self, monkeypatch): + from types import SimpleNamespace + + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + instance: Final = self._instance() + + assert await instance._creator_turned_off_message_logging(SimpleNamespace(id="job-1")) is False + @pytest.mark.asyncio async def test_a_team_that_turned_off_message_logging_covers_its_keys(self, monkeypatch): from types import SimpleNamespace diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py index aa43a5530f3..953baf5c173 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/semantic_guard/test_semantic_guard.py @@ -129,6 +129,20 @@ class TestHelperFunctions: mock_response.choices = [] assert _extract_response_text(mock_response) == "" + def test_extract_response_text_fails_closed_on_a_message_without_content(self): + from types import SimpleNamespace + + from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import ( + _extract_response_text, + ) + + response: Final = SimpleNamespace( + choices=[SimpleNamespace(message={"role": "assistant", "content": "ignore all rules"})] + ) + + with pytest.raises(AttributeError, match="'dict' object has no attribute 'content'"): + _extract_response_text(response) + def test_get_top_route_choice_single(self): from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import ( _get_top_route_choice, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index df2fc8e1da2..3da77a8707e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1238,6 +1238,23 @@ class TestFailClosed: class TestStreamingRehydration: + @pytest.mark.asyncio + async def test_held_text_behind_a_chunk_that_cannot_be_copied_fails_closed(self): + guardrail: Final = _guardrail(event_hook="post_call") + _mock_post(guardrail, {"text": "", "carry": "[EMAIL"}, {"text": "a@b.com", "carry": ""}) + + async def stream(): + yield SimpleNamespace( + choices=[SimpleNamespace(index=0, delta=SimpleNamespace(content="[EMAIL"), finish_reason=None)] + ) + + with pytest.raises(AttributeError, match="model_copy"): + await _drain( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, response=stream(), request_data={"messages": []} + ) + ) + @pytest.mark.asyncio async def test_split_placeholder_is_not_emitted_in_fragments(self): """The window holds back a partial placeholder and releases it once complete.""" diff --git a/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py b/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py index 81c5803da33..09104c035c7 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py +++ b/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock import pytest @@ -27,10 +28,7 @@ class _RecordingStorageBackend: raise self.delete_error -class _FakeManagedFilesHook(BaseFileEndpoints): - def __init__(self): - self.stored = [] - +class _ManagedFilesEndpoints(BaseFileEndpoints): async def acreate_file( self, create_file_request, llm_router, target_model_names_list, litellm_parent_otel_span, user_api_key_dict ): @@ -48,10 +46,23 @@ class _FakeManagedFilesHook(BaseFileEndpoints): async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): raise NotImplementedError + +class _FakeManagedFilesHook(_ManagedFilesEndpoints): + def __init__(self): + self.stored = [] + async def store_unified_file_id(self, **kwargs): self.stored.append(kwargs) +class _ForwardingManagedFilesHook(_ManagedFilesEndpoints): + def __init__(self, writer: _FakeManagedFilesHook): + self._writer = writer + + def __getattr__(self, name: str) -> object: + return getattr(self._writer, name) + + class _FailingManagedFilesHook(_FakeManagedFilesHook): async def store_unified_file_id(self, **kwargs): raise RuntimeError("db down") @@ -142,6 +153,24 @@ async def test_upload_with_target_model_names_and_hook_stores_unified_id(monkeyp } +@pytest.mark.asyncio +async def test_upload_stores_the_unified_id_through_a_hook_that_forwards_its_writer(monkeypatch): + backend: Final = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend) + writer: Final = _FakeManagedFilesHook() + + file_object: Final = await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=_ForwardingManagedFilesHook(writer)), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert [stored["file_id"] for stored in writer.stored] == [file_object.id] + + @pytest.mark.asyncio async def test_upload_hands_the_prisma_client_to_the_storage_backend_factory(monkeypatch: pytest.MonkeyPatch): backend = _RecordingStorageBackend() diff --git a/tests/unit/proxy/proxy_server/test_openapi_customization.py b/tests/unit/proxy/proxy_server/test_openapi_customization.py index 6c588976efc..0eacf333960 100644 --- a/tests/unit/proxy/proxy_server/test_openapi_customization.py +++ b/tests/unit/proxy/proxy_server/test_openapi_customization.py @@ -36,6 +36,7 @@ from litellm.proxy.proxy_server import ( get_openapi_schema, mount_swagger_ui, ) +from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder from .conftest import normalize @@ -209,6 +210,17 @@ def test_inject_websocket_stubs_into_openapi_schema_adds_stub(): } +def test_inject_websocket_stubs_into_openapi_schema_documents_query_params_that_forward_their_fields(): + query_param: Final = AttributeForwarder(SimpleNamespace(name="session_id", required=True)) + route: Final = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=SimpleNamespace(query_params=[query_param])) + + result: Final = _inject_websocket_stubs_into_openapi_schema({"paths": {}}, [route]) + + assert result["paths"]["/ws/chat"]["get"]["parameters"] == [ + {"name": "session_id", "in": "query", "required": True, "schema": {"type": "string"}} + ] + + def test_inject_websocket_stubs_into_openapi_schema_does_not_overwrite_existing_get(): # Existing GET on the same path must not be replaced by the stub. existing_get = {"summary": "real http get", "operationId": "real_get"} From 4710c3b30e5c1459dbb11a060ad733ee06f06c18 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:34:32 -0700 Subject: [PATCH 05/43] fix(router): route decisions calls only to deployments whose provider supports decisions (#48175) * fix(router): route decisions calls only to deployments whose provider supports decisions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): account for indirect route helper coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): pin the chat leg so the mixed-group check is deterministic Weights pick the Bedrock leg for chat instead of relying on simple-shuffle landing on both legs. The route filter now passes the candidate list through unchanged in type, so the pre-existing selector lines lose their pyright ignores Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 91 +++++++++++- .../router_code_coverage.py | 3 + tests/unit/test_router/test_router.py | 133 +++++++++++++++++- 3 files changed, 220 insertions(+), 7 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 40c5c8afa6c..25591a435f9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -438,6 +438,7 @@ from litellm.utils import ( EmbeddingResponse, ModelInfoLookupMiss, ModelResponse, + ProviderConfigManager, Rules, TextCompletionResponse, function_setup, @@ -567,6 +568,13 @@ CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_b _PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) _video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) +_routed_call_type: Final[ContextVar[str | None]] = ContextVar("routed_call_type", default=None) +_ROUTE_CAPABILITY_RESOLVERS: Final[Mapping[str, Callable[[str, LlmProviders], object | None]]] = MappingProxyType( + { + "adecisions": ProviderConfigManager.get_provider_decisions_config, + "decisions": ProviderConfigManager.get_provider_decisions_config, + } +) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -823,6 +831,7 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: _RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) +_RouteDeploymentsT: Final = TypeVar("_RouteDeploymentsT", bound=Sequence[Mapping[str, object]] | Mapping[str, object]) @dataclass(frozen=True, slots=True) @@ -6193,7 +6202,7 @@ class Router: try: parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) try: - with self._video_call_pin(handler_name, kwargs): + with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, request_kwargs=kwargs, @@ -6921,6 +6930,15 @@ class Router: finally: _video_call_deployment_id.reset(token) + @staticmethod + @contextlib.contextmanager + def _route_capability_pin(handler_name: str | None) -> Generator[None]: + token: Final = _routed_call_type.set(handler_name) + try: + yield + finally: + _routed_call_type.reset(token) + @staticmethod def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]: """ @@ -6954,7 +6972,7 @@ class Router: kwargs=kwargs, metadata_variable_name=metadata_variable_name, ) - with self._video_call_pin(handler_name, kwargs): + with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): deployment: Final = self.get_available_deployment( model=model, messages=kwargs.get("messages", None), @@ -14209,6 +14227,52 @@ class Router: ) return False + @staticmethod + def _deployment_serves_route( + deployment: DeploymentTypedDict, + group: str, + resolver: Callable[[str, LlmProviders], object | None], + ) -> bool: + if "litellm_params" not in deployment: + return True + + try: + model, resolved_custom_llm_provider = Router._capability_model_and_provider( + deployment["litellm_params"], group + ) + provider: Final = next( + (member for member in LlmProviders if member.value == resolved_custom_llm_provider), + None, + ) + return provider is not None and resolver(model, provider) is not None + except Exception as e: # noqa: BLE001 # best-effort capability filter + verbose_router_logger.debug( + "litellm.router.py::_deployment_serves_route: keeping deployment for model=%s. Got - %s", group, e + ) + return True + + @staticmethod + def _without_deployments_that_cannot_serve_route( + healthy_deployments: _RouteDeploymentsT, group: str + ) -> _RouteDeploymentsT: + handler_name: Final = _routed_call_type.get() + resolver: Final = _ROUTE_CAPABILITY_RESOLVERS.get(handler_name) if handler_name is not None else None + if resolver is None or isinstance(healthy_deployments, Mapping): + return healthy_deployments + + filtered_deployments: Final = [ + deployment + for deployment in healthy_deployments + if Router._deployment_serves_route( + cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape + group, + resolver, + ) + ] + if not filtered_deployments: + return healthy_deployments + return cast(_RouteDeploymentsT, filtered_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 @@ -14845,6 +14909,13 @@ class Router: if verbose_router_logger.isEnabledFor(logging.DEBUG): verbose_router_logger.debug("healthy_deployments after web search filter: %s", healthy_deployments) + healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows + self._without_deployments_that_cannot_serve_route( + healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type + model, + ) + ) + if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: raise litellm.ServiceUnavailableError( @@ -14859,10 +14930,12 @@ class Router: ) # Health-check-based filtering (before cooldown) - healthy_deployments = await self._async_filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, + healthy_deployments = ( # rebind-ok: filtered candidates enter health checks + await self._async_filter_health_check_unhealthy_deployments( + healthy_deployments=healthy_deployments, + parent_otel_span=parent_otel_span, + health_check_probe=health_check_probe, + ) ) routing_read_batch: Final = RoutingReadBatch.active() @@ -16113,6 +16186,12 @@ class Router: specific_deployment=specific_deployment, request_kwargs=request_kwargs, ) + healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows + self._without_deployments_that_cannot_serve_route( + healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type + model, + ) + ) strategy, strategy_selector = self._get_routing_context(model, request_kwargs) routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index af2f386a362..b5e36055354 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -138,6 +138,9 @@ ignored_function_names = [ "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py "_capability_model_and_provider", # Tested through the authenticating-provider pre-call check test in test_router.py + "_route_capability_pin", # Exercised through Decisions routing tests in test_router.py + "_deployment_serves_route", # Exercised through mixed-group Decisions tests in test_router.py + "_without_deployments_that_cannot_serve_route", # Exercised through Decisions routing tests in test_router.py "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index f06f9e08f7e..2288d4311a9 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -10,7 +10,7 @@ import os import sys import threading import warnings -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timedelta, timezone from pathlib import Path from types import MappingProxyType, SimpleNamespace @@ -108,6 +108,137 @@ if TYPE_CHECKING: from litellm.integrations.otel.logger import OpenTelemetryV2 +_ROUTER_DECISIONS_RESPONSE: Final[Mapping[str, object]] = { + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "damaged", "probability": 0.9}], + "usage": { + "input_tokens": 1, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2, + }, +} + + +_OPENAI_DECISIONS_PARAMS: Final[Mapping[str, object]] = { + "model": "openai/gpt-6-luna", + "api_base": "https://api.openai.com/v1", + "api_key": "k", +} +_BEDROCK_DECISIONS_PARAMS: Final[Mapping[str, object]] = { + "model": "bedrock/us.openai.gpt-5.6-terra", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", +} + + +def _decisions_leg(leg_id: str, params: Mapping[str, object], weight: int | None = None) -> dict[str, object]: + return { + "model_name": "decision-group", + "litellm_params": params if weight is None else {**params, "weight": weight}, + "model_info": {"id": leg_id}, + } + + +_OPENAI_DECISIONS_LEG: Final = _decisions_leg("openai-leg", _OPENAI_DECISIONS_PARAMS) +_BEDROCK_DECISIONS_LEG: Final = _decisions_leg("bedrock-leg", _BEDROCK_DECISIONS_PARAMS) + + +def _mixed_decisions_group_router(*legs: dict[str, object]) -> Router: + return Router(model_list=list(legs), routing_strategy="simple-shuffle") + + +@pytest.fixture +def _router_httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_decisions_on_a_mixed_group_never_routes_to_a_leg_without_decisions_support( + respx_mock: respx.MockRouter, + _router_httpx_transport: None, +) -> None: + router: Final = _mixed_decisions_group_router(_OPENAI_DECISIONS_LEG, _BEDROCK_DECISIONS_LEG) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_ROUTER_DECISIONS_RESPONSE) + + for _ in range(20): + await router.adecisions( + model="decision-group", + input=[{"role": "user", "content": "The package arrived with a broken screen."}], + questions=[ + {"type": "predicate", "name": "damaged", "instructions": "Does the customer report a damaged item?"} + ], + ) + + assert route.call_count == 20 + + +def test_decisions_on_a_mixed_group_never_routes_to_a_leg_without_decisions_support_sync( + respx_mock: respx.MockRouter, +) -> None: + router: Final = _mixed_decisions_group_router(_OPENAI_DECISIONS_LEG, _BEDROCK_DECISIONS_LEG) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_ROUTER_DECISIONS_RESPONSE) + + for _ in range(20): + router.decisions( + model="decision-group", + input=[{"role": "user", "content": "The package arrived with a broken screen."}], + questions=[ + {"type": "predicate", "name": "damaged", "instructions": "Does the customer report a damaged item?"} + ], + ) + + assert route.call_count == 20 + + +@pytest.mark.asyncio +async def test_decisions_on_a_group_with_no_capable_leg_still_returns_the_provider_error() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "bedrock-only", + "litellm_params": { + "model": "bedrock/us.openai.gpt-5.6-terra", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + }, + "model_info": {"id": "bedrock-leg"}, + } + ] + ) + + with pytest.raises(litellm.BadRequestError, match="Unknown Decisions provider 'bedrock'"): + await router.adecisions( + model="bedrock-only", + input=[{"role": "user", "content": "The package arrived with a broken screen."}], + questions=[ + {"type": "predicate", "name": "damaged", "instructions": "Does the customer report a damaged item?"} + ], + ) + + +@pytest.mark.asyncio +async def test_chat_on_a_mixed_group_still_reaches_the_leg_without_decisions_support() -> None: + router: Final = _mixed_decisions_group_router( + _decisions_leg("openai-leg", _OPENAI_DECISIONS_PARAMS, weight=0), + _decisions_leg("bedrock-leg", _BEDROCK_DECISIONS_PARAMS, weight=1), + ) + + response: Final = await router.acompletion( + model="decision-group", + messages=[{"role": "user", "content": "Say hi"}], + mock_response="hi", + ) + + assert response._hidden_params["model_id"] == "bedrock-leg" + + def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): # initialize a real Router (env‑vars can be empty) router = litellm.Router( From 08ef4f40b184e03bb3a27594d83a355aded5466e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 24 Oct 2026 12:20:13 -0700 Subject: [PATCH 06/43] feat(ui): configure the default Usage time range (#48207) * feat(ui): configure the default Usage time range * fix(ui): retain the selected calendar preset during input sync * fix(ui): preserve fallback dates when applying across DST --- .../proxy_setting_endpoints.py | 7 ++ .../test_proxy_setting_endpoints.py | 102 ++++++++++++++++++ ui/litellm-dashboard/AGENTS.md | 2 +- .../components/UsagePageView.test.tsx | 55 ++++++++++ .../_components/components/UsagePageView.tsx | 22 ++-- .../usage/_components/usageUrlState.test.ts | 79 +++++++++++++- .../usage/_components/usageUrlState.ts | 11 +- .../src/app/(dashboard)/usage/page.test.tsx | 51 +++++++++ .../src/app/(dashboard)/usage/page.tsx | 24 ++++- .../UISettings/UISettings.test.tsx | 30 ++++++ .../AdminSettings/UISettings/UISettings.tsx | 47 ++++++++ .../shared/advanced_date_picker.test.tsx | 26 +++++ .../shared/advanced_date_picker.tsx | 62 ++--------- .../components/shared/date_range_presets.ts | 56 ++++++++++ 14 files changed, 508 insertions(+), 66 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/page.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/date_range_presets.ts diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index d51e1167888..e04ddbc75cc 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import ( Annotated, Final, + Literal, NamedTuple, Protocol, cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read @@ -256,6 +257,11 @@ class UISettings(LiteLLMBaseModel): model_config = ConfigDict(extra="allow") + usage_default_time_range: Literal["today", "last_7_days", "last_30_days", "month_to_date", "year_to_date"] = Field( + default="last_7_days", + description="Default date range when opening Usage without an explicit URL range. Calendar ranges use the browser timezone.", + ) + disable_model_add_for_internal_users: bool = Field( default=False, description=( @@ -386,6 +392,7 @@ class UISettingsResponse(SettingsResponse): # Allowlist of UI settings that can be stored ALLOWED_UI_SETTINGS_FIELDS: Final = { + "usage_default_time_range", "disable_model_add_for_internal_users", "disable_team_admin_delete_team_user", "enabled_ui_pages_internal_users", 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 512f58380ea..0d558f206d8 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 @@ -4345,3 +4345,105 @@ class TestMoyaiUrlSetting: await update_ui_settings({"moyai_url": None}, actor) assert persisted["moyai_url"] is None + + +@pytest.mark.parametrize("preset", ["today", "last_7_days", "last_30_days", "month_to_date", "year_to_date"]) +def test_usage_default_range_round_trips_without_overwriting_other_settings( + mock_auth, monkeypatch: pytest.MonkeyPatch, preset: str, +) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.config_resolvers import SettingsStore + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=SimpleNamespace( + ui_settings={"disable_custom_api_keys": True}, + )) + prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server.proxy_config, "settings", SettingsStore("general_settings")) + actor: Final = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + app.dependency_overrides[user_api_key_auth] = lambda: actor + try: + initial: Final = client.get("/get/ui_settings") + assert initial.json()["values"]["usage_default_time_range"] == "last_7_days" + response: Final = client.patch("/update/ui_settings", json={"usage_default_time_range": preset}) + assert response.status_code == 200, response.text + stored: Final = json.loads(prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["update"]["ui_settings"]) + assert stored == {"disable_custom_api_keys": True, "usage_default_time_range": preset} + prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace(ui_settings=stored) + fetched: Final = client.get("/get/ui_settings") + assert fetched.json()["values"]["usage_default_time_range"] == preset + assert fetched.json()["source"]["usage_default_time_range"] == "db" + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize("submitted", ["invalid", None, 7]) +def test_usage_default_range_rejects_invalid_values_without_writing( + mock_auth, monkeypatch: pytest.MonkeyPatch, submitted: object, +) -> None: + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + actor: Final = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + app.dependency_overrides[user_api_key_auth] = lambda: actor + try: + response: Final = client.patch("/update/ui_settings", json={"usage_default_time_range": submitted}) + assert response.status_code == 422, response.text + prisma.db.litellm_uisettings.upsert.assert_not_called() + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) +def test_usage_default_range_cannot_be_changed_by_non_admins( + mock_auth, monkeypatch: pytest.MonkeyPatch, role: LitellmUserRoles, +) -> None: + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + actor: Final = UserAPIKeyAuth(user_id="test-user", user_role=role) + app.dependency_overrides[user_api_key_auth] = lambda: actor + try: + response: Final = client.patch("/update/ui_settings", json={"usage_default_time_range": "month_to_date"}) + assert response.status_code == 403, response.text + prisma.db.litellm_uisettings.upsert.assert_not_called() + finally: + app.dependency_overrides.clear() + + +def test_usage_default_range_reports_config_precedence(mock_auth, monkeypatch: pytest.MonkeyPatch) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy import proxy_server + from litellm.proxy.config_resolvers import SettingsStore + + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=SimpleNamespace( + ui_settings={"usage_default_time_range": "today"}, + )) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + settings: Final = SettingsStore("general_settings") + settings.load_yaml({"usage_default_time_range": "month_to_date"}) + monkeypatch.setattr(proxy_server.proxy_config, "settings", settings) + response: Final = client.get("/get/ui_settings") + assert response.status_code == 200, response.text + assert response.json()["values"]["usage_default_time_range"] == "month_to_date" + assert response.json()["source"]["usage_default_time_range"] == "config" diff --git a/ui/litellm-dashboard/AGENTS.md b/ui/litellm-dashboard/AGENTS.md index b6fe5219022..bbf2e49aad7 100644 --- a/ui/litellm-dashboard/AGENTS.md +++ b/ui/litellm-dashboard/AGENTS.md @@ -30,7 +30,7 @@ Type tests are `*.test-d.ts` files run by the `types` vitest project (`npm run t -# This is NOT the Next.js you know +## This is NOT the Next.js you know This version has breaking changes — APIs, conventions, and file structure may all differ from your training data. Read the relevant guide in `node_modules/next/dist/docs/` (resolved from this file's directory; in monorepos the `next` package may not be visible from the repo root) before writing any code. Heed deprecation notices. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index 24e9936a1d0..adb6d8791a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -1914,6 +1914,61 @@ describe("UsagePage", () => { ); }); + it("queries the configured month and preserves later manual selections", async () => { + const now = new Date(2026, 9, 24, 12); + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(now); + try { + renderWithProviders(); + await waitFor(() => + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith( + expect.objectContaining({ + startTime: new Date(2026, 9, 1), + endTime: new Date(2026, 9, 24, 23, 59, 59, 999), + }), + ), + ); + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1); + fireEvent.click(screen.getByTestId("pick-a-different-range")); + await waitFor(() => + expect(mockUserDailyActivityAggregatedCall).toHaveBeenLastCalledWith( + expect.objectContaining({ + startTime: new Date(2024, 0, 1), + endTime: new Date(2024, 0, 8, 23, 59, 59, 999), + }), + ), + ); + } finally { + vi.useRealTimers(); + } + }); + + it.each(["?range=today", "?from=2025-03-01&to=2025-03-05"])( + "preserves explicit dates over configured defaults for %s", + async (searchParams) => { + vi.useFakeTimers({ toFake: ["Date"] }); + vi.setSystemTime(new Date(2026, 9, 24, 12)); + try { + const view = renderWithProviders(, { + searchParams, + }); + const startTime = searchParams.includes("range=today") ? new Date(2026, 9, 24) : new Date(2025, 2, 1); + const endTime = searchParams.includes("range=today") + ? new Date(2026, 9, 24, 23, 59, 59, 999) + : new Date(2025, 2, 5, 23, 59, 59, 999); + await waitFor(() => + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledWith( + expect.objectContaining({ startTime, endTime }), + ), + ); + view.rerender(); + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }, + ); + it("fetches the custom range in the URL and writes a newly picked one", async () => { const { lastQuery } = renderAt("?from=2025-03-01&to=2025-03-05"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 3ff758a3c18..d7d3b502b6f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -82,6 +82,7 @@ import { clientDateRangePatch, clientPanelPatch, dateRangeFromParams, + defaultUsageDateRange, dateRangePatch, entitySelectionPatch, selectedEntitiesFromParams, @@ -108,23 +109,24 @@ interface UsagePageProps { teams: Team[]; organizations: Organization[]; teamsLoading?: boolean; + defaultTimeRange?: string; } -const UsagePage: React.FC = ({ teams, organizations, teamsLoading = false }) => { +const UsagePage: React.FC = ({ teams, organizations, teamsLoading = false, defaultTimeRange }) => { const { accessToken, userRole, userId: userID, premiumUser, isLoading: authorizationLoading } = useAuthorized(); const [gatewayActivityData, setGatewayActivityData] = useState(null); - // Create initial dates outside of state to prevent recreation - const initialFromDate = useMemo(() => new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), []); - const initialToDate = useMemo(() => new Date(), []); + const defaultRange = useMemo(() => defaultUsageDateRange(defaultTimeRange), [defaultTimeRange]); const [urlParams, setUrlParams] = useQueryStates(usageUrlParsers); const { range, from, to } = urlParams; const dateValue = useMemo( - () => dateRangeFromParams({ range, from, to }) ?? { from: initialFromDate, to: initialToDate }, - [range, from, to, initialFromDate, initialToDate], + () => dateRangeFromParams({ range, from, to }) ?? defaultRange.value, + [range, from, to, defaultRange], ); + const presetShortLabel = range ?? (from === null && to === null ? defaultRange.preset : null); + const { client_range, client_from, client_to } = urlParams; const clientDateValue = useMemo( () => clientDateRangeFromParams({ client_range, client_from, client_to }) ?? dateValue, @@ -481,7 +483,11 @@ const UsagePage: React.FC = ({ teams, organizations, teamsLoadin )}
- + - - - )} - {models.isSuccess && models.data.length === 0 && ( - - No chat models are available for your account. - - )} - + + + Could not load models.{" "} + + + ); } From 5a8ab74d10984c977d0defe18a21470868d9ff5b Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 24 Oct 2026 12:30:01 -0700 Subject: [PATCH 09/43] fix(chat): remove hosted telemetry on chat submissions (#48225) --- .../src/app/chat/layout.integration.test.tsx | 5 - .../src/app/chat/layout.test.tsx | 15 +-- .../app/chat/page.images.integration.test.tsx | 28 ++--- .../app/chat/page.mcp.integration.test.tsx | 1 - ui/litellm-dashboard/src/app/chat/page.tsx | 4 - .../src/lib/http/chatTelemetry.test.ts | 43 ------- .../src/lib/http/chatTelemetry.ts | 16 --- .../src/lib/telemetry/chat.test.ts | 117 ------------------ .../src/lib/telemetry/chat.ts | 20 --- 9 files changed, 8 insertions(+), 241 deletions(-) delete mode 100644 ui/litellm-dashboard/src/lib/http/chatTelemetry.test.ts delete mode 100644 ui/litellm-dashboard/src/lib/http/chatTelemetry.ts delete mode 100644 ui/litellm-dashboard/src/lib/telemetry/chat.test.ts delete mode 100644 ui/litellm-dashboard/src/lib/telemetry/chat.ts diff --git a/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx index 12acf9192aa..a7189ea6cd7 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx @@ -2,11 +2,8 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, testQueryClient } from "@/../tests/test-utils"; import { setServerRootPath } from "@/lib/serverRootPath"; -import { recordChatQuery } from "@/lib/telemetry/chat"; import ChatLayout from "./layout"; -vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery: vi.fn(async () => {}) })); - const { router, route } = vi.hoisted(() => ({ router: { push: vi.fn(), replace: vi.fn() }, route: { pathname: "/ui/chat" }, @@ -94,7 +91,6 @@ describe("Chat layout product navigation", () => { ); expect(screen.getByRole("heading", { name: "Chat content" })).toBeVisible(); - expect(recordChatQuery).not.toHaveBeenCalled(); await user.click(screen.getByRole("button", { name: "Switch workspace, current: Chat" })); expect(await screen.findByRole("menuitem", { name: "Observability" })).toBeInTheDocument(); expect(localStorage.getItem("litellm_plugin_mode")).toBe(plugin.name); @@ -103,7 +99,6 @@ describe("Chat layout product navigation", () => { expect(localStorage.getItem("litellm_plugin_mode")).toBe("ai-gateway"); expect(assign).toHaveBeenCalledWith(gateway); - expect(recordChatQuery).not.toHaveBeenCalled(); }, ); diff --git a/ui/litellm-dashboard/src/app/chat/layout.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.test.tsx index 0b0cb96ddc9..a741a7725dd 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.test.tsx @@ -3,7 +3,7 @@ import { act, fireEvent, screen } from "@testing-library/react"; import { renderWithProviders as render } from "@/../tests/test-utils"; import ChatLayout from "./layout"; -const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, recordChatQuery, state } = vi.hoisted(() => { +const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = vi.hoisted(() => { const state = { enableChatUI: false, isUISettingsLoading: false, @@ -14,7 +14,6 @@ const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, recordCha return { state, mockReplace: vi.fn(), - recordChatQuery: vi.fn(async () => {}), mockUiHref: vi.fn((segment: string) => `/mocked-ui/${segment}`), mockUseAuthorized: vi.fn(() => ({ accessToken: "token-123", @@ -38,7 +37,6 @@ vi.mock("next/navigation", () => ({ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: mockUseAuthorized })); vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: mockUseUISettings })); vi.mock("@/utils/uiHref", () => ({ uiHref: mockUiHref })); -vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery })); vi.mock("@/components/navbar", () => ({ default: () =>
})); vi.mock("@/contexts/ThemeContext", () => ({ ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, @@ -66,7 +64,6 @@ describe("ChatLayout", () => { state.userId = "user-1"; mockReplace.mockClear(); mockUiHref.mockClear(); - recordChatQuery.mockClear(); vi.restoreAllMocks(); vi.unstubAllGlobals(); }); @@ -81,7 +78,6 @@ describe("ChatLayout", () => { expect(screen.getByTestId("chat-shell")).toBeInTheDocument(); expect(screen.getByTestId("page-content")).toBeInTheDocument(); expect(mockReplace).not.toHaveBeenCalled(); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should redirect without recording a query when enable_chat_ui is off", () => { @@ -93,7 +89,6 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(mockReplace).toHaveBeenCalledWith("/mocked-ui/"); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should wait for UI settings without recording a query", () => { @@ -105,7 +100,6 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(mockReplace).not.toHaveBeenCalled(); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should wait for the signed-in user's identity without recording a query", () => { @@ -118,7 +112,6 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should hide Chat without recording a query when the signed-in user's identity cannot be confirmed", () => { @@ -131,7 +124,6 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should not record a query or fetch telemetry when opening, focusing, or revisiting Chat", () => { @@ -140,11 +132,9 @@ describe("ChatLayout", () => { const view = render(Chat content); act(() => fireEvent.focus(window)); - expect(recordChatQuery).not.toHaveBeenCalled(); visibility.mockReturnValue("visible"); act(() => fireEvent(document, new Event("visibilitychange"))); - expect(recordChatQuery).not.toHaveBeenCalled(); view.rerender(Another Chat page); act(() => fireEvent.focus(window)); @@ -155,14 +145,12 @@ describe("ChatLayout", () => { fireEvent.focus(window); fireEvent(document, new Event("visibilitychange")); }); - expect(recordChatQuery).not.toHaveBeenCalled(); expect(fetch).not.toHaveBeenCalled(); }); it("should not record a query or fetch telemetry when the signed-in account changes", () => { state.enableChatUI = true; const view = render(Chat content); - expect(recordChatQuery).not.toHaveBeenCalled(); state.userId = "user-2"; view.rerender(Chat content); @@ -172,7 +160,6 @@ describe("ChatLayout", () => { state.isAuthorized = false; view.rerender(Chat content); act(() => fireEvent.focus(window)); - expect(recordChatQuery).not.toHaveBeenCalled(); expect(fetch).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx index f7bbf59c0cb..7e4401fc3ab 100644 --- a/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx @@ -3,11 +3,8 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, testQueryClient } from "@/../tests/test-utils"; import { ChatShellProvider } from "@/contexts/ChatShellContext"; -import { recordChatQuery } from "@/lib/telemetry/chat"; import ChatConversationPage from "./page"; -vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery: vi.fn(async () => {}) })); - vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn(), replace: vi.fn() }), })); @@ -67,10 +64,10 @@ async function chooseModel(user: ReturnType, current: st } beforeEach(() => { + vi.stubEnv("NODE_ENV", "production"); localStorage.clear(); testQueryClient.clear(); requests.length = 0; - vi.mocked(recordChatQuery).mockClear(); models = [ { model_group: "research-chat", mode: "chat", providers: ["openai"] }, { model_group: "studio-image", mode: "image_generation", providers: ["openai"] }, @@ -109,6 +106,11 @@ beforeEach(() => { }); afterEach(() => { + for (const [input] of vi.mocked(fetch).mock.calls) { + const url = input instanceof Request ? input.url : String(input); + expect(new URL(url, window.location.origin).origin).toBe(window.location.origin); + } + vi.unstubAllEnvs(); vi.unstubAllGlobals(); localStorage.clear(); }); @@ -121,10 +123,9 @@ describe("chat image workflows", () => { await userEvent.setup().click(screen.getByRole("button", { name: "research-chat", exact: true })); expect(await screen.findByRole("button", { name: /^studio-image/ })).toBeInTheDocument(); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); - expect(recordChatQuery).not.toHaveBeenCalled(); }); - it("should upload a document as a file part, preserve it on edit, and keep contents out of metrics", async () => { + it("should upload a document as a file part, preserve it on edit", async () => { const user = userEvent.setup(); renderChat(); await screen.findByRole("button", { name: "research-chat" }); @@ -145,7 +146,6 @@ describe("chat image workflows", () => { "data:application/vnd.openxmlformats-officedocument.wordprocessingml.document;base64,cHJpdmF0ZSBmaXh0dXJl", }; expect(requests[0].body).toMatchObject({ input: [{ role: "user", content: [filePart] }] }); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); await user.click(screen.getByRole("button", { name: "Edit message" })); act(() => fireEvent.change(screen.getByRole("textbox", { name: "Edit message" }), { @@ -175,7 +175,6 @@ describe("chat image workflows", () => { await user.click(await screen.findByRole("button", { name: "Remove note-0.txt" })); expect(screen.queryByText("note-0.txt")).not.toBeInTheDocument(); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should restore the selected chat model after switching modes and reopening the page", async () => { @@ -201,7 +200,6 @@ describe("chat image workflows", () => { expect(screen.queryByRole("img", { name: "diagram.png" })).not.toBeInTheDocument(); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); expect(requests).toHaveLength(0); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should send the image and preserve it when editing and resending the message", async () => { @@ -212,7 +210,6 @@ describe("chat image workflows", () => { await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "Explain this diagram"); await user.click(screen.getByRole("button", { name: "Send" })); await screen.findByText("I can see the uploaded image."); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); expect(requests[0].body).toMatchObject({ model: "research-chat", input: [ @@ -246,7 +243,6 @@ describe("chat image workflows", () => { }); expect(requests[1].body).not.toHaveProperty("previous_response_id"); expect(screen.queryByText("Explain this diagram", { selector: "p" })).not.toBeInTheDocument(); - expect(recordChatQuery).toHaveBeenCalledTimes(2); }); it("should preserve a chat turn when editing it while an image model is selected", async () => { @@ -267,7 +263,6 @@ describe("chat image workflows", () => { expect(screen.getByRole("heading", { name: "Explain the experiment" })).toBeInTheDocument(); expect(screen.getByText("I can see the uploaded image.")).toBeInTheDocument(); expect(requests).toHaveLength(1); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); }); it("should continue an image conversation using the completed response context", async () => { @@ -288,7 +283,6 @@ describe("chat image workflows", () => { }); it("should generate and render an image through the selected Model Hub model", async () => { - vi.mocked(recordChatQuery).mockReturnValueOnce(new Promise(() => {})); const user = userEvent.setup(); renderChat(); await screen.findByRole("button", { name: "research-chat" }); @@ -301,7 +295,6 @@ describe("chat image workflows", () => { expect(requests).toEqual([ { pathname: "/images/generations", body: { model: "studio-image", prompt: "A blue research lab" } }, ]); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); }); it("should list no image models when the account has none permitted", async () => { @@ -327,7 +320,6 @@ describe("chat image workflows", () => { await user.click(screen.getByRole("button", { name: "Generate image" })); expect(await screen.findByText(/Image access is disabled for this key/)).toBeInTheDocument(); expect(screen.queryByRole("img", { name: "Generated image 1" })).not.toBeInTheDocument(); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); }); it("should abort image generation when the user stops the request", async () => { @@ -346,12 +338,10 @@ describe("chat image workflows", () => { fireEvent.keyDown(screen.getByRole("textbox", { name: "Message", exact: true }), { key: "Enter" }); }); expect(requests).toHaveLength(1); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); await user.click(screen.getByRole("button", { name: "Stop request" })); expect(generationSignal?.aborted).toBe(true); expect(await screen.findByText("[stopped]")).toBeInTheDocument(); expect(screen.queryByRole("img", { name: "Generated image 1" })).not.toBeInTheDocument(); - expect(recordChatQuery).toHaveBeenCalledTimes(1); }); it("should reject unsupported pasted files without issuing a request", async () => { @@ -364,7 +354,6 @@ describe("chat image workflows", () => { }); expect(await screen.findByRole("alert")).toHaveTextContent("Choose a PNG, JPEG, WebP, or GIF image."); expect(requests).toHaveLength(0); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should not count a blank submission as a query", async () => { @@ -377,7 +366,6 @@ describe("chat image workflows", () => { fireEvent.keyDown(screen.getByRole("textbox", { name: "Message", exact: true }), { key: "Enter" }); }); expect(requests).toHaveLength(0); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it("should count one query for a message that streams text and executes multiple MCP tools", async () => { @@ -390,7 +378,6 @@ describe("chat image workflows", () => { await user.click(serverSwitch); await waitFor(() => expect(screen.getByRole("switch")).toBeChecked()); await user.click(screen.getByRole("button", { name: "Connect tools" })); - expect(recordChatQuery).not.toHaveBeenCalled(); act(() => { fireEvent.change(screen.getByRole("textbox", { name: "Message", exact: true }), { @@ -404,6 +391,5 @@ describe("chat image workflows", () => { expect(requests[0].body).toMatchObject({ tools: [{ type: "mcp", server_label: "research-tools", require_approval: "never" }], }); - expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); }); }); diff --git a/ui/litellm-dashboard/src/app/chat/page.mcp.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.mcp.integration.test.tsx index 4c4612f1ba8..367935a42a2 100644 --- a/ui/litellm-dashboard/src/app/chat/page.mcp.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.mcp.integration.test.tsx @@ -7,7 +7,6 @@ import type { Conversation, MCPEvent } from "@/components/chat/types"; import ChatConversationPage from "./page"; vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn(), replace: vi.fn() }) })); -vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery: vi.fn(async () => {}) })); const userId = "mcp-replay-user"; const storageKey = `litellm_chat_history_v1:${userId}`; diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx index 9b8f0275a07..85fc7133215 100644 --- a/ui/litellm-dashboard/src/app/chat/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.tsx @@ -19,7 +19,6 @@ import { runChatRequest } from "@/components/chat/runChatRequest"; import { readChatAttachment, MAX_CHAT_ATTACHMENTS, CHAT_ATTACHMENT_ACCEPT } from "@/components/chat/attachmentHelpers"; import ChatFileChip from "@/components/chat/ChatFileChip"; import { useQueryState } from "nuqs"; -import { recordChatQuery } from "@/lib/telemetry/chat"; const LOCALSTORAGE_MODEL_KEY = "litellm_chat_selected_model"; @@ -75,7 +74,6 @@ export default function ChatConversationPage() { const router = useRouter(); const { accessToken, - userId, conversations, selectedMCPServers, setSelectedMCPServers, @@ -306,7 +304,6 @@ export default function ChatConversationPage() { onUpdate: (updates: AssistantMessageUpdate) => updateLastAssistantMessage(convId, updates), }; try { - void recordChatQuery(userId); await runChatRequest(request); } finally { if (abortControllerRef.current === requestController) { @@ -324,7 +321,6 @@ export default function ChatConversationPage() { mode, attachments, accessToken, - userId, createConversation, appendMessage, updateLastAssistantMessage, diff --git a/ui/litellm-dashboard/src/lib/http/chatTelemetry.test.ts b/ui/litellm-dashboard/src/lib/http/chatTelemetry.test.ts deleted file mode 100644 index 389406dd8b7..00000000000 --- a/ui/litellm-dashboard/src/lib/http/chatTelemetry.test.ts +++ /dev/null @@ -1,43 +0,0 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; -import { captureChatRequest } from "./chatTelemetry"; - -afterEach(() => { - vi.restoreAllMocks(); - vi.unstubAllGlobals(); -}); - -describe("Chat telemetry transport", () => { - it("should send only pseudonymous user and request IDs without gateway credentials or referrer", async () => { - const send = vi.fn().mockResolvedValue(new Response("{}")); - const timeout = vi.spyOn(AbortSignal, "timeout"); - vi.stubGlobal("fetch", send); - await captureChatRequest("opaque-user", "unique-request"); - - expect(send).toHaveBeenCalledOnce(); - const [url, request] = send.mock.calls[0]; - expect(url).toBe("https://chat-telemetry.litellm.ai/v1/chat/events"); - const expectedRequest = { - method: "POST", - headers: { "Content-Type": "application/json" }, - credentials: "omit", - referrerPolicy: "no-referrer", - keepalive: true, - signal: expect.any(AbortSignal), - body: expect.any(String), - }; - expect(request).toEqual(expectedRequest); - expect(timeout).toHaveBeenCalledExactlyOnceWith(5000); - expect(JSON.parse(request.body)).toEqual({ - user_id: "opaque-user", - request_id: "unique-request", - }); - }); - - it("should reject an unsuccessful ingestion response for the caller to handle", async () => { - vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response("", { status: 503 }))); - await expect(captureChatRequest("opaque-user", "unique-request")).rejects.toThrow( - "Chat telemetry was not accepted", - ); - expect(fetch).toHaveBeenCalledOnce(); - }); -}); diff --git a/ui/litellm-dashboard/src/lib/http/chatTelemetry.ts b/ui/litellm-dashboard/src/lib/http/chatTelemetry.ts deleted file mode 100644 index c2f172250ee..00000000000 --- a/ui/litellm-dashboard/src/lib/http/chatTelemetry.ts +++ /dev/null @@ -1,16 +0,0 @@ -export async function captureChatRequest(userId: string, requestId: string): Promise { - const request: RequestInit = { - method: "POST", - headers: { "Content-Type": "application/json" }, - credentials: "omit", - referrerPolicy: "no-referrer", - keepalive: true, - signal: AbortSignal.timeout(5000), - body: JSON.stringify({ - user_id: userId, - request_id: requestId, - }), - }; - const response = await fetch("https://chat-telemetry.litellm.ai/v1/chat/events", request); - if (!response.ok) throw new Error("Chat telemetry was not accepted"); -} diff --git a/ui/litellm-dashboard/src/lib/telemetry/chat.test.ts b/ui/litellm-dashboard/src/lib/telemetry/chat.test.ts deleted file mode 100644 index 599290ec7e4..00000000000 --- a/ui/litellm-dashboard/src/lib/telemetry/chat.test.ts +++ /dev/null @@ -1,117 +0,0 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; - -const mocks = vi.hoisted(() => ({ capture: vi.fn(), fetch: vi.fn() })); -vi.mock("@/lib/http/chatTelemetry", () => ({ captureChatRequest: mocks.capture })); - -let chat: typeof import("./chat"); - -beforeEach(async () => { - vi.resetModules(); - vi.resetAllMocks(); - vi.stubEnv("NODE_ENV", "production"); - vi.stubGlobal("window", { location: { origin: "https://gateway.example.com" } }); - vi.stubGlobal("fetch", mocks.fetch); - mocks.capture.mockResolvedValue(undefined); - chat = await import("./chat"); -}); - -afterEach(() => { - vi.useRealTimers(); - vi.unstubAllEnvs(); - vi.unstubAllGlobals(); -}); - -describe("Chat request metrics", () => { - it("should do nothing when imported without a submission", () => { - expect(mocks.fetch).not.toHaveBeenCalled(); - expect(mocks.capture).not.toHaveBeenCalled(); - }); - - it("should emit by default with a stable user ID, fresh request IDs, and no backend settings request", async () => { - await chat.recordChatQuery("alice@example.com"); - await chat.recordChatQuery("alice@example.com"); - - expect(mocks.capture).toHaveBeenCalledTimes(2); - const ids = mocks.capture.mock.calls.map(([id]) => id); - expect(ids[0]).toMatch(/^[\da-f]{8}-[\da-f]{4}-5[\da-f]{3}-[\da-f]{4}-[\da-f]{12}$/); - expect(ids[0]).toBe(ids[1]); - const requestIds = mocks.capture.mock.calls.map(([, requestId]) => requestId); - for (const requestId of requestIds) { - expect(requestId).toMatch(/^[\da-f]{8}-[\da-f]{4}-4[\da-f]{3}-[\da-f]{4}-[\da-f]{12}$/); - } - expect(new Set(requestIds).size).toBe(2); - expect(mocks.capture.mock.calls.every((args) => args.length === 2)).toBe(true); - expect(JSON.stringify(mocks.capture.mock.calls)).not.toContain("alice@example.com"); - expect(JSON.stringify(mocks.capture.mock.calls)).not.toContain("gateway.example.com"); - expect(mocks.fetch).not.toHaveBeenCalled(); - }); - - it("should use the same user ID across reloads and browsers without stored identifiers", async () => { - await chat.recordChatQuery("alice"); - vi.resetModules(); - chat = await import("./chat"); - vi.stubGlobal("window", { location: { origin: "https://gateway.example.com" } }); - await chat.recordChatQuery("alice"); - expect(mocks.capture.mock.calls[0][0]).toBe(mocks.capture.mock.calls[1][0]); - }); - - it("should distinguish accounts and deployments", async () => { - await chat.recordChatQuery("alice"); - await chat.recordChatQuery("bob"); - vi.stubGlobal("window", { location: { origin: "https://another.example.com" } }); - await chat.recordChatQuery("alice"); - const ids = mocks.capture.mock.calls.map(([id]) => id); - expect(new Set(ids).size).toBe(3); - }); - - it("should preserve user identity across days while leaving event timestamps to the collector", async () => { - vi.useFakeTimers(); - vi.setSystemTime(new Date("2026-10-23T23:59:00Z")); - await chat.recordChatQuery("alice"); - vi.setSystemTime(new Date("2026-10-24T00:01:00Z")); - await chat.recordChatQuery("alice"); - expect(mocks.capture.mock.calls[0][0]).toBe(mocks.capture.mock.calls[1][0]); - expect(mocks.capture.mock.calls[0][1]).not.toBe(mocks.capture.mock.calls[1][1]); - expect(JSON.stringify(mocks.capture.mock.calls)).not.toContain("2026-10-"); - }); - - it("should not collect from development, unauthenticated, or server rendering", async () => { - vi.stubEnv("NODE_ENV", "development"); - await chat.recordChatQuery("alice"); - vi.stubEnv("NODE_ENV", "production"); - await chat.recordChatQuery(""); - vi.stubGlobal("window", undefined); - await chat.recordChatQuery("alice"); - expect(mocks.fetch).not.toHaveBeenCalled(); - expect(mocks.capture).not.toHaveBeenCalled(); - }); - - it("should count requests on HTTP deployments without secure-context APIs or browser storage", async () => { - vi.stubGlobal("window", { location: { origin: "http://gateway.internal:4000" } }); - vi.stubGlobal("crypto", { getRandomValues: crypto.getRandomValues.bind(crypto) }); - vi.stubGlobal("localStorage", undefined); - await chat.recordChatQuery("alice"); - expect(mocks.capture).toHaveBeenCalledOnce(); - }); - - it("should separate deployments mounted under different root paths", async () => { - const { setServerRootPath } = await import("@/lib/serverRootPath"); - setServerRootPath("/team-a/"); - await chat.recordChatQuery("alice"); - setServerRootPath("team-a"); - await chat.recordChatQuery("alice"); - setServerRootPath("/team-b"); - await chat.recordChatQuery("alice"); - const ids = mocks.capture.mock.calls.map(([id]) => id); - expect(ids[0]).toBe(ids[1]); - expect(ids[0]).not.toBe(ids[2]); - }); - - it("should swallow transport failures without retrying requests", async () => { - mocks.capture.mockRejectedValueOnce(new Error("blocked")); - await expect(chat.recordChatQuery("alice")).resolves.toBeUndefined(); - expect(mocks.capture).toHaveBeenCalledTimes(1); - await chat.recordChatQuery("alice"); - expect(mocks.capture).toHaveBeenCalledTimes(2); - }); -}); diff --git a/ui/litellm-dashboard/src/lib/telemetry/chat.ts b/ui/litellm-dashboard/src/lib/telemetry/chat.ts deleted file mode 100644 index 706c49527fd..00000000000 --- a/ui/litellm-dashboard/src/lib/telemetry/chat.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { v4 as uuidv4, v5 as uuidv5 } from "uuid"; -import { serverRootPath } from "@/lib/serverRootPath"; -import { captureChatRequest } from "@/lib/http/chatTelemetry"; - -// Called only by Chat's accepted-submit handler. Daily unique users = DAU; -// event count = total requests. Page visits never count as activity. -export async function recordChatQuery(userId: string): Promise { - // Match docs: development previews never contribute to production metrics. - if (process.env.NODE_ENV !== "production" || typeof window === "undefined" || !userId) return; - try { - // Stable across browsers for an account on this deployment. The payload - // excludes its raw account ID, email, and gateway origin. - const root = serverRootPath.replace(/^\/+|\/+$/g, ""); - const identity = JSON.stringify([window.location.origin, root, userId]); - const distinctId = uuidv5(identity, uuidv5.URL); - await captureChatRequest(distinctId, uuidv4()); - } catch { - // Analytics failures must never interrupt a request or trigger a retry. - } -} From e6a32bbeebc35bce73c21a9a1c64cb5f3bc3b2b4 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 12:34:13 -0700 Subject: [PATCH 10/43] fix(proxy): buffer tool discoveries in Redis and write LiteLLM_ToolTable in one statement per drain, on its own lock (#48089) * fix(proxy): buffer tool discoveries in redis and let the leader write LiteLLM_ToolTable in one statement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): declare the redis family on the tool discovery buffer and drop the blind except and cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): satisfy LIT010 and LIT014 in the tool discovery buffer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark the aggregation fold Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): requeue drained tool rows to redis when the write provably did not apply Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): drain tool discoveries on their own lock and scheduler job, not the spend commit * chore(proxy): suppress the untyped apscheduler add_job on the tool discovery job * refactor(proxy): move the LiteLLM_ToolTable bulk upsert into its own module next to the other bulk upserts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): name the LiteLLM_ToolTable bulk write span so the prisma span check resolves it * refactor(proxy): move the tool discovery drain out of DBSpendUpdateWriter into tool_registry_writer * refactor(proxy): share db_write_failure_is_requeue_safe from exception_handler so the spend writer stays untouched --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 + litellm/proxy/_types.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 41 +-- .../redis_update_buffer.py | 60 +++++ .../tool_discovery_queue.py | 61 ++--- litellm/proxy/db/exception_handler.py | 12 + litellm/proxy/db/tool_registry_writer.py | 120 +++++---- litellm/proxy/db/tool_table_bulk_upsert.py | 71 ++++++ litellm/proxy/proxy_server.py | 13 + litellm/proxy/utils.py | 14 ++ .../test_redis_update_buffer.py | 16 +- .../test_tool_discovery_queue.py | 73 +++--- .../proxy/db/test_db_spend_update_writer.py | 35 +-- .../proxy/db/test_tool_registry_writer.py | 234 ++++++++++++++---- .../proxy/db/test_tool_table_bulk_upsert.py | 38 +++ 15 files changed, 553 insertions(+), 238 deletions(-) create mode 100644 litellm/proxy/db/tool_table_bulk_upsert.py create mode 100644 tests/unit/proxy/db/test_tool_table_bulk_upsert.py diff --git a/litellm/constants.py b/litellm/constants.py index c721fb94c42..d1daf2ee2d1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -480,6 +480,7 @@ REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_sp REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" +REDIS_TOOL_DISCOVERY_BUFFER_KEY: Final = "litellm_tool_discovery_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100) REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer" REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000 @@ -1844,6 +1845,7 @@ LITELLM_UI_SESSION_DURATION: Final = os.getenv("LITELLM_UI_SESSION_DURATION", "2 ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME: Final = "db_spend_update_job" +TOOL_DISCOVERY_COMMIT_JOB_NAME: Final = "tool_discovery_commit_job" DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 28d3c08bb31..3a3306dfdc0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5262,6 +5262,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): team_id: str | None # team that triggered discovery key_alias: str | None # human-readable key alias user_agent: str | None # HTTP User-Agent of the caller + call_count: ReadOnly[int] # invocations folded into this row since the previous flush from litellm.models.managed_files import ( # noqa: E402 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b290e0f85b4..24ab09283f1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -37,7 +37,6 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -74,7 +73,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.request_error_tracking import request_error_accumulator from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout @@ -251,16 +250,6 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... -_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) - - -def _spend_commit_failure_is_requeue_safe(e: Exception) -> bool: - if isinstance(e, DB_CONNECTION_ERROR_TYPES): - return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) - sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) - return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES - - _SpendTableName = Literal[ "user_list_transactions", "end_user_list_transactions", @@ -298,7 +287,7 @@ def _spend_tables_left_to_send( _SPEND_TABLE_COMMIT_ORDER[len(committed)] if len(committed) < len(_SPEND_TABLE_COMMIT_ORDER) else None ) dropped: Final[frozenset[_SpendTableName]] = ( - frozenset() if in_flight is None or _spend_commit_failure_is_requeue_safe(failure) else frozenset({in_flight}) + frozenset() if in_flight is None or db_write_failure_is_requeue_safe(failure) else frozenset({in_flight}) ) if dropped and in_flight is not None: spend_log_error( @@ -1635,8 +1624,6 @@ class DBSpendUpdateWriter: window_spend_update_queue=self.window_spend_update_queue, ) - await self._flush_tool_discovery_queue(prisma_client=prisma_client) - # Only commit from redis to db if this pod is the leader if await self.pod_lock_manager.acquire_lock( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, @@ -1781,7 +1768,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: - if not _spend_commit_failure_is_requeue_safe(e): + if not db_write_failure_is_requeue_safe(e): uncommitted.pop("window_spend_update_transactions", None) spend_log_error( "Spend tracking - dropped %d budget window increments: the failed statement may have " @@ -1947,7 +1934,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run - if _spend_commit_failure_is_requeue_safe(e): + if db_write_failure_is_requeue_safe(e): spend_log_error( "Spend tracking - failed to commit budget window spend updates. " "Re-queued %d window increments for retry on next tick. Error: %s", @@ -1965,9 +1952,6 @@ class DBSpendUpdateWriter: exc=e, ) - ################## Tool Registry Upserts ################## - await self._flush_tool_discovery_queue(prisma_client=prisma_client) - async def commit_daily_tag_spend_to_db( self, prisma_client: PrismaClient, @@ -2099,20 +2083,6 @@ class DBSpendUpdateWriter: ) raise - async def _flush_tool_discovery_queue( - self, - prisma_client: PrismaClient, - ) -> None: - """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable.""" - from litellm.proxy.db.tool_registry_writer import batch_upsert_tools - - try: - items: Final = self.tool_discovery_queue.flush() - if items: - await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: - verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) - @staticmethod async def handle_spend_update_failure( e: Exception, @@ -2124,7 +2094,6 @@ class DBSpendUpdateWriter: """Retry a failed spend-update transaction on connection errors, deadlocks or a rollup ``lock_timeout`` (55P03), else re-raise. All three roll the transaction back before any increment applied, so re-sending the same batch cannot double-count.""" - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import raise_failed_update_spend_exception is_retryable = ( @@ -2633,7 +2602,7 @@ class DBSpendUpdateWriter: _mark_daily_spend_commit_started() _mark_daily_spend_commit_finished() except Exception as batch_error: - if _spend_commit_failure_is_requeue_safe(batch_error): + if db_write_failure_is_requeue_safe(batch_error): spend_log_error( "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 66833fd688a..3a4375ac54a 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -9,6 +9,7 @@ import json from collections.abc import Mapping, Sequence from datetime import datetime from functools import reduce +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import TypeAdapter @@ -27,6 +28,7 @@ from litellm.constants import ( REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, + REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) @@ -42,12 +44,17 @@ from litellm.proxy._types import ( DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem, + ToolDiscoveryQueueItem, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, @@ -135,6 +142,9 @@ def _merged_entity_transactions( ) +_TOOL_DISCOVERY_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem]) + + class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database @@ -559,6 +569,56 @@ class RedisUpdateBuffer: str(e), ) + @with_service_target(SPEND_QUEUE_TARGET) + async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None: + """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable.""" + if self.redis_cache is None: + return + items: Final = tool_discovery_queue.flush() + if not items: + return + try: + _ = await self.redis_cache.async_rpush_pipeline( + rpush_list=[ + RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)]) + ] + ) + except (RedisError, OSError, asyncio.TimeoutError) as e: + verbose_proxy_logger.error( + "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s", + len(items), + str(e), + ) + for item in items: + tool_discovery_queue.add_update(item) + + @with_service_target(SPEND_QUEUE_TARGET) + async def restore_tool_discoveries_to_redis(self, items: Sequence[ToolDiscoveryQueueItem]) -> None: + """Put drained tool rows back for the next drain after a write that provably did not apply.""" + if self.redis_cache is None or not items: + return + _ = await self.redis_cache.async_rpush_pipeline( + rpush_list=[ + RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(list(items))]) + ] + ) + + @with_service_target(SPEND_QUEUE_TARGET) + async def get_all_tool_discoveries_from_redis_buffer(self) -> tuple[ToolDiscoveryQueueItem, ...]: + """Drain every worker's tool rows and fold them to one row per tool_name.""" + if self.redis_cache is None: + return () + popped: Final = await self.redis_cache.async_lpop_pipeline( + lpop_list=[ + RedisPipelineLpopOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT) + ] + ) + entries: Final = popped[0] if popped else None + if not entries: + return () + batches: Final = tuple(_TOOL_DISCOVERY_ITEMS.validate_json(entry) for entry in entries) + return aggregate_tool_discoveries(chain.from_iterable(batches)) + @with_service_target(SPEND_QUEUE_TARGET) async def store_spend_logs_in_redis( self, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index d5f6a447131..e9ec040023c 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -1,52 +1,57 @@ """ -In-memory buffer for tool registry upserts. - -Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue -uses set-deduplication: each unique tool_name is only queued once per flush -cycle (~30s). The seen-set is cleared on every flush so that call_count -increments in subsequent cycles rather than stopping after the first flush. +In-memory buffer for tool registry upserts, aggregated by tool_name so one flush carries one row per +tool with the number of times it was invoked since the previous flush. """ +from collections.abc import Iterable from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +_ITEM: Final = TypeAdapter(ToolDiscoveryQueueItem) + + +def _merged(existing: ToolDiscoveryQueueItem | None, item: ToolDiscoveryQueueItem) -> ToolDiscoveryQueueItem: + """The first sighting keeps its origin and caller fields, later sightings only add to call_count.""" + increment: Final = item.get("call_count", 1) + if existing is None: + return _ITEM.validate_python({**item, "call_count": increment}) + return _ITEM.validate_python({**existing, "call_count": existing.get("call_count", 1) + increment}) + + +def aggregate_tool_discoveries(items: Iterable[ToolDiscoveryQueueItem]) -> tuple[ToolDiscoveryQueueItem, ...]: + """Collapse items to one per tool_name, summing call_count, in first-seen order.""" + by_name: Final[dict[str, ToolDiscoveryQueueItem]] = {} # mutable-ok: single pass fold keyed by tool_name + for item in items: + if not (tool_name := item.get("tool_name", "")): + continue + by_name[tool_name] = _merged(by_name.get(tool_name), item) + return tuple(by_name.values()) + class ToolDiscoveryQueue: - """ - In-memory buffer for tool registry upserts. - - Deduplicates by tool_name within each flush cycle: a tool is only queued - once per ~30s batch, so call_count increments once per flush cycle the - tool appears in (not once per invocation, but not once per pod lifetime - either). The seen-set is cleared on flush so subsequent batches can - re-count the same tool. - """ + """In-memory buffer for tool registry upserts, one pending row per tool_name with a running call_count.""" def __init__(self) -> None: - self._seen_tool_names: set[str] = set() - self._pending: list[ToolDiscoveryQueueItem] = [] + self._pending: dict[str, ToolDiscoveryQueueItem] = {} # mutable-ok: folded per add_update, cleared on flush def add_update(self, item: ToolDiscoveryQueueItem) -> None: - """Enqueue a tool discovery item if tool_name has not been seen before.""" + """Count one more invocation of item's tool_name, keeping the first sighting's origin and caller fields.""" tool_name: Final = item.get("tool_name", "") if not tool_name: return - if tool_name in self._seen_tool_names: - verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name) - return - self._seen_tool_names.add(tool_name) - self._pending.append(item) + self._pending[tool_name] = _merged(self._pending.get(tool_name), item) verbose_proxy_logger.debug( - "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + "ToolDiscoveryQueue: tool %s call_count=%s (origin=%s)", tool_name, + self._pending[tool_name].get("call_count"), item.get("origin"), ) def flush(self) -> list[ToolDiscoveryQueueItem]: - """Return and clear all pending items. Resets seen-set so the next - flush cycle can re-count the same tools.""" - items, self._pending = self._pending, [] - self._seen_tool_names.clear() + """Return and clear all pending rows.""" + items, self._pending = list(self._pending.values()), {} return items diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 379e0012ea8..5dfcaa4b2f8 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,7 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, ProxyErrorTypes, ProxyException, ) @@ -637,3 +638,14 @@ async def call_with_db_reconnect_retry( # propagate — repeated reconnect-loops are the watchdog's job, not # this helper's. return await coro_factory() + + +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def db_write_failure_is_requeue_safe(e: Exception) -> bool: + """Whether a failed write provably never applied, so its rows can be sent again without double counting.""" + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 2e3530d7613..f7ec88945a1 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -14,8 +14,17 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.constants import TOOL_DISCOVERY_COMMIT_JOB_NAME from litellm.proxy._types import ToolDiscoveryQueueItem -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ToolRepository @@ -105,58 +114,75 @@ def _row_to_model(row: object) -> LiteLLM_ToolTableRow: async def batch_upsert_tools( prisma_client: "PrismaClient", - items: list[ToolDiscoveryQueueItem], + items: Sequence[ToolDiscoveryQueueItem], ) -> None: - """ - Batch-upsert tool registry rows via Prisma. + """Write a batch of discovered tools to LiteLLM_ToolTable in a single statement. Raises on failure so the + caller decides whether the rows can be retried.""" + rows: Final = aggregate_tool_discoveries(items) + if not rows: + return + sql, params = build_tool_bulk_upsert(rows) + async with db_span("commit_tool_discoveries", "LiteLLM_ToolTable"): + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client + verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(rows)) - On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1. - On conflict: increments call_count; preserves existing policies. - """ + +async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discovery_queue: ToolDiscoveryQueue) -> None: + """One worker's own rows straight to LiteLLM_ToolTable, the path without the Redis buffer.""" + items: Final = tool_discovery_queue.flush() if not items: return try: - data: Final = [item for item in items if item.get("tool_name")] - if not data: - return - now: Final = datetime.now(timezone.utc) - table: Final = _tool_table_actions(prisma_client) - for item in data: - tool_name = item.get("tool_name", "") - origin = item.get("origin") or "user_defined" - created_by = item.get("created_by") or "system" - key_hash = item.get("key_hash") - team_id = item.get("team_id") - key_alias = item.get("key_alias") - user_agent = item.get("user_agent") - await table.upsert( - where={"tool_name": tool_name}, - data={ - "create": { - "tool_id": str(uuid.uuid4()), - "tool_name": tool_name, - "origin": origin, - "input_policy": "untrusted", - "output_policy": "untrusted", - "call_count": 1, - "created_by": created_by, - "updated_by": created_by, - "key_hash": key_hash, - "team_id": team_id, - "key_alias": key_alias, - "user_agent": user_agent, - "last_used_at": now, - }, - "update": { - "call_count": {"increment": 1}, - "updated_at": now, - "last_used_at": now, - }, - }, + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail + verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + + +async def commit_tool_discoveries( + prisma_client: "PrismaClient", + tool_discovery_queue: ToolDiscoveryQueue, + redis_update_buffer: RedisUpdateBuffer, + pod_lock_manager: PodLockManager, + use_redis_buffer: bool, +) -> None: + """Scheduled alongside update_spend on its own lock, so tool table writes never ride the spend commit. + + With the Redis buffer every worker hands its rows to Redis and the holder of the tool lock writes them all in + one statement. Without it each worker writes its own rows.""" + if not use_redis_buffer: + await flush_tool_discovery_queue(prisma_client=prisma_client, tool_discovery_queue=tool_discovery_queue) + return + await redis_update_buffer.store_tool_discoveries_in_redis(tool_discovery_queue=tool_discovery_queue) + if not await pod_lock_manager.acquire_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME): + return + try: + await _commit_tool_discoveries_from_redis(prisma_client=prisma_client, redis_update_buffer=redis_update_buffer) + finally: + await pod_lock_manager.release_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME) + + +async def _commit_tool_discoveries_from_redis( + prisma_client: "PrismaClient", redis_update_buffer: RedisUpdateBuffer +) -> None: + """Tool lock holder only: drain every worker's tool rows from Redis and write them in one statement.""" + items: Final = await redis_update_buffer.get_all_tool_discoveries_from_redis_buffer() + if not items: + return + try: + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # the rows go back to redis or are dropped, the scheduler job must not fail + if db_write_failure_is_requeue_safe(e): + await redis_update_buffer.restore_tool_discoveries_to_redis(items) + verbose_proxy_logger.error( + "Tool discovery - %d tool rows went back to Redis for the next tick. Error: %s", len(items), str(e) ) - verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(data)) - except Exception as e: - verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) _NO_OWNERS: Final[Mapping[str, ToolDiscoveryUser]] = MappingProxyType({}) diff --git a/litellm/proxy/db/tool_table_bulk_upsert.py b/litellm/proxy/db/tool_table_bulk_upsert.py new file mode 100644 index 00000000000..e586042ae7c --- /dev/null +++ b/litellm/proxy/db/tool_table_bulk_upsert.py @@ -0,0 +1,71 @@ +"""One multi-row ``INSERT ... ON CONFLICT ("tool_name") DO UPDATE`` per batch of discovered tools. + +LiteLLM_ToolTable holds one row per tool name the proxy has seen. A batch arrives already folded to one +entry per tool name with a summed call_count, and this statement inserts the names the table has not seen +with the untrusted policies the columns default to, and for the names it has adds the batch's call_count +while leaving the policies and the discovering caller's fields alone. Same shape as daily_spend_bulk_upsert. +""" + +import uuid +from collections.abc import Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.proxy._types import ToolDiscoveryQueueItem + +SqlValue: TypeAlias = str | int | None + +_TOOL_ROW_COLUMNS: Final[tuple[str, ...]] = ( + "tool_id", + "tool_name", + "origin", + "call_count", + "created_by", + "updated_by", + "key_hash", + "team_id", + "key_alias", + "user_agent", +) +_CASTS: Final[Mapping[str, str]] = MappingProxyType({"call_count": "integer"}) + + +def _row_params(item: ToolDiscoveryQueueItem) -> tuple[SqlValue, ...]: + created_by: Final = item.get("created_by") or "system" + return ( + str(uuid.uuid4()), + item.get("tool_name", ""), + item.get("origin") or "user_defined", + item.get("call_count", 1), + created_by, + created_by, + item.get("key_hash"), + item.get("team_id"), + item.get("key_alias"), + item.get("user_agent"), + ) + + +def build_tool_bulk_upsert(items: Sequence[ToolDiscoveryQueueItem]) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one batch of distinct tool_names, plus its positional arguments.""" + width: Final = len(_TOOL_ROW_COLUMNS) + rows: Final = ", ".join( + "(" + + ", ".join( + f"${row_index * width + offset + 1}::{_CASTS.get(column, 'text')}" + for offset, column in enumerate(_TOOL_ROW_COLUMNS) + ) + + ", (NOW() AT TIME ZONE 'UTC'), (NOW() AT TIME ZONE 'UTC'))" + for row_index in range(len(items)) + ) + columns: Final = ", ".join(f'"{column}"' for column in _TOOL_ROW_COLUMNS) + sql: Final = ( + f'INSERT INTO "LiteLLM_ToolTable" ({columns}, "last_used_at", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("tool_name") DO UPDATE SET\n' + ' "call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count",\n' + ' "last_used_at" = EXCLUDED."last_used_at",\n' + ' "updated_at" = EXCLUDED."updated_at"' + ) + return sql, tuple(chain.from_iterable(map(_row_params, items))) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 44f43808dc8..df948b69a67 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11296,6 +11296,19 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + ### TOOL DISCOVERY (own lock, same interval, never rides the spend commit) ### + from litellm.proxy.utils import commit_tool_discoveries + + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + commit_tool_discoveries, + "interval", + seconds=batch_writing_interval, + args=[prisma_client, proxy_logging_obj], + id="commit_tool_discoveries_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ### ## Reduces QPS as there are more tags for a single request tag_spend_update_interval: Final = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6cf274884c7..8212e7cbe74 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -184,6 +184,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db import tool_registry_writer from litellm.proxy.db.create_views import ( create_missing_views, create_view_tolerating_race, @@ -191,6 +192,7 @@ from litellm.proxy.db.create_views import ( ) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_url_settings import ( DatabaseURLSettings, add_missing_query_params, @@ -8182,6 +8184,18 @@ class ProxyUpdateSpend: return False +async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: + """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + await tool_registry_writer.commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=writer.tool_discovery_queue, + redis_update_buffer=writer.redis_update_buffer, + pod_lock_manager=writer.pod_lock_manager, + use_redis_buffer=RedisUpdateBuffer.should_commit_spend_updates_to_redis(), + ) + + async def update_spend( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index f30606ac2f1..24c12e06c9c 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction @@ -769,3 +768,18 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False assert redis.rows == [] + + +@pytest.mark.asyncio +async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(): + from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue + + redis_cache = MagicMock() + redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue) + + assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}] diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index e400ad16e84..0a015255f2e 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -5,9 +5,9 @@ Unit tests for ToolDiscoveryQueue. import pytest - from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, + aggregate_tool_discoveries, ) @@ -16,57 +16,50 @@ def queue(): return ToolDiscoveryQueue() -def test_add_single_tool(queue): +def test_a_tool_seen_once_flushes_with_call_count_one(queue): queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "my_tool" - assert items[0]["origin"] == "user_defined" + assert queue.flush() == [{"tool_name": "my_tool", "origin": "user_defined", "call_count": 1}] -def test_deduplication_same_name(queue): - """Adding the same tool_name twice should only keep the first.""" - queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) - queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["origin"] == "mcp_server" # first wins +def test_repeated_sightings_fold_into_one_row_counting_every_call(queue): + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob"}) + queue.add_update({"tool_name": "tool_a"}) + assert queue.flush() == [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice", "call_count": 3}] -def test_deduplication_different_names(queue): +def test_distinct_tools_each_get_their_own_row(queue): queue.add_update({"tool_name": "tool_a"}) queue.add_update({"tool_name": "tool_b"}) - items = queue.flush() - assert len(items) == 2 - names = {i["tool_name"] for i in items} - assert names == {"tool_a", "tool_b"} - - -def test_flush_clears_pending(queue): - queue.add_update({"tool_name": "tool_x"}) - items1 = queue.flush() - assert len(items1) == 1 - items2 = queue.flush() - assert len(items2) == 0 - - -def test_seen_names_reset_after_flush(queue): - """Seen-set is cleared on flush so the same tool can re-enter the next cycle.""" queue.add_update({"tool_name": "tool_a"}) - queue.flush() - queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "tool_a" + assert queue.flush() == [ + {"tool_name": "tool_a", "call_count": 2}, + {"tool_name": "tool_b", "call_count": 1}, + ] + + +def test_flush_clears_pending_and_the_next_cycle_counts_from_zero(queue): + queue.add_update({"tool_name": "tool_x"}) + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 2}] + assert queue.flush() == [] + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 1}] def test_empty_tool_name_ignored(queue): queue.add_update({"tool_name": ""}) queue.add_update({"tool_name": None}) # type: ignore[arg-type] - items = queue.flush() - assert len(items) == 0 + assert queue.flush() == [] -def test_flush_returns_list(queue): - result = queue.flush() - assert isinstance(result, list) +def test_aggregate_sums_call_counts_across_workers_flushes(): + worker_a = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 3}] + worker_b = [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 2}, + {"tool_name": "search", "call_count": 1}, + ] + assert aggregate_tool_discoveries([*worker_a, *worker_b]) == ( + {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, + {"tool_name": "search", "call_count": 1}, + ) 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 9b2a96bbc58..6071569b806 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -19,7 +19,7 @@ from redis.exceptions import DataError import litellm from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes -from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_UPDATE_BUFFER_KEY, SPEND_ROLLUP_LOCK_TIMEOUT_MS +from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, @@ -29,7 +29,6 @@ from litellm.proxy.db.db_spend_update_writer import ( _SpendTableName, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( @@ -2514,31 +2513,6 @@ class _InMemoryRedis: return [self.lists.pop(op["key"], None) for op in lpop_list] -@pytest.mark.asyncio -@pytest.mark.parametrize("lock_held_by_another_pod", [False, True]) -async def test_commit_with_redis_writes_discovered_tools_on_leader_and_non_leader(lock_held_by_another_pod: bool): - redis = _InMemoryRedis() - if lock_held_by_another_pod: - assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) - db_writer = DBSpendUpdateWriter(redis_cache=redis) - db_writer.pod_lock_manager.redis_cache = redis - db_writer.tool_discovery_queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) - prisma_client = MagicMock() - prisma_client.db.litellm_tooltable.upsert = AsyncMock() - - await db_writer._commit_spend_updates_to_db_with_redis( - prisma_client=prisma_client, - n_retry_times=0, - proxy_logging_obj=MagicMock(), - ) - - upserted = [c.kwargs["where"] for c in prisma_client.db.litellm_tooltable.upsert.await_args_list] - assert upserted == [{"tool_name": "get_weather"}] - assert (REDIS_UPDATE_BUFFER_KEY in redis.lists) is lock_held_by_another_pod, ( - f"spend buffer left in redis={list(redis.lists)}, only the lock holder may drain it" - ) - - @pytest.mark.asyncio async def test_commit_with_redis_requeues_all_on_db_failure(): """ @@ -3384,7 +3358,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() db.query_raw = AsyncMock(side_effect=Exception("connection reset")) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), @@ -3392,7 +3365,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ proxy_logging_obj=MagicMock(), ) - db_writer._flush_tool_discovery_queue.assert_called_once() requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() assert requeued == (transaction,) @@ -3468,7 +3440,6 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov db_writer = DBSpendUpdateWriter() await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3494,7 +3465,6 @@ async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} ) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3522,7 +3492,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"} await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend") - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3533,7 +3502,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == [] (team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend") assert _row_values(team_upsert, "team_id") == ["team-1"] - db_writer._flush_tool_discovery_queue.assert_called_once() db.failing_table = None await db_writer._commit_spend_updates_to_db_without_redis_buffer( @@ -3760,7 +3728,6 @@ async def test_failed_window_spend_commit_is_requeued_only_when_the_rows_are_pro ) await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFailureDB(failure) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index c9df665741d..76898b0eeab 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -3,22 +3,28 @@ Unit tests for tool_registry_writer.py — uses a mock prisma client that exposes litellm_tooltable.upsert / find_many / find_unique. """ +import json from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock import pytest -from prisma.errors import PrismaError - +from prisma.errors import PrismaError, RawQueryError +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_TOOL_DISCOVERY_BUFFER_KEY, TOOL_DISCOVERY_COMMIT_JOB_NAME +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue from litellm.proxy.db.tool_registry_writer import ( ToolPolicyRegistry, batch_upsert_tools, + commit_tool_discoveries, get_tool, get_tool_policy_registry, get_tools_by_names, list_tools, update_tool_policy, ) +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation def _mock_row(**kwargs): @@ -62,6 +68,7 @@ def _make_prisma( prisma = MagicMock() prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) + prisma.db.execute_raw = AsyncMock(return_value=1) prisma.db.litellm_tooltable.find_many = AsyncMock( return_value=find_many_rows if find_many_rows is not None else [] ) @@ -73,28 +80,32 @@ def _make_prisma( return prisma +def _statement(prisma) -> tuple[str, tuple[object, ...]]: + prisma.db.execute_raw.assert_awaited_once() + sql, *params = prisma.db.execute_raw.call_args.args + return sql, tuple(params) + + @pytest.mark.asyncio -async def test_batch_upsert_tools_calls_upsert(): +async def test_batch_upsert_tools_inserts_a_new_tool_with_the_batch_call_count_and_adds_it_on_conflict(): prisma = _make_prisma() - items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] - await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_awaited_once() - call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs - assert call_kw["where"] == {"tool_name": "tool_a"} - assert call_kw["data"]["create"]["tool_name"] == "tool_a" - assert call_kw["data"]["create"]["origin"] == "mcp_server" - assert call_kw["data"]["create"]["input_policy"] == "untrusted" - assert call_kw["data"]["create"]["output_policy"] == "untrusted" - assert call_kw["data"]["create"]["call_count"] == 1 - assert call_kw["data"]["update"]["call_count"] == {"increment": 1} - assert "updated_at" in call_kw["data"]["update"] + await batch_upsert_tools( + prisma, [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}] + ) + sql, params = _statement(prisma) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"input_policy"' not in sql, "policies come from the table defaults and are never overwritten" + assert params[1:] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + prisma.db.litellm_tooltable.upsert.assert_not_awaited() @pytest.mark.asyncio async def test_batch_upsert_tools_empty_list(): prisma = _make_prisma() await batch_upsert_tools(prisma, []) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio @@ -102,21 +113,22 @@ async def test_batch_upsert_tools_skips_empty_names(): prisma = _make_prisma() items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item] await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio -async def test_batch_upsert_multiple_tools_calls_upsert_per_tool(): +async def test_batch_upsert_many_tools_is_one_statement_with_one_row_per_distinct_tool(): prisma = _make_prisma() items = [ - {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, - {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 2}, + {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice", "call_count": 1}, + {"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob", "call_count": 5}, ] await batch_upsert_tools(prisma, items) - assert prisma.db.litellm_tooltable.upsert.await_count == 2 - calls = prisma.db.litellm_tooltable.upsert.call_args_list - assert calls[0].kwargs["where"]["tool_name"] == "tool_a" - assert calls[1].kwargs["where"]["tool_name"] == "tool_b" + sql, params = _statement(prisma) + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert params[1:10] == ("tool_a", "mcp_server", 7, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", None, None, None, None) @pytest.mark.asyncio @@ -264,20 +276,14 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): assert call_kw["where"] == {"tool_name": "my_tool"} assert call_kw["data"]["update"]["input_policy"] == "blocked" assert call_kw["data"]["update"]["updated_by"] == "admin" - prisma.db.litellm_tooltable.find_unique.assert_awaited_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio async def test_get_tools_by_names_returns_policy_map(): rows = [ - _mock_row( - tool_name="tool_a", input_policy="trusted", output_policy="untrusted" - ), - _mock_row( - tool_name="tool_b", input_policy="blocked", output_policy="untrusted" - ), + _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"), + _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"), ] prisma = _make_prisma(find_many_rows=rows) result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) @@ -285,9 +291,7 @@ async def test_get_tools_by_names_returns_policy_map(): "tool_a": ("trusted", "untrusted"), "tool_b": ("blocked", "untrusted"), } - prisma.db.litellm_tooltable.find_many.assert_awaited_once_with( - where={"tool_name": {"in": ["tool_a", "tool_b"]}} - ) + prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(where={"tool_name": {"in": ["tool_a", "tool_b"]}}) @pytest.mark.asyncio @@ -382,12 +386,8 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): return [] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock( - side_effect=_flaky_find_many - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(side_effect=_flaky_find_many) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(return_value=[]) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -397,10 +397,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_tools_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_tools_lookup_failure" assert registry.is_initialized() @@ -421,9 +418,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - side_effect=_flaky_perms_find_many - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(side_effect=_flaky_perms_find_many) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -433,7 +428,142 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() assert len(perms_invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_perms_lookup_failure" + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_perms_lookup_failure" + + +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lock_held_by_another_pod", + [None, DB_SPEND_UPDATE_JOB_NAME, TOOL_DISCOVERY_COMMIT_JOB_NAME], + ids=["free", "spend_lock_held", "tool_lock_held"], +) +async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_from_the_tool_lock_holder_only( + lock_held_by_another_pod: str | None, +): + redis = _InMemoryRedis() + if lock_held_by_another_pod is not None: + assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=lock_held_by_another_pod) + other_worker = ToolDiscoveryQueue() + other_worker.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + other_worker.add_update({"tool_name": "search", "origin": "user_defined"}) + await RedisUpdateBuffer(redis_cache=redis).store_tool_discoveries_in_redis(other_worker) + + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + prisma_client.db.litellm_tooltable.upsert = AsyncMock() + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + assert queue.flush() == [], "every worker hands its tools to redis each tick" + prisma_client.db.litellm_tooltable.upsert.assert_not_awaited() + if lock_held_by_another_pod == TOOL_DISCOVERY_COMMIT_JOB_NAME: + prisma_client.db.execute_raw.assert_not_awaited() + assert len(redis.lists[REDIS_TOOL_DISCOVERY_BUFFER_KEY]) == 2, "non-holders leave their tools for the holder" + return + assert REDIS_TOOL_DISCOVERY_BUFFER_KEY not in redis.lists, "the spend lock must not gate the tool drain" + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') and 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 3, "search": 1}, f"one row per tool summed across workers, got {params}" + + +@pytest.mark.asyncio +async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statement_and_never_touches_redis(): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=False, + ) + + assert redis.lists == {}, "without the buffer flag nothing goes through redis" + assert queue.flush() == [] + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 2, "search": 1}, f"got {params}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (Exception("connection reset"), True), + (_postgres_rejection("23502"), False), + ], + ids=["connection_error_goes_back_to_redis", "rejected_row_is_dropped"], +) +async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + prisma_client.db.execute_raw.assert_awaited_once() + left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] + expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] + assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" diff --git a/tests/unit/proxy/db/test_tool_table_bulk_upsert.py b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py new file mode 100644 index 00000000000..d1e79030f9d --- /dev/null +++ b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py @@ -0,0 +1,38 @@ +"""Tests for the single-statement LiteLLM_ToolTable upsert (LIT-9997).""" + +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert + + +def test_one_row_per_item_and_the_conflict_update_only_moves_the_counters(): + sql, params = build_tool_bulk_upsert( + [ + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}, + { + "tool_name": "tool_b", + "origin": "user_defined", + "created_by": "alice", + "call_count": 1, + "key_hash": "hash", + "team_id": "team", + "key_alias": "alias", + "user_agent": "agent", + }, + ] + ) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert "$4::integer" in sql and "$14::integer" in sql + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"last_used_at" = EXCLUDED."last_used_at"' in sql + assert '"input_policy"' not in sql and '"output_policy"' not in sql, "policies stay on the table defaults" + assert '"origin" =' not in sql and '"key_hash" =' not in sql, "an existing row keeps who discovered it" + assert len(params) == 20 + assert params[1:10] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", "hash", "team", "alias", "agent") + assert params[0] != params[10], "each row gets its own tool_id" + + +def test_defaults_fill_origin_created_by_and_call_count(): + _, params = build_tool_bulk_upsert([{"tool_name": "tool_a"}]) + assert params[1:6] == ("tool_a", "user_defined", 1, "system", "system") From d0671f75a4bf87e9fcaea38498be005aa3931b12 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:39:27 -0700 Subject: [PATCH 11/43] refactor(typing): widen delete_nested_value to Mapping and drop two unused Span aliases (#48001) * refactor(typing): replace pass-through Any with object in 13 annotations * docs(client): say dict[str, object] in the models get docstring * refactor(typing): drop the unused Span aliases in base_cache and levo * refactor(typing): take a Mapping in delete_nested_value and cover the fresh copy * fix(typing): deep-copy a read-only Mapping in delete_nested_value through a dict --- litellm/caching/base_cache.py | 9 +-------- litellm/integrations/levo/levo.py | 4 ---- .../litellm_core_utils/dot_notation_indexing.py | 4 ++-- litellm/proxy/client/models.py | 2 +- .../test_dot_notation_indexing.py | 14 ++++++++++++++ 5 files changed, 18 insertions(+), 15 deletions(-) diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py index fb91d13a8c9..4dc894a64ba 100644 --- a/litellm/caching/base_cache.py +++ b/litellm/caching/base_cache.py @@ -9,14 +9,7 @@ Has 4 methods: """ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final - -if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - Span = _Span -else: - Span = Any +from typing import Final class BaseCache(ABC): diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py index 6273f39e09c..c1852e9bfe4 100644 --- a/litellm/integrations/levo/levo.py +++ b/litellm/integrations/levo/levo.py @@ -4,8 +4,6 @@ from typing import TYPE_CHECKING, Any, Final from litellm.integrations.opentelemetry import OpenTelemetry if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - from litellm.integrations.opentelemetry import ( OpenTelemetryConfig as _OpenTelemetryConfig, ) @@ -13,11 +11,9 @@ if TYPE_CHECKING: Protocol = _Protocol OpenTelemetryConfig = _OpenTelemetryConfig - Span = _Span else: Protocol = Any OpenTelemetryConfig = Any - Span = Any class LevoConfig: diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 09720fa413c..818cf1b9695 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -181,7 +181,7 @@ def _delete_nested_value_custom( def delete_nested_value( - data: dict[str, object], + data: Mapping[str, object], path: str, depth: int = 0, max_depth: int = 20, @@ -214,7 +214,7 @@ def delete_nested_value( """ import copy - result: Final = copy.deepcopy(data) + result: Final = copy.deepcopy(data if isinstance(data, dict) else dict(data)) try: # Parse path into segments diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 7342974c0cc..1361a3e6364 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -168,7 +168,7 @@ class ModelsManagementClient: return_request (bool): If True, returns the prepared request object instead of executing it Returns: - Union[Dict[str, Any], requests.Request]: Either the model information from the server or + dict[str, object] | requests.Request: Either the model information from the server or a prepared request object if return_request is True Raises: diff --git a/tests/unit/litellm_core_utils/test_dot_notation_indexing.py b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py index ae529a71009..7c670f1dac9 100644 --- a/tests/unit/litellm_core_utils/test_dot_notation_indexing.py +++ b/tests/unit/litellm_core_utils/test_dot_notation_indexing.py @@ -2,6 +2,9 @@ Tests for litellm.litellm_core_utils.dot_notation_indexing module. """ +from types import MappingProxyType +from typing import Final + import pytest from litellm.litellm_core_utils.dot_notation_indexing import ( @@ -125,3 +128,14 @@ class TestDeleteNestedValue: data = {"items": [{"a": 1, "b": 2}, {"a": 3, "b": 4}]} result = delete_nested_value(data, "items[0].b") assert result == {"items": [{"a": 1}, {"a": 3, "b": 4}]} + + def test_delete_returns_a_fresh_dict_and_leaves_a_narrowly_typed_input_untouched(self): + payload: Final[dict[str, list[str]]] = {"tags": ["a", "b"], "secret": ["s"]} + result = delete_nested_value(payload, "secret") + assert result == {"tags": ["a", "b"]} + assert payload == {"tags": ["a", "b"], "secret": ["s"]} + assert result["tags"] is not payload["tags"] + + def test_delete_accepts_a_read_only_mapping(self): + payload: Final = MappingProxyType({"tags": ["a"], "secret": ["s"]}) + assert delete_nested_value(payload, "secret") == {"tags": ["a"]} From 953ee1914a3d70602008d18674ba6516c4ed1f33 Mon Sep 17 00:00:00 2001 From: hgn <35766626+hgnnmnn@users.noreply.github.com> Date: Sat, 24 Oct 2026 21:43:38 +0200 Subject: [PATCH 12/43] fix(ui): wire up Download CSV Template button in bulk invite (#48199) * fix(ui): wire up Download CSV Template button in bulk invite * refactor(ui): reuse shared downloadBlob helper for bulk invite CSV template --------- Co-authored-by: ryan-crabbe-berri --- .../bulk_create_users_button.test.tsx | 30 +++++++++++++++++++ .../components/bulk_create_users_button.tsx | 12 +++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx index e1faf8be1df..51d7521c909 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx @@ -30,6 +30,36 @@ describe("BulkCreateUsersButton", () => { expect(screen.getByText("+ Bulk Invite Users")).toBeInTheDocument(); }); + it("downloads a CSV template with the columns the upload accepts", async () => { + const createObjectURL = vi.fn((_blob: Blob) => "blob:template"); + const revokeObjectURL = vi.fn(); + const { createObjectURL: originalCreate, revokeObjectURL: originalRevoke } = URL; + URL.createObjectURL = createObjectURL; + URL.revokeObjectURL = revokeObjectURL; + const anchorClick = vi.spyOn(HTMLAnchorElement.prototype, "click").mockImplementation(() => {}); + + const user = await openUploadStep(); + await user.click(screen.getByRole("button", { name: /download csv template/i })); + + expect(anchorClick).toHaveBeenCalledOnce(); + const anchor = anchorClick.mock.contexts[0] as HTMLAnchorElement; + expect(anchor.download).toBe("bulk_users_template.csv"); + expect(anchor.href).toBe("blob:template"); + + const csv = await new Promise((resolve) => { + const reader = new FileReader(); + reader.onload = () => resolve(reader.result as string); + reader.readAsText(createObjectURL.mock.calls[0][0]); + }); + const [headerRow] = csv.split(/\r?\n/); + expect(headerRow).toBe("user_email,user_role,teams,max_budget,budget_duration,models"); + expect(revokeObjectURL).toHaveBeenCalledWith("blob:template"); + + anchorClick.mockRestore(); + URL.createObjectURL = originalCreate; + URL.revokeObjectURL = originalRevoke; + }); + it("parses a CSV chosen through the file input", async () => { await openUploadStep(); diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx index 8b823ca2f73..840cc1af00e 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx @@ -16,6 +16,7 @@ import Papa from "papaparse"; import { CheckCircleIcon, XCircleIcon, ExclamationIcon } from "@heroicons/react/outline"; import { CopyToClipboard } from "react-copy-to-clipboard"; import { toast } from "@/lib/toast"; +import { downloadBlob } from "@/components/EntityUsageExport/utils"; interface BulkCreateUsersProps { accessToken: string; @@ -87,6 +88,15 @@ const BulkCreateUsersButton: React.FC = ({ setBaseUrl(base.toString()); }, [accessToken]); + const downloadTemplate = () => { + const template = [ + ["user_email", "user_role", "teams", "max_budget", "budget_duration", "models"], + ["user@example.com", "internal_user", "team-id-1,team-id-2", "100", "30d", "gpt-3.5-turbo,gpt-4"], + ]; + + downloadBlob(new Blob([Papa.unparse(template)], { type: "text/csv" }), "bulk_users_template.csv"); + }; + const handleFileUpload = (file: File) => { // Reset all error states setParseError(null); @@ -609,7 +619,7 @@ const BulkCreateUsersButton: React.FC = ({
- From 55479821e22f07c64fb5ecdbbc708b980da9ec2d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 24 Oct 2026 13:00:58 -0700 Subject: [PATCH 13/43] feat(router): add opt-in session repinning after successful failover (#48206) * feat(router): add opt-in session repinning after successful failover * fix(router): respect shared session pins and split metadata * fix(router): preserve session routing during Redis failures * test(router): inject the delayed Redis reader --- litellm/caching/affinity_cache.py | 29 +- litellm/caching/redis_cache.py | 6 +- litellm/router.py | 42 +++ .../deployment_affinity_check.py | 105 +++++- .../integration/sdk/test_dual_cache_redis.py | 64 ++++ .../test_session_id_affinity.py | 348 ++++++++++++++++++ 6 files changed, 587 insertions(+), 7 deletions(-) diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2d27574b098..d9a40c44d2b 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -24,6 +24,16 @@ if current == false then end if ARGV[3] then local decoded, stored = pcall(cjson.decode, current) + if ARGV[4] then + local legacy = current + if decoded and type(stored) == 'string' then legacy = stored end + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + if eligible['model_id'] == legacy then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end if decoded and type(stored) == 'table' then for _, eligible in ipairs(cjson.decode(ARGV[3])) do local matches = true @@ -68,6 +78,7 @@ def claim_affinity_pin_in_memory( ttl_seconds: int, *, eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, ) -> object: """No await between read and write, so same-loop claims agree during a Redis outage.""" existing: Final[object] = cache.in_memory_cache.get_cache(cache_key) @@ -75,7 +86,12 @@ def claim_affinity_pin_in_memory( if _legacy_pin_matches(existing, pin_value): set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds) return existing - winner: Final = existing if existing is not None and existing in (eligible_values or ()) else pin_value + eligible_existing: Final = existing in (eligible_values or ()) or ( + preserve_legacy_model_ids + and isinstance(existing, str) + and any("model_id" in value and _legacy_pin_matches(existing, value) for value in (eligible_values or ())) + ) + winner: Final = existing if existing is not None and eligible_existing else pin_value set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) return winner @@ -94,6 +110,7 @@ async def claim_affinity_pin( ttl_seconds: int, *, eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, ) -> object: """Return the authoritative first writer, replacing it only when it becomes ineligible. @@ -112,6 +129,7 @@ async def claim_affinity_pin( if eligible_values is not None else () ), + *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), ) raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here object, await claim_script(keys=(cache_key,), args=args) @@ -124,4 +142,11 @@ async def claim_affinity_pin( return winner except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) - return claim_affinity_pin_in_memory(cache, cache_key, pin_value, ttl_seconds, eligible_values=eligible_values) + return claim_affinity_pin_in_memory( + cache, + cache_key, + pin_value, + ttl_seconds, + eligible_values=eligible_values, + preserve_legacy_model_ids=preserve_legacy_model_ids, + ) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index eb17facfb12..afe6f8a1eab 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1738,7 +1738,9 @@ class RedisCache(BaseCache): return key_value_dict @_redis_circuit_breaker_guard - async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): + async def async_get_cache( + self, key, parent_otel_span: Span | None = None, *, raise_on_error: bool = False, **kwargs + ): from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1781,6 +1783,8 @@ class RedisCache(BaseCache): ) ) print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}") + if raise_on_error: + raise _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard diff --git a/litellm/router.py b/litellm/router.py index 25591a435f9..96102b678c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1247,6 +1247,7 @@ class Router: fallback_access_check: FallbackAccessCheck | None = None, fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, + session_affinity_repin_on_failover: bool = False, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -1283,6 +1284,8 @@ class Router: alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. + session_affinity_repin_on_failover (bool): After a successful replacement response, move a session pin + whose deployment is in cooldown or removed. Async routes only. Defaults to False. Applied at startup. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. @@ -1467,6 +1470,7 @@ class Router: self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds + self.session_affinity_repin_on_failover = session_affinity_repin_on_failover self.model_group_affinity_config = model_group_affinity_config warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) @@ -2862,6 +2866,7 @@ class Router: enable_session_id_affinity=False, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -2908,6 +2913,7 @@ class Router: enable_session_id_affinity=enable_session_id_affinity, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -9300,6 +9306,12 @@ class Router: if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None: return + if self.session_affinity_repin_on_failover: + await self._repin_session_after_success( + _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), + _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), + ) + # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits increment_deployment_successes_for_current_minute( litellm_router_instance=self, @@ -9323,6 +9335,36 @@ class Router: "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) + async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: + if not self.session_affinity_repin_on_failover: + return + try: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import completed_affinity_response + + if not completed_affinity_response(response): + return + callback: Final = next( + (logger for logger in self._optional_loggers() if isinstance(logger, DeploymentAffinityCheck)), + None, + ) + if callback is None: + return + model_group: Final = callback.get_repin_model_group(litellm_params) + if model_group is None or self._is_priority_routing_group(model_group): + return + registered_ids: Final = self.get_candidate_model_ids_for_route( + model=model_group, team_id=get_request_team_id(litellm_params) + ) + cooldown_ids: Final = frozenset( + model_id + for model_id, _ in await self.cooldown_cache.async_get_active_cooldowns( + model_ids=list(registered_ids), parent_otel_span=None + ) + ) + await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) + except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting + verbose_router_logger.debug("Session re-pinning failed: %s", error) + async def increment_deployment_usage_for_response( self, response: object, diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 88fb8876175..8693eaa5151 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -16,6 +16,7 @@ import hashlib from collections.abc import Callable, Mapping, Sequence from typing import Final, cast +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target @@ -38,6 +39,27 @@ class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] +def _affinity_response_fields(response: object) -> Mapping[str, object]: + return ( + TypeAdapter(Mapping[str, object]).validate_python(response) + if isinstance(response, Mapping) + else {name: getattr(response, name, None) for name in ("response", "status", "error")} + ) + + +def completed_affinity_response(response: object) -> bool: + if response is None: + return False + envelope: Final = _affinity_response_fields(response) + if envelope.get("error") is not None or envelope.get("status") not in (None, "completed"): + return False + nested: Final = envelope.get("response") + fields: Final = _affinity_response_fields(nested) if nested is not None else envelope + return ( + fields.get("response") is None and fields.get("error") is None and fields.get("status") in (None, "completed") + ) + + VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( { "deployment_affinity", @@ -86,6 +108,7 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id_affinity: bool = False, model_group_affinity_config: dict[str, list[str]] | None = None, is_priority_group: Callable[[str], bool] | None = None, + session_affinity_repin_on_failover: bool = False, ): super().__init__() self.cache = cache @@ -95,6 +118,7 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} self.is_priority_group = is_priority_group + self.session_affinity_repin_on_failover = session_affinity_repin_on_failover def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -297,7 +321,7 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}" @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_session_id_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): session_id = DeploymentAffinityCheck._get_session_id_from_metadata_dict(metadata=metadata) if session_id is not None: @@ -305,7 +329,7 @@ class DeploymentAffinityCheck(CustomLogger): return None @staticmethod - def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None: + def _get_marker_session_affinity_ttl(request_kwargs: Mapping[str, object]) -> int | None: """TTL from the session-affinity marker the Router stamps at pre-routing time when an auto-router routed this request with session_affinity enabled. Marker presence enables session pinning for this request only; anything that @@ -357,7 +381,7 @@ class DeploymentAffinityCheck(CustomLogger): model: str, healthy_deployments: list, messages: list[AllMessageValues] | None, - request_kwargs: dict | None = None, + request_kwargs: dict[str, object] | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: """ @@ -418,7 +442,15 @@ class DeploymentAffinityCheck(CustomLogger): session_cache_key: Final = self.get_session_affinity_cache_key( model_group=stable_model_map_key, session_id=session_id, user_key=user_key ) - session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key) + session_cache_result: Final = ( + await self._read_shared_session_pin( + session_cache_key, + self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds, + parent_otel_span, + ) + if self.session_affinity_repin_on_failover + else await self.cache.async_get_cache(key=session_cache_key) + ) session_model_id: str | None = None if isinstance(session_cache_result, dict): @@ -482,6 +514,71 @@ class DeploymentAffinityCheck(CustomLogger): ) return [deployment] + async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: + redis_cache: Final = self.cache.redis_cache + if redis_cache is None: + return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key)) + local_before: Final = TypeAdapter(object).validate_python(self.cache.in_memory_cache.get_cache(cache_key)) + try: + shared: Final[object] = await redis_cache.async_get_cache( + key=cache_key, parent_otel_span=parent_otel_span, raise_on_error=True + ) + except Exception: # noqa: BLE001 # Optional affinity falls back locally on Redis admission or read failure + verbose_router_logger.debug("Session pin read failed, using pod-local affinity") + return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True)) + if self.cache.in_memory_cache.get_cache(cache_key) is local_before: + if shared is None: + self.cache.in_memory_cache.delete_cache(cache_key) + else: + self._set_local_pin(cache_key, shared, ttl_seconds) + return shared + + def get_repin_model_group(self, request_kwargs: Mapping[str, object]) -> str | None: + model_group: Final = next( + ( + value + for metadata in self._iter_metadata_dicts(request_kwargs) + if isinstance(value := metadata.get("deployment_model_name"), str) and value + ), + None, + ) + if model_group is None or self._get_session_id_from_request_kwargs(request_kwargs) is None: + return None + if ( + not self._get_effective_flags(model_group)[2] + and self._get_marker_session_affinity_ttl(request_kwargs) is None + ): + return None + return model_group + + async def async_repin_session_after_success( + self, + request_kwargs: Mapping[str, object], + model_group: str, + eligible_deployment_ids: frozenset[str], + ) -> None: + marker_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs) + if not self._get_effective_flags(model_group)[2] and marker_ttl is None: + return + session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) + if session_id is None: + return + model_info: Final = TypeAdapter(Mapping[str, object]).validate_python(request_kwargs.get("model_info") or {}) + model_id: Final = model_info.get("id") + if model_id is None or str(model_id) not in eligible_deployment_ids: + return + cache_key: Final = self.get_session_affinity_cache_key( + model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) + ) + await claim_affinity_pin( + self.cache, + cache_key, + DeploymentAffinityCacheValue(model_id=str(model_id)), + marker_ttl if marker_ttl is not None else self.ttl_seconds, + eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids), + preserve_legacy_model_ids=True, + ) + async def async_pre_call_deployment_hook( self, kwargs: dict[str, object], call_type: CallTypes | None ) -> dict | None: diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index ddd01480d36..3527dbf91ee 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/tests/integration/sdk/test_dual_cache_redis.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest +from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -92,3 +93,66 @@ async def test_sync_and_async_batch_reads_share_one_redis_without_sync_reads_goi assert in_loop_results == [[expected[key]] for key in in_loop_keys] assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]] + + +@pytest.mark.asyncio +async def test_concurrent_session_repin_claims_keep_the_shared_winner_after_recovery() -> None: + redis: Final = _redis_cache() + workers: Final = (DualCache(redis_cache=redis), DualCache(redis_cache=redis)) + key: Final = f"session-repin-{uuid.uuid4().hex}" + choices: Final = ({"model_id": "replacement"}, {"model_id": "other"}) + try: + await redis.async_set_cache(key, {"model_id": "original"}, ttl=60) + for worker in workers: + worker.in_memory_cache.set_cache(key, {"model_id": "original"}, ttl=60) + winners: Final = await asyncio.gather( + *( + claim_affinity_pin(worker, key, choice, 60, eligible_values=choices, preserve_legacy_model_ids=True) + for worker, choice in zip(workers, choices, strict=True) + ) + ) + assert winners[0] == winners[1] and winners[0] in choices + late_original: Final = await claim_affinity_pin( + workers[0], + key, + {"model_id": "original"}, + 60, + eligible_values=({"model_id": "original"}, *choices), + preserve_legacy_model_ids=True, + ) + assert late_original == winners[0] + assert await redis.async_get_cache(key) == winners[0] + finally: + await redis.async_delete_cache(key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored", ["original", '"original"', '{"model_id":"original"}']) +async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unavailable_ones(stored: str) -> None: + redis: Final = _redis_cache() + cache: Final = DualCache(redis_cache=redis) + key: Final = f"session-repin-legacy-{uuid.uuid4().hex}" + replacement: Final = {"model_id": "replacement"} + try: + redis.redis_client.set(redis.check_and_fix_namespace(key), stored, ex=60) + preserved: Final = await claim_affinity_pin( + cache, + key, + replacement, + 60, + eligible_values=({"model_id": "original"}, replacement), + preserve_legacy_model_ids=True, + ) + assert preserved in ("original", {"model_id": "original"}) + moved: Final = await claim_affinity_pin( + cache, + key, + replacement, + 60, + eligible_values=(replacement,), + preserve_legacy_model_ids=True, + ) + assert moved == replacement + assert await redis.async_get_cache(key) == replacement + finally: + await redis.async_delete_cache(key) diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 9bbaed0ae1b..6b4ea74894a 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,20 +1,368 @@ import asyncio import json from typing import Final +from collections.abc import AsyncIterator, Mapping +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection +from redis.asyncio import BlockingConnectionPool, Redis import litellm from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) +def _repinning_router(enabled: bool | None = True) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "repin-group", + "litellm_params": {"model": "openai/test-model", "api_key": "fake-api-key"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("original", "replacement", "other") + ] + + [ + { + "model_name": "unrelated-group", + "litellm_params": {"model": "openai/test-model", "api_key": "fake-api-key"}, + "model_info": {"id": "unrelated"}, + } + ], + optional_pre_call_checks=["session_affinity", "deployment_affinity"], + **({"session_affinity_repin_on_failover": enabled} if enabled is not None else {}), + ) + + +def _repin_params(deployment_id: str, caller: str = "caller-1") -> dict[str, object]: + return { + "model_info": {"id": deployment_id}, + "metadata": { + "model_group": "repin-group", + "deployment_model_name": "repin-group", + "session_id": "repin-session", + "user_api_key_hash": caller, + }, + } + + +def _repin_key(caller: str = "caller-1") -> str: + return DeploymentAffinityCheck.get_session_affinity_cache_key("repin-group", "repin-session", caller) + + +def _cool_down_original(router: litellm.Router) -> None: + router.cooldown_cache.add_deployment_to_cooldown( + model_id="original", original_exception=RuntimeError("test outage"), exception_status=503, cooldown_time=3600 + ) + + +async def _record_replacement_success( + router: litellm.Router, deployment_id: str = "replacement", response: object = None +) -> None: + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + await router.deployment_callback_on_success( + {"litellm_params": _repin_params(deployment_id), "standard_logging_object": {"total_tokens": 0}}, + {"choices": [{"message": {"content": "ok"}}]} if response is None else response, + timestamp, + timestamp, + ) + + +@pytest.mark.parametrize("enabled", [None, False, True], ids=["default", "disabled", "enabled"]) +@pytest.mark.parametrize("unavailable", ["cooldown", "removed", "temporary-filter"]) +async def test_session_repin_requires_opt_in_unavailable_original_and_success( + enabled: bool | None, unavailable: str +) -> None: + router: Final = _repinning_router(enabled) + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + await router.cache.async_set_cache(_repin_key("caller-2"), {"model_id": "original"}) + user_key: Final = DeploymentAffinityCheck.get_affinity_cache_key("repin-group", "caller-1") + await router.cache.async_set_cache(user_key, {"model_id": "original"}) + if unavailable == "cooldown": + _cool_down_original(router) + elif unavailable == "removed": + router.delete_deployment("original") + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + await callback.async_pre_call_deployment_hook(_repin_params("replacement"), None) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + await _record_replacement_success(router) + expected: Final = "replacement" if enabled and unavailable != "temporary-filter" else "original" + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected} + assert await router.cache.async_get_cache(_repin_key("caller-2")) == {"model_id": "original"} + assert await router.cache.async_get_cache(user_key) == {"model_id": "original"} + router.cooldown_cache.cooldown_store.flush_cache() + await _record_replacement_success(router, "other") + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected} + finally: + router.discard() + + +@pytest.mark.parametrize("status", ["failed", "incomplete", "in_progress", "cancelled"]) +@pytest.mark.parametrize("wrapped", [False, True]) +async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + response: Final[Mapping[str, object]] = {"status": status} + await _record_replacement_success(router, response={"response": response} if wrapped else response) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + finally: + router.discard() + + +async def test_concurrent_successful_replacements_keep_first_winner() -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + await asyncio.gather(_record_replacement_success(router), _record_replacement_success(router, "other")) + winner: Final = await router.cache.async_get_cache(_repin_key()) + assert winner in ({"model_id": "replacement"}, {"model_id": "other"}) + router.cooldown_cache.cooldown_store.flush_cache() + await _record_replacement_success(router, "original") + assert await router.cache.async_get_cache(_repin_key()) == winner + finally: + router.discard() + + +async def test_healthy_legacy_string_pin_is_preserved_with_repinning_enabled() -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), "original") + await _record_replacement_success(router) + assert await router.cache.async_get_cache(_repin_key()) == "original" + finally: + router.discard() + + +@pytest.mark.parametrize("route_bucket", ["metadata", "litellm_metadata"]) +async def test_repin_session_after_success_reads_split_metadata(route_bucket: str) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + session_bucket: Final = "litellm_metadata" if route_bucket == "metadata" else "metadata" + await router._repin_session_after_success( + { + "model_info": {"id": "replacement"}, + route_bucket: {"deployment_model_name": "repin-group"}, + session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"}, + }, + {"status": "completed"}, + ) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"} + finally: + router.discard() + + +@pytest.mark.parametrize("enabled", [False, True]) +async def test_session_selection_observes_another_workers_repin_only_when_enabled(enabled: bool) -> None: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + router: Final = _repinning_router(enabled) + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url=f"redis://repin-cache.example:6379/{int(enabled)}") + router.cache.attach_redis_cache(redis) + other_worker: Final = DualCache(redis_cache=redis) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + await other_worker.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + callback: Final = next( + cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck) + ) + selected: Final = await callback.async_filter_deployments( + model="repin-group", + healthy_deployments=router.get_model_list(model_name="repin-group") or [], + messages=None, + request_kwargs=_repin_params("replacement"), + ) + assert [deployment["model_info"]["id"] for deployment in selected] == [ + "replacement" if enabled else "original" + ] + finally: + router.discard() + await pool.aclose() + sync_client.close() + + +@pytest.fixture +async def shared_session_router() -> AsyncIterator[tuple[litellm.Router, FakeServer, RedisCache]]: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + router: Final = _repinning_router() + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url="redis://session-outage.example:6379/0") + router.cache.attach_redis_cache(redis) + yield router, server, redis + finally: + router.discard() + await pool.aclose() + sync_client.close() + + +async def _selected_session_ids(router: litellm.Router) -> tuple[str, ...]: + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + selected: Final = await callback.async_filter_deployments( + model="repin-group", + healthy_deployments=router.get_model_list(model_name="repin-group") or [], + messages=None, + request_kwargs=_repin_params("replacement"), + ) + return tuple(deployment["model_info"]["id"] for deployment in selected) + + +@pytest.mark.parametrize("has_pin", [False, True]) +async def test_redis_read_failures_and_open_breaker_preserve_session_routing( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], has_pin: bool +) -> None: + router, server, redis = shared_session_router + if has_pin: + await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + assert await _selected_session_ids(router) == ("replacement",) + server.connected = False + with patch("time.time", return_value=1790000000.0): + for _ in range(redis._circuit_breaker.failure_threshold + 1): + assert set(await _selected_session_ids(router)) == ( + {"replacement"} if has_pin else {"original", "replacement", "other"} + ) + assert redis._circuit_breaker.is_open() + + +async def test_confirmed_redis_miss_does_not_revive_a_local_pin_during_outage( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + router, server, _ = shared_session_router + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"} + server.connected = False + assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"} + + +async def test_shared_session_pin_resumes_after_a_transient_redis_failure( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + router, server, redis = shared_session_router + await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + server.connected = False + assert await _selected_session_ids(router) == ("original",) + server.connected = True + assert await _selected_session_ids(router) == ("replacement",) + + +@pytest.mark.parametrize("shared_value", [None, {"model_id": "original"}]) +async def test_delayed_shared_pin_read_preserves_a_newer_local_write( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], shared_value: Mapping[str, str] | None +) -> None: + router, server, redis = shared_session_router + if shared_value is not None: + await redis.async_set_cache(_repin_key(), shared_value, ttl=3600) + observed: Final = asyncio.Event() + resume: Final = asyncio.Event() + + class DelayedReader(Redis): + async def get(self, name: str) -> object: + value: Final = await super().get(name) + observed.set() + await resume.wait() + return value + + reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool) + litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader) + pending: Final = asyncio.create_task(_selected_session_ids(router)) + try: + await observed.wait() + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + finally: + resume.set() + await pending + server.connected = False + assert await _selected_session_ids(router) == ("replacement",) + + +@pytest.mark.parametrize( + "response", + [ + {"status": "failed", "response": {"status": "completed"}}, + {"error": {"message": "test failure"}, "response": {"status": "completed"}}, + {"response": {"response": {"status": "completed"}}}, + ], +) +async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: Mapping[str, object]) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + await _record_replacement_success(router, response=response) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + finally: + router.discard() + + +class _ReadCountingCache(InMemoryCache): + def __init__(self) -> None: + super().__init__() + self.read_keys: tuple[str, ...] = () + + def get_cache(self, key: str, **kwargs: object) -> object: + self.read_keys += (key,) + return super().get_cache(key, **kwargs) + + +@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"]) +async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None: + router: Final = _repinning_router(case != "disabled") + reads: Final = _ReadCountingCache() + try: + router.cooldown_cache.cooldown_store.in_memory_cache = reads + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + if case == "group-disabled": + callback.model_group_affinity_config = {"repin-group": []} + await router._repin_session_after_success( + { + "model_info": {"id": "replacement"}, + "metadata": { + "deployment_model_name": "repin-group", + **({"session_id": "repin-session"} if case != "no-session" else {}), + SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session", + }, + }, + {"status": "completed"}, + ) + expected: Final = ( + { + router.cooldown_cache.get_cooldown_cache_key(model_id) + for model_id in ("original", "replacement", "other") + } + if case == "active" + else set() + ) + assert set(reads.read_keys) == expected + finally: + router.discard() + + class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data From 49f4a59cc86ce83ab5d3cd7b4a878412f186e015 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:43:07 -0700 Subject: [PATCH 14/43] fix(router): keep latency routing alive when latency samples are unusable (#48216) * fix(router): keep latency routing alive when latency samples are unusable latency-based-routing ranks deployments by the average of the cached samples and keeps everything within lowest_latency_buffer times the lowest. A negative or NaN sample breaks that invariant, because the minimum no longer satisfies its own bound. The candidate list then comes back empty and random.choice raises IndexError, which turns every request to the group into a 500 until the cache entry expires. Samples come from completion_start_time - start_time, and that difference is negative whenever the two timestamps invert. Production recorded samples around -600s from successful calls. Drop non-finite and negative samples before ranking, and fall back to every candidate if the window is empty anyway, for example with a negative buffer. Co-Authored-By: Claude Code Co-Authored-By: AaronHowell <111272993+AaronHowell@users.noreply.github.com> * test(router): seed latency_normalized in the unusable TTFT fallback test Main reads latency_normalized since #47727, so seeding latency left both deployments ranked at zero and the test picked at random. Co-authored-by: Rick <26716961+Bytechoreographer@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): assert ranked latency after dropping unusable TTFT samples Greptile on #48216: the new tests passed with the sample filter removed because the empty-window fallback alone saved them. They now assert _latency_per_deployment so unfiltered negative or NaN samples fail deterministically. Also adds Final to ttft_samples and drops a comment that restated the code Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): drop Final on a loop-local in latency ranking basedpyright rejects Final on a variable assigned inside a loop, and the type_check_gate counts it as a new reportGeneralTypeIssues error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): type the request kwargs in the unusable TTFT tests Greptile on #48216: bare dict annotations. The tests now hold the metadata dict in a typed local and assert on it directly Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Rick <26716961+Bytechoreographer@users.noreply.github.com> Co-authored-by: Claude Code Co-authored-by: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/lowest_latency.py | 24 +++++-- .../router_strategy/test_lowest_latency.py | 65 +++++++++++++++++++ 2 files changed, 82 insertions(+), 7 deletions(-) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 28ed062f1ec..4e994d3fdc6 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -3,7 +3,7 @@ import random from collections.abc import Sequence from datetime import datetime, timedelta -from math import ceil +from math import ceil, isfinite from typing import TYPE_CHECKING, Any, Final from pydantic import Field @@ -46,6 +46,16 @@ def _percentile_latency(samples: Sequence[float], percentile: float) -> float: return values[index] +def _usable_latency_samples(samples: Sequence[float]) -> tuple[float, ...]: + """Drop samples that cannot be a real elapsed time. + + ``completion_start_time - start_time`` goes negative when the timestamps invert. + A negative or NaN sample pushes ``lowest_latency + buffer`` below every candidate, + leaving the deployment filter empty. + """ + return tuple(sample for sample in samples if isfinite(sample) and sample >= 0) + + def _ttft_seconds(elapsed: timedelta | float) -> float: if isinstance(elapsed, timedelta): return elapsed.total_seconds() @@ -469,16 +479,17 @@ class LowestLatencyLoggingHandler(CustomLogger): item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) + ttft_samples = _usable_latency_samples(item_ttft_latency) use_ttft = ( request_kwargs is not None and request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True - and len(item_ttft_latency) > 0 + and len(ttft_samples) > 0 ) selected_latency = ( - _percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile) + _percentile_latency(ttft_samples, self.routing_args.ttft_percentile) if use_ttft and self.routing_args.ttft_percentile is not None - else _average_latency(item_ttft_latency if use_ttft else item_latency) + else _average_latency(ttft_samples if use_ttft else _usable_latency_samples(item_latency)) ) # -------------- # @@ -512,10 +523,9 @@ class LowestLatencyLoggingHandler(CustomLogger): # Find deployments within buffer of lowest latency buffer: Final = self.routing_args.lowest_latency_buffer * lowest_latency - valid_deployments: Final = [x for x in sorted_deployments if x[1] <= lowest_latency + buffer] + valid_deployments: Final = tuple(x for x in sorted_deployments if x[1] <= lowest_latency + buffer) - # Pick a random deployment from valid deployments - random_valid_deployment: Final = random.choice(valid_deployments) + random_valid_deployment: Final = random.choice(valid_deployments or sorted_deployments) deployment = random_valid_deployment[0] metadata_field: Final = self._select_metadata_field(request_kwargs) if request_kwargs is not None and metadata_field in request_kwargs: diff --git a/tests/unit/router_strategy/test_lowest_latency.py b/tests/unit/router_strategy/test_lowest_latency.py index 9af5def0d56..706ec567710 100644 --- a/tests/unit/router_strategy/test_lowest_latency.py +++ b/tests/unit/router_strategy/test_lowest_latency.py @@ -176,6 +176,8 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): MODEL_GROUP = "gpt-4o-mini" FAST_TTFT_ID = "fast-ttft-short-output" SLOW_TTFT_ID = "slow-ttft-long-output" +FAST_API_BASE = "https://fast.example.test/v1" +SLOW_API_BASE = "https://slow.example.test/v1" STREAMING_DEPLOYMENTS = [ {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}}, {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}}, @@ -407,6 +409,69 @@ async def test_async_get_available_deployments_treats_missing_samples_as_zero_la assert picked["model_info"]["id"] == DEPLOYMENT_ID +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("ttft_samples", "expected_latency"), + [([-602.0], 0.0), ([10.0, 10.0, -40.0], 10.0), ([float("nan"), 2.0], 2.0)], + ids=["negative", "negative_average", "nan"], +) +async def test_unusable_ttft_samples_do_not_leave_the_buffer_window_empty( + ttft_samples: list[float], expected_latency: float +): + """``completion_start_time - start_time`` can be negative or NaN. With a non-zero + ``lowest_latency_buffer`` a negative average pushes ``lowest_latency + buffer`` + below every candidate, so the group must still route instead of raising + IndexError from ``random.choice``, and the bad samples must not reach the ranking.""" + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"lowest_latency_buffer": 5}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={FAST_TTFT_ID: {"time_to_first_token_seconds": ttft_samples}}, + ) + metadata: Final[dict[str, object]] = {} + request_kwargs: Final[dict[str, object]] = {"stream": True, "metadata": metadata} + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=[{"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {"api_base": FAST_API_BASE}}], + request_kwargs=request_kwargs, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + assert metadata["_latency_per_deployment"] == {FAST_API_BASE: expected_latency} + + +@pytest.mark.asyncio +async def test_unusable_ttft_samples_fall_back_to_recorded_latency(): + """A deployment whose TTFT samples are all unusable is ranked by its recorded + latency, so dirty samples cannot make it look like the fastest deployment.""" + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"lowest_latency_buffer": 5}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"latency_normalized": [0.1], "time_to_first_token_seconds": [-600.0]}, + SLOW_TTFT_ID: {"latency_normalized": [5.0], "time_to_first_token_seconds": [-300.0]}, + }, + ) + metadata: Final[dict[str, object]] = {} + request_kwargs: Final[dict[str, object]] = {"stream": True, "metadata": metadata} + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=[ + {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {"api_base": FAST_API_BASE}}, + {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {"api_base": SLOW_API_BASE}}, + ], + request_kwargs=request_kwargs, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + assert metadata["_latency_per_deployment"] == {FAST_API_BASE: 0.1, SLOW_API_BASE: 5.0} + + def _latency_router(routing_strategy_args: dict) -> Router: return Router( model_list=[ From 9b11729712b427372b1caafa7df9af0fc13e6ab8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:44:19 -0700 Subject: [PATCH 15/43] fix(vector_stores): put azure vector store id before the api-version query on retrieve, update and delete (#48224) * fix(vector_stores): put azure vector store id before api-version query Azure OpenAI's get_complete_url returns a url ending in ?api-version=..., and the vector store handlers appended the id with an f-string, producing .../openai/vector_stores?api-version=X/vs_abc and a 404 from Azure. Add BaseVectorStoreConfig.get_vector_store_url, defaulting to the previous behaviour so other providers are unchanged, and override it in AzureOpenAIVectorStoreConfig to keep the query string last. The six retrieve/update/delete handlers now call it. Co-Authored-By: Claude Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(vector_stores): trim comments on get_vector_store_url Drops the inline comment in the Azure override and shortens the base docstring, per AGENTS.md Co-authored-by: Keith-Hardaway <276775292+hardawaykwork@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(vector_stores): capture the Azure retrieve request without a mutable url list Greptile on #48224: the handler test now follows the file's received list pattern with Final bindings Co-authored-by: Keith-Hardaway <276775292+hardawaykwork@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Keith-Hardaway Co-authored-by: Claude Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: joshua Co-authored-by: Keith-Hardaway <276775292+hardawaykwork@users.noreply.github.com> --- .../azure/vector_stores/transformation.py | 7 +++++ .../base_llm/vector_store/transformation.py | 5 ++++ litellm/llms/custom_httpx/llm_http_handler.py | 12 ++++---- ...test_azure_vector_stores_transformation.py | 29 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 30 +++++++++++++++++++ 5 files changed, 77 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure/vector_stores/transformation.py b/litellm/llms/azure/vector_stores/transformation.py index 9510c021c7d..5a2961cd550 100644 --- a/litellm/llms/azure/vector_stores/transformation.py +++ b/litellm/llms/azure/vector_stores/transformation.py @@ -1,3 +1,5 @@ +from typing import Final + from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig from litellm.types.router import GenericLiteLLMParams @@ -15,5 +17,10 @@ class AzureOpenAIVectorStoreConfig(OpenAIVectorStoreConfig): route="/openai/vector_stores", ) + def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> str: + path, _, query = api_base.partition("?") + url: Final = f"{path}/{encoded_vector_store_id}" + return f"{url}?{query}" if query else url + def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: return BaseAzureLLM.base_validate_azure_environment(headers=headers, litellm_params=litellm_params) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 138700b3212..6885183aeba 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -221,6 +221,11 @@ class BaseVectorStoreConfig: raise ValueError("api_base is required") return api_base + def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> str: + """Url of one vector store for retrieve, update and delete. Providers whose + ``get_complete_url`` adds a query string, like Azure, override this.""" + return f"{api_base}/{encoded_vector_store_id}" + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: from ..chat.transformation import BaseLLMException diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 616f1ef559b..1b491f1614d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -10891,7 +10891,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", @@ -10956,7 +10956,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", @@ -11159,7 +11159,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) request_body: Final[dict[str, object]] = dict(vector_store_update_optional_params) metadata: Final = vector_store_update_optional_params.get("metadata") @@ -11238,7 +11238,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) request_body: Final[dict[str, object]] = dict(vector_store_update_optional_params) metadata: Final = vector_store_update_optional_params.get("metadata") @@ -11306,7 +11306,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", @@ -11368,7 +11368,7 @@ class BaseLLMHTTPHandler: ) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}" + url: Final = vector_store_provider_config.get_vector_store_url(api_base, encoded_vector_store_id) logging_obj.pre_call( input="", diff --git a/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py index 59bec08fca6..3713b1cc895 100644 --- a/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py +++ b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py @@ -1,4 +1,7 @@ +from unittest.mock import Mock + from litellm.llms.azure.vector_stores.transformation import AzureOpenAIVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig def test_transform_search_vector_store_request_preserves_azure_query_string(): @@ -18,3 +21,29 @@ def test_transform_search_vector_store_request_preserves_azure_query_string(): ) assert url == "https://x.openai.azure.com/openai/vector_stores/vs_1/search?api-version=2024-10-21" + + +def test_get_vector_store_url_puts_id_before_api_version_query(): + config = AzureOpenAIVectorStoreConfig() + api_base = config.get_complete_url( + api_base="https://x.openai.azure.com", + litellm_params={"api_version": "2025-04-01-preview"}, + ) + + url = config.get_vector_store_url(api_base, "vs_abc") + + assert url == "https://x.openai.azure.com/openai/vector_stores/vs_abc?api-version=2025-04-01-preview" + + +def test_get_vector_store_url_without_query_is_plain_path(): + url = AzureOpenAIVectorStoreConfig().get_vector_store_url( + "https://x.openai.azure.com/openai/vector_stores", "vs_abc" + ) + + assert url == "https://x.openai.azure.com/openai/vector_stores/vs_abc" + + +def test_base_get_vector_store_url_appends_id_for_providers_without_query_string(): + url = BaseVectorStoreConfig.get_vector_store_url(Mock(), "https://api.example.test/v1/vector_stores", "vs_abc") + + assert url == "https://api.example.test/v1/vector_stores/vs_abc" diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 8b32ab2e6b8..bbcd7f70e63 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -24,6 +24,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.azure.vector_stores.transformation import AzureOpenAIVectorStoreConfig from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -5624,3 +5625,32 @@ async def test_messages_opencode_session_header( "stream": stream, } assert received.empty() + + +def test_vector_store_retrieve_handler_sends_id_before_azure_api_version_query() -> None: + received: Final[list[httpx.Request]] = [] + + def answer(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(200, json={"id": "vs_abc", "object": "vector_store"}) + + handler: Final = BaseLLMHTTPHandler() + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(answer))) + + handler.vector_store_retrieve_handler( + vector_store_id="vs_abc", + vector_store_provider_config=AzureOpenAIVectorStoreConfig(), + custom_llm_provider="azure", + litellm_params=GenericLiteLLMParams( + api_base="https://x.openai.azure.com", + api_key="azure-test-key", + api_version="2025-04-01-preview", + ), + logging_obj=Mock(), + client=client, + _is_async=False, + ) + + sent: Final = received[0] + assert str(sent.url) == "https://x.openai.azure.com/openai/vector_stores/vs_abc?api-version=2025-04-01-preview" + assert len(received) == 1 From bcd541918a1c52fbb4dd6e8eb901b3b864a9465c Mon Sep 17 00:00:00 2001 From: tbartik Date: Sat, 24 Oct 2026 22:48:29 +0200 Subject: [PATCH 16/43] fix(anthropic): stop injecting empty-content placeholder into tool-call-only assistant turns (#34822) * fix(anthropic): stop injecting empty-content placeholder into tool-call-only assistant turns _sanitize_empty_text_content replaced empty or whitespace-only text with "[System: Empty message content sanitised to satisfy protocol]" for every user and assistant turn. anthropic_messages_pt calls it unconditionally, independent of modify_params, because an empty text block always makes Anthropic 400. An assistant turn carrying tool_calls and no text is the normal shape an OpenAI-format agent emits, and Anthropic accepts an assistant turn whose content is just tool_use, so there was nothing to substitute there. Substituting anyway put a protocol diagnostic in the content field, which is the assistant's utterance as far as the model and the client are concerned, so it got persisted into conversation history and echoed back on later turns. Empty text is now dropped rather than replaced when the assistant turn already carries substantive content (tool_calls, or non-text blocks such as tool_use). Turns that would otherwise be left with no content at all keep the placeholder, since Anthropic rejects those too. Fixes #24498 * fix(anthropic): only drop empty text when a sibling block survives conversion _carries_non_text_content treated any non-text block as substantive, but the assistant list-content converter only keeps thinking (when signable), server_tool_use and *_tool_result blocks. For an assistant turn holding empty text next to a block it discards, such as an image_url or a signature-less thinking block, dropping the empty text left nothing behind and the whole turn disappeared from the request. The predicate now mirrors what the conversion actually keeps, so those turns retain the placeholder and survive. This also improves on the previous behaviour, where such a turn was dropped outright. * test(anthropic): cover malformed content blocks in the survivor predicate _survives_assistant_conversion had a separate non-dict guard that no test reached, since _is_empty_text_block already rejects non-dict blocks before the predicate runs on them in practice. Folding the check into the type lookup keeps the same behaviour for bare strings and blocks without a type key, with no unreachable branch. Extends the sibling-block test with a block missing its type key. * test(anthropic): set modify_params through monkeypatch in sanitization tests The test-quality gate flags direct writes to litellm module globals. Swap the class setup/teardown for an autouse monkeypatch fixture, drop writes that just repeat its default, and route the False cases through monkeypatch --------- Co-authored-by: tbartik <22518863+tbartik@users.noreply.github.com> Co-authored-by: ryan-crabbe-berri --- .../prompt_templates/factory.py | 136 ++++++-- .../anthropic/test_message_sanitization.py | 298 ++++++++++++++++-- 2 files changed, 380 insertions(+), 54 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index cfb6e3ab45f..0aed2956b0a 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2205,15 +2205,97 @@ def anthropic_process_openai_file_message( _EMPTY_TEXT_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]" +def _dict_block(block: object) -> Mapping[str, object] | None: + if not isinstance(block, dict): + return None + return cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance; the converters only read dict blocks + + +def _is_empty_text_block(block: object) -> bool: + match _dict_block(block): + case {"type": "text", "text": str() as text}: + return not text.strip() + case {"type": "text"}: + return True + case _: + return False + + +def _with_placeholder_text(block: object) -> object: + fields: Final = _dict_block(block) + if fields is None or not _is_empty_text_block(fields): + return block + return {**fields, "text": _EMPTY_TEXT_PLACEHOLDER} + + +def _survives_assistant_conversion(block: object) -> bool: + """Whether ``anthropic_messages_pt`` keeps this list-content block of an assistant turn. + + Mirrors the assistant converter: a ``thinking`` / ``redacted_thinking`` block is kept + unless it is unsignable (a signed block is kept even with empty text), and + ``server_tool_use`` and ``*_tool_result`` blocks pass through. A redacted block with + blank data is dropped under ``drop_params``, which is not known here, so it never counts. + """ + fields: Final = _dict_block(block) + if fields is None: + return False + if (thinking_block := _as_thinking_block(fields)) is not None: + return not (is_unsignable_thinking_block(thinking_block) or has_blank_redacted_data(thinking_block)) + match fields: + case {"type": "server_tool_use"}: + return True + case {"type": str() as block_type}: + return block_type.endswith("_tool_result") + case _: + return False + + +def _tool_call_survives_conversion(tool_call: object) -> bool: + """Whether ``convert_to_anthropic_tool_invoke`` can emit a block for this tool call. + + Non-function calls are skipped there, and a tool search call is dropped when its + result is not replayed, which depends on the rest of the conversation, so neither counts. + """ + if message_field(tool_call, "type") != "function": + return False + match message_field(tool_call, "id"), message_field(message_field(tool_call, "function"), "name"): + case (str() as tool_id, str() as tool_name): + return not is_anthropic_tool_search_use(tool_id, tool_name) + case _: + return True + + +def _carries_non_text_content(message: AllMessageValues) -> bool: + if any(_tool_call_survives_conversion(tool_call) for tool_call in parts_of(message_field(message, "tool_calls"))): + return True + content: Final = message_field(message, "content") + blocks: Final = parts_of(content) + return isinstance(content, list) and any(_survives_assistant_conversion(block) for block in blocks) + + +def _with_content(message: AllMessageValues, content: object, action: str) -> AllMessageValues: + sanitized: Final = cast(AllMessageValues, {**message, "content": content}) + verbose_logger.debug("_sanitize_empty_text_content: %s in %s message", action, message_field(message, "role")) + return sanitized + + def _sanitize_empty_text_content( message: AllMessageValues, ) -> AllMessageValues: """ Case C: Sanitize empty text content - - Replace empty or whitespace-only text content with a placeholder message. - - Handles both string content and list-of-blocks content (rewriting only - the empty text blocks in place; non-text blocks like images are left - untouched). + + Anthropic rejects empty text blocks, so empty or whitespace-only text must never be + forwarded verbatim. What replaces it depends on what else the turn carries. + + Assistant turns that already carry substantive content (``tool_calls``, or non-text + blocks such as ``server_tool_use``) get the empty text dropped. Anthropic accepts an + assistant turn whose content is just ``tool_use``, so there is nothing to substitute, and + substituting anyway leaks a protocol diagnostic into the conversation as though the + assistant had said it, which the model then echoes on later turns. + + Every other empty turn would otherwise end up with no content at all, which Anthropic + also rejects, so the placeholder is still substituted there. Returns: The message with sanitized content if needed, otherwise the original message @@ -2221,42 +2303,26 @@ def _sanitize_empty_text_content( if message.get("role") not in ["user", "assistant"]: return message - content: Final = message.get("content") + content: Final = message_field(message, "content") + drop_empty_text: Final = message_field(message, "role") == "assistant" and _carries_non_text_content(message) if isinstance(content, str): - if not content or not content.strip(): - message = cast(AllMessageValues, dict(message)) # Make a copy - message["content"] = _EMPTY_TEXT_PLACEHOLDER - verbose_logger.debug( - "_sanitize_empty_text_content: Replaced empty text content in %s message", message.get("role") - ) + if content.strip(): + return message + if drop_empty_text: + return _with_content(message, None, "Dropped empty text content") + return _with_content(message, _EMPTY_TEXT_PLACEHOLDER, "Replaced empty text content") + + blocks: Final = parts_of(content) + if not isinstance(content, list) or not any(_is_empty_text_block(block) for block in blocks): return message - if isinstance(content, list): - # Walk the blocks and rewrite any empty text blocks. We rewrite (rather - # than drop) so callers don't end up with an entirely empty content - # list, which Anthropic also rejects. - new_blocks: Final[list[Any]] = [] - rewrote_any = False - for block in content: - if isinstance(block, dict) and block.get("type") == "text": - text = block.get("text") - if not isinstance(text, str) or not text or not text.strip(): - new_block = dict(block) - new_block["text"] = _EMPTY_TEXT_PLACEHOLDER - new_blocks.append(new_block) - rewrote_any = True - continue - new_blocks.append(block) + if drop_empty_text: + kept: Final = [block for block in blocks if not _is_empty_text_block(block)] + return _with_content(message, kept, "Dropped empty text block(s)") - if rewrote_any: - message = cast(AllMessageValues, dict(message)) # Make a copy - message["content"] = new_blocks - verbose_logger.debug( - "_sanitize_empty_text_content: Replaced empty text block(s) in %s message", message.get("role") - ) - - return message + rewritten: Final = [_with_placeholder_text(block) for block in blocks] + return _with_content(message, rewritten, "Replaced empty text block(s)") def _add_missing_tool_results( diff --git a/tests/unit/llms/anthropic/test_message_sanitization.py b/tests/unit/llms/anthropic/test_message_sanitization.py index 7afa60baf7e..0afd443d07a 100644 --- a/tests/unit/llms/anthropic/test_message_sanitization.py +++ b/tests/unit/llms/anthropic/test_message_sanitization.py @@ -7,6 +7,8 @@ B. Orphaned tool_result without matching tool_use C. Empty text content """ +import copy +import json import pytest import sys import os @@ -16,6 +18,8 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../. import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( + _EMPTY_TEXT_PLACEHOLDER as PLACEHOLDER, + _sanitize_empty_text_content, sanitize_messages_for_tool_calling, anthropic_messages_pt, ) @@ -24,16 +28,9 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( class TestMessageSanitization: """Test message sanitization for tool calling scenarios""" - def setup_method(self): - """Setup for each test""" - # Save original modify_params value - self.original_modify_params = litellm.modify_params - litellm.modify_params = True - - def teardown_method(self): - """Cleanup after each test""" - # Restore original modify_params value - litellm.modify_params = self.original_modify_params + @pytest.fixture(autouse=True) + def _modify_params_on(self, monkeypatch): + monkeypatch.setattr(litellm, "modify_params", True) def test_case_a_orphaned_tool_call_single(self): """ @@ -255,11 +252,11 @@ class TestMessageSanitization: assert sanitized[3]["content"] == "[System: Empty message content sanitised to satisfy protocol]" assert sanitized[4]["role"] == "assistant" - def test_modify_params_false_no_sanitization(self): + def test_modify_params_false_no_sanitization(self, monkeypatch): """ Test that sanitization is skipped when modify_params=False """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ {"role": "user", "content": ""}, @@ -287,7 +284,6 @@ class TestMessageSanitization: """ Test that sanitization is integrated into anthropic_messages_pt """ - litellm.modify_params = True messages = [ {"role": "user", "content": "What is the weather in Nashik?"}, @@ -316,7 +312,7 @@ class TestMessageSanitization: assert result[0]["role"] == "user" assert result[1]["role"] == "assistant" - def test_empty_string_content_sanitized_without_modify_params(self): + def test_empty_string_content_sanitized_without_modify_params(self, monkeypatch): """ Regression: An empty user message ({"role": "user", "content": ""}) must be rewritten to a non-empty placeholder *before* it reaches Anthropic, @@ -324,7 +320,7 @@ class TestMessageSanitization: "messages: text content blocks must be non-empty" Reproduces a real failure from the pr-review agent (pydantic-ai). """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ {"role": "user", "content": "First message"}, @@ -344,13 +340,13 @@ class TestMessageSanitization: assert block["text"].strip() != "" assert text_blocks[2]["text"] == ("[System: Empty message content sanitised to satisfy protocol]") - def test_empty_text_block_in_list_content_sanitized(self): + def test_empty_text_block_in_list_content_sanitized(self, monkeypatch): """ Same regression for the list-of-blocks form: {"role": "user", "content": [{"type": "text", "text": ""}]} Empty text *blocks* must be rewritten too, regardless of modify_params. """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ { @@ -372,12 +368,12 @@ class TestMessageSanitization: for block in text_blocks[1:]: assert block["text"].strip() != "" - def test_non_empty_content_unchanged_without_modify_params(self): + def test_non_empty_content_unchanged_without_modify_params(self, monkeypatch): """ Sanity check: when nothing is empty, the messages flow through unchanged even with modify_params disabled. """ - litellm.modify_params = False + monkeypatch.setattr(litellm, "modify_params", False) messages = [ {"role": "user", "content": "Hello"}, @@ -393,6 +389,270 @@ class TestMessageSanitization: assert result[1]["content"][0]["text"] == "Hi there" assert result[2]["content"][0]["text"] == "How are you?" + @pytest.mark.parametrize("modify_params", [True, False]) + @pytest.mark.parametrize( + "empty_content", + [None, "", " \n \t ", [], [{"type": "text", "text": ""}]], + ids=["none", "empty_str", "whitespace_str", "empty_list", "empty_text_block"], + ) + def test_tool_call_only_assistant_turn_never_gets_placeholder(self, monkeypatch, modify_params, empty_content): + """ + Regression for the placeholder leaking into tool-call-only assistant turns. + + An OpenAI-format assistant turn that carries tool_calls but no text is the normal + shape emitted by coding agents. Anthropic accepts an assistant turn whose content + is just tool_use, so there is nothing to substitute; injecting the placeholder made + it look like the assistant had said it, and the model echoed it on later turns. + """ + monkeypatch.setattr(litellm, "modify_params", modify_params) + + messages = [ + {"role": "user", "content": "list the files"}, + { + "role": "assistant", + "content": empty_content, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": '{"cmd": "ls"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "a.txt"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assistant_turns = [m for m in result if m["role"] == "assistant"] + assert len(assistant_turns) == 1 + blocks = assistant_turns[0]["content"] + + assert [b["type"] for b in blocks] == ["tool_use"] + assert blocks[0]["id"] == "call_1" + assert blocks[0]["name"] == "run_bash" + + assert PLACEHOLDER not in json.dumps(result) + + def test_empty_text_dropped_alongside_non_text_blocks(self): + """ + An assistant turn already carrying a non-text block has enough content for + Anthropic once the empty text block is dropped, so the placeholder must not be + substituted in its place. + """ + + message = { + "role": "assistant", + "content": [ + {"type": "text", "text": " "}, + {"type": "thinking", "thinking": "weighing options", "signature": "sig"}, + ], + } + + sanitized = _sanitize_empty_text_content(message) + + assert [b["type"] for b in sanitized["content"]] == ["thinking"] + assert PLACEHOLDER not in json.dumps(sanitized) + + @pytest.mark.parametrize( + "sibling_block", + [ + {"type": "thinking", "thinking": "", "signature": "sig"}, + {"type": "redacted_thinking", "data": "opaque"}, + ], + ids=["signed_empty_thinking", "redacted_thinking"], + ) + def test_empty_text_dropped_alongside_kept_thinking_block(self, sibling_block): + """ + The converter keeps signed thinking blocks even with empty text, and keeps + redacted_thinking blocks, so the turn still reaches Anthropic without the + empty text and the placeholder must not be substituted. + """ + + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}, sibling_block]}, + {"role": "user", "content": "and now?"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert [b["type"] for b in result[1]["content"]] == [sibling_block["type"]] + assert PLACEHOLDER not in json.dumps(result) + + def test_placeholder_kept_when_only_tool_call_is_dropped_tool_search(self, monkeypatch): + """ + A tool search call whose result is not replayed is dropped by the converter, so + it does not count as content: the empty text keeps its placeholder instead of + the assistant turn vanishing from the request. modify_params stays off so no + dummy tool result answers the call and keeps it. + """ + monkeypatch.setattr(litellm, "modify_params", False) + + messages = [ + {"role": "user", "content": "find a tool"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "srvtoolu_1", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": '{"query": "x"}'}, + } + ], + }, + {"role": "user", "content": "go on"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert result[1]["content"] == [{"type": "text", "text": PLACEHOLDER}] + + @pytest.mark.parametrize( + "sibling_block", + [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}}, + {"type": "thinking", "thinking": ""}, + {"type": "redacted_thinking", "data": " "}, + {"type": "unknown_future_block"}, + {"no_type_key": "malformed"}, + ], + ids=["image_url", "unsignable_thinking", "blank_redacted_thinking", "unknown_type", "missing_type"], + ) + def test_empty_text_kept_when_sibling_block_would_not_survive(self, sibling_block): + """ + Dropping empty text is only safe when something else in the turn actually reaches + Anthropic. Blocks the assistant conversion discards (assistant-side images, + signature-less thinking, unrecognised types) would leave an empty turn that gets + dropped from the request entirely, so the placeholder has to stay. + """ + + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}, sibling_block]}, + {"role": "user", "content": "and now?"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-haiku-4-5", llm_provider="anthropic") + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert PLACEHOLDER in json.dumps(result) + + def test_assistant_text_preserved_alongside_tool_calls(self): + """ + Dropping empty text must not drop real text: an assistant turn that says something + *and* calls a tool keeps both blocks. + """ + + messages = [ + {"role": "user", "content": "list the files"}, + { + "role": "assistant", + "content": "Let me check.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": '{"cmd": "ls"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "a.txt"}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assistant_turns = [m for m in result if m["role"] == "assistant"] + blocks = assistant_turns[0]["content"] + + assert [b["type"] for b in blocks] == ["text", "tool_use"] + assert blocks[0]["text"] == "Let me check." + assert PLACEHOLDER not in json.dumps(result) + + @pytest.mark.parametrize( + "messages", + [ + [{"role": "user", "content": ""}], + [{"role": "user", "content": [{"type": "text", "text": " "}]}], + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": ""}], + ], + ids=["user_str", "user_block", "assistant_without_tool_calls"], + ) + def test_placeholder_still_used_when_turn_would_be_empty(self, messages): + """ + The placeholder is still required where dropping the empty text would leave the + turn with no content at all, which Anthropic also rejects. + """ + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert PLACEHOLDER in json.dumps(result) + for message in result: + assert message["content"] + + def test_no_empty_text_block_or_content_reaches_anthropic(self): + """ + The invariant Anthropic actually enforces, over a conversation mixing every shape: + no empty content list and no empty text block may survive the conversion. + """ + + messages = [ + {"role": "user", "content": ""}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "out"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}]}, + {"role": "user", "content": [{"type": "text", "text": "thanks"}]}, + ] + + result = anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + for message in result: + content = message["content"] + assert content + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + assert block["text"].strip() + + def test_sanitization_does_not_mutate_caller_messages(self): + """ + Callers reuse their message list across retries and fallbacks, so sanitization has + to stay copy-on-write. + """ + + messages = [ + {"role": "user", "content": "go"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run_bash", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "out"}, + ] + before = copy.deepcopy(messages) + + anthropic_messages_pt(messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic") + + assert messages == before + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 560372456eff83f4d57e4247b41d04111514b3b8 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 13:54:18 -0700 Subject: [PATCH 17/43] fix(bedrock): price claude_platform calls from their own cost map rows (#48150) * fix(bedrock): price claude_platform calls from the anthropic cost map rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: avoid recursive Claude Platform lookups Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): price claude_platform calls from their own cost map rows Replace the claude_platform -> anthropic lookup redirect with bedrock/claude_platform/* rows, and apply the inference_geo multiplier in the bedrock cost path so US-only inference is priced at 1.1x Co-Authored-By: Claude Opus 5.5 * test(bedrock): exempt claude_platform rows from the Converse routing guard The claude_platform/ route is explicit in the row key, so these rows are not accidental non-Converse Bedrock chat rows Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- litellm/llms/bedrock/cost_calculation.py | 14 +- ...odel_prices_and_context_window_backup.json | 600 ++++++++++++++++++ model_prices_and_context_window.json | 600 ++++++++++++++++++ .../llms/bedrock/test_cost_calculation.py | 43 ++ tests/unit/test_utils_get_model_info.py | 1 + 5 files changed, 1255 insertions(+), 3 deletions(-) create mode 100644 tests/unit/llms/bedrock/test_cost_calculation.py diff --git a/litellm/llms/bedrock/cost_calculation.py b/litellm/llms/bedrock/cost_calculation.py index 0a2e39ab972..1b45531dda4 100644 --- a/litellm/llms/bedrock/cost_calculation.py +++ b/litellm/llms/bedrock/cost_calculation.py @@ -3,9 +3,13 @@ Helper util for handling bedrock-specific cost calculation - e.g.: prompt caching """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + generic_cost_per_token, + get_provider_specific_geo_multiplier, +) if TYPE_CHECKING: from litellm.types.utils import Usage @@ -17,9 +21,13 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) Follows the same logic as Anthropic's cost per token calculation. """ - return generic_cost_per_token( + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="bedrock") + prompt_cost, completion_cost = generic_cost_per_token( model=model, usage=usage, custom_llm_provider="bedrock", service_tier=service_tier, + model_info=model_info, ) + geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + return prompt_cost * geo_multiplier, completion_cost * geo_multiplier diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index eb9bf6aa8bf..c7dfce0712f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -77996,5 +77996,605 @@ "supports_reasoning": true, "supports_vision": false, "source": "https://opencode.ai/v2/docs/console/go/" + }, + "bedrock/claude_platform/claude-fable-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-fable-5-1": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_computer_use": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-5-5": { + "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, + "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, + "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_1hr": 2e-07, + "cache_creation_input_token_cost_batches": 6.25e-08, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_batches": 5e-09, + "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_batches": 2.5e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": true, + "thinking_always_on": false, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true, + "input_cost_per_token_above_100k_tokens": 5e-07, + "output_cost_per_token_above_100k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, + "cache_read_input_token_cost_above_100k_tokens": 5e-08 + }, + "bedrock/claude_platform/claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-7": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-8": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost_batches": 1.5e-07, + "deprecation_date": "2026-11-30", + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "provider_specific_entry": { + "us": 1.1 + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5-5": { + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": false, + "thinking_always_on": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index eb9bf6aa8bf..c7dfce0712f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -77996,5 +77996,605 @@ "supports_reasoning": true, "supports_vision": false, "source": "https://opencode.ai/v2/docs/console/go/" + }, + "bedrock/claude_platform/claude-fable-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-fable-5-1": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_computer_use": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-5-5": { + "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, + "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, + "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_1hr": 2e-07, + "cache_creation_input_token_cost_batches": 6.25e-08, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_batches": 5e-09, + "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_batches": 2.5e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": true, + "thinking_always_on": false, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true, + "input_cost_per_token_above_100k_tokens": 5e-07, + "output_cost_per_token_above_100k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, + "cache_read_input_token_cost_above_100k_tokens": 5e-08 + }, + "bedrock/claude_platform/claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-7": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-8": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost_batches": 1.5e-07, + "deprecation_date": "2026-11-30", + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "provider_specific_entry": { + "us": 1.1 + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5-5": { + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": false, + "thinking_always_on": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true } } diff --git a/tests/unit/llms/bedrock/test_cost_calculation.py b/tests/unit/llms/bedrock/test_cost_calculation.py new file mode 100644 index 00000000000..f4718344e7e --- /dev/null +++ b/tests/unit/llms/bedrock/test_cost_calculation.py @@ -0,0 +1,43 @@ +import json +from importlib.resources import files +from typing import Final + +import pytest + +import litellm +from litellm.types.utils import Usage + +CLAUDE_PLATFORM_PREFIX: Final = "bedrock/claude_platform/" +CLAUDE_PLATFORM_KEYS: Final = tuple( + key + for key in json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) + if key.startswith(CLAUDE_PLATFORM_PREFIX) +) + + +def test_claude_platform_rows_exist() -> None: + assert CLAUDE_PLATFORM_KEYS + + +@pytest.mark.parametrize("inference_geo", ["global", "us"]) +@pytest.mark.parametrize("key", CLAUDE_PLATFORM_KEYS) +def test_claude_platform_cost_matches_claude_api_cost( + local_model_cost_map: None, key: str, inference_geo: str +) -> None: + usage: Final = Usage( + prompt_tokens=1_000, + completion_tokens=200, + cache_read_input_tokens=300, + cache_creation_input_tokens=250, + inference_geo=inference_geo, + ) + + claude_platform_cost: Final = litellm.cost_per_token( + model=key.removeprefix("bedrock/"), custom_llm_provider="bedrock", usage_object=usage + ) + claude_api_cost: Final = litellm.cost_per_token( + model=key.removeprefix(CLAUDE_PLATFORM_PREFIX), custom_llm_provider="anthropic", usage_object=usage + ) + + assert claude_platform_cost == pytest.approx(claude_api_cost) + assert sum(claude_platform_cost) > 0 diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index 4bd16a2997f..302a7d60cdd 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -160,6 +160,7 @@ def _enforce_bedrock_converse_models( info.get("litellm_provider") == "bedrock" and info.get("mode") == "chat" and model not in whitelist_models + and BedrockModelInfo.get_bedrock_route(model) != "claude_platform" and not ( (base_model := BedrockModelInfo.get_base_model(model)) != model and model_cost.get(base_model, {}).get("litellm_provider") == "bedrock_converse" From 287abaf372b337dc3a81cda7b6b91ba40e6cbf4e Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:07:26 -0700 Subject: [PATCH 18/43] chore(cost-map): add openai gpt-5.6-cyber 12-hour cache write rate from the pricing page (#48246) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c7dfce0712f..a32ee86dd74 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -37401,6 +37401,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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c7dfce0712f..a32ee86dd74 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -37401,6 +37401,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, From c0dde21c3b4a0892a178a3f71d2f86b6edfa7889 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:07:48 -0700 Subject: [PATCH 19/43] fix(bedrock): take Nova Sonic context and output limits from the Bedrock model cards (#48247) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 9 ++++++++- model_prices_and_context_window.json | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a32ee86dd74..630b1ed9441 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 }, @@ -75363,10 +75367,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/model_prices_and_context_window.json b/model_prices_and_context_window.json index a32ee86dd74..630b1ed9441 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 }, @@ -75363,10 +75367,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 }, From fac7459b3ef2005bdbb6d521efe08c7b107f2636 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:08:07 -0700 Subject: [PATCH 20/43] fix(gemini): bill gemini-2.5-flash-native-audio video input at the audio / video rate (#48248) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 630b1ed9441..a3fff5ca26d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59519,6 +59519,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, @@ -59714,6 +59715,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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 630b1ed9441..a3fff5ca26d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59519,6 +59519,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, @@ -59714,6 +59715,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, From 015096af0b9b24eb1b1d6a957650685499aecdd7 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:08:23 -0700 Subject: [PATCH 21/43] chore(pricing): add github_copilot/claude-opus-4.8-fast deprecation date (#48251) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a3fff5ca26d..058d49133f0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33216,6 +33216,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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a3fff5ca26d..058d49133f0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33216,6 +33216,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, From 29edcee9ebff4b7d759eafa2610d90633ca4b652 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:08:37 -0700 Subject: [PATCH 22/43] fix(pricing): set azure gpt-chat-latest max_input_tokens to 128000 (#48252) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 058d49133f0..02eaa906522 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4343,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", @@ -9236,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", @@ -9683,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", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 058d49133f0..02eaa906522 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4343,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", @@ -9236,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", @@ -9683,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", From 172aa72ad2dc9f044b38d0fe7f250293a8fda2ef Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Sat, 24 Oct 2026 15:12:07 -0700 Subject: [PATCH 23/43] feat(responses): honor max_agentic_loops for the streamed MCP auto-execute round cap (#47688) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/mcp/mcp_streaming_iterator.py | 13 ++++-- .../mcp/test_aresponses_api_with_mcp.py | 32 +++++++++++--- .../mcp/test_mcp_streaming_iterator.py | 43 +++++++++++++------ 3 files changed, 65 insertions(+), 23 deletions(-) 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/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): """ From 1418c1990fcc403d7598d97832f1272f45049299 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:33:17 -0700 Subject: [PATCH 24/43] feat(guardrails): add native Reco guardrail provider (pre_call), fail closed by default (#48243) * feat(guardrails): add native Reco guardrail provider (pre_call) Registers guardrail: reco as a first-class vendor integration, built on the existing generic guardrail API wire contract rather than reimplementing the HTTP client and payload mapping. Maps reco_tenant_id to the X-Reco-Tenant-Id header and hardcodes fail_open/fail_on_error since those should not be customer-configurable. v1 only supports mode: pre_call; post_call and streaming redaction already work against the existing generic pilot and are deferred. Also overrides apply_guardrail to re-raise blocked-request errors with the guardrail's own configured name: the base class hardcodes a module-level name in that exception, which would otherwise mislabel every blocked Reco request in logs. * feat(guardrails): configurable Reco fail modes Expose unreachable_fallback and fail_on_error on the Reco guardrail and keep fail-open as the default. Values under optional_params win. Top-level litellm_params values apply only when set explicitly, since their defaults are fail-closed and would otherwise flip the default Also type RecoGuardrail's forwarded kwargs with an Unpack'd TypedDict instead of silencing ANN003 Signed-off-by: MickeyShnaiderman-RecoLabs Co-authored-by: Cursor * chore(guardrails): align Reco with main's CI gates Upstream moved unit tests from tests/test_litellm to tests/unit, so test_reco.py was invoked by no CI job. Move it next to the other guardrail tests Replace RecoGuardrail's forwarded kwargs with explicit keyword params, read optional_params through a pydantic adapter instead of untyped dict/getattr access, freeze RecoOptionalParams, and drop mutable-ok markers that suppress nothing, to stay within main's type-discipline and basedpyright budgets Signed-off-by: MickeyShnaiderman-RecoLabs Co-authored-by: Cursor * feat(guardrails): fail closed by default in the Reco guardrail Matches GenericGuardrailAPI and every other guardrail here: a Reco error or outage now blocks the request unless the admin sets fail_on_error: false. Tests cover the default, explicit fail-open, and top-level vs optional_params precedence Co-authored-by: RecoLabs-YaronBratspiess <208601538+RecoLabs-YaronBratspiess@users.noreply.github.com> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): honor explicit fail-open on Reco when optional_params is a model model_dump(exclude_unset=True) so pydantic defaults in optional_params no longer override an explicit top-level fail_on_error. Also drop the duplicate log_guardrail_information decorator, call the parent by keyword, freeze RecoConfigModel, and type the tests Co-authored-by: RecoLabs-YaronBratspiess <208601538+RecoLabs-YaronBratspiess@users.noreply.github.com> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): keep Reco config model assignable and the hook decorator the gate requires frozen=True on RecoConfigModel froze the shared LitellmParams mixin, which the guardrail registry assigns resolved secrets on, so every guardrail failed to load. Dropped it and added a test that LitellmParams stays assignable. Restored log_guardrail_information on the Reco override since tests/code_coverage_tests/check_guardrail_apply_decorator.py requires it on every hook. Test params typed as Mapping for LIT001 Co-authored-by: RecoLabs-YaronBratspiess <208601538+RecoLabs-YaronBratspiess@users.noreply.github.com> Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(guardrails): format the Reco test file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Signed-off-by: MickeyShnaiderman-RecoLabs Co-authored-by: Yaron Bratspiess Co-authored-by: MickeyShnaiderman-RecoLabs Co-authored-by: Cursor Co-authored-by: joshua Co-authored-by: MickeyShnaiderman-RecoLabs <77006830+MickeyShnaiderman-RecoLabs@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/reco/__init__.py | 72 +++++ .../guardrails/guardrail_hooks/reco/reco.py | 84 ++++++ litellm/types/guardrails.py | 5 + .../proxy/guardrails/guardrail_hooks/reco.py | 48 +++ .../guardrails/guardrail_hooks/test_reco.py | 285 ++++++++++++++++++ 5 files changed, 494 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/reco/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/reco/reco.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/reco.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_reco.py 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/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/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"] From 34ec3fcdfae89cac4edb873b6e80027789a05ea4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 16:11:33 -0700 Subject: [PATCH 25/43] fix(router): skip declared text-only deployments for image requests on every routing strategy (#48241) * fix(router): skip declared text-only deployments for image requests without enable_pre_call_checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): run the vision filter after the cooldown filter so a cooled vision leg falls back to the text leg Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): apply the vision filter after pre-call checks and tidy the new tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): mark the remaining locals in the vision filter tests Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 13 ++- litellm/router.py | 54 +++++++---- .../proxy/test_common_request_processing.py | 21 +++++ tests/unit/test_router/test_router.py | 93 ++++++++++++++----- 4 files changed, 135 insertions(+), 46 deletions(-) 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/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/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/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( From 6268c0a281f05671c98528cb6b59e9c3fd2a06f0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 16:15:10 -0700 Subject: [PATCH 26/43] feat(proxy): add disable_pre_routing_error_logs to skip spend rows for pre-routing rejections (#48239) * feat(proxy): add disable_pre_routing_error_logs to skip spend rows for pre-routing rejections Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): decide pre-routing rejections by pipeline stage, not exception type Guardrails such as AIM raise ProxyException(400), the same shape as an unparsable body, so the status code list skipped their spend rows. The pipeline stamps litellm_call_id only after auth, body parsing and model resolution, so a failure reaching the hook without one died before any provider was considered. _handle_logging_proxy_only_error flags those requests when it mints the call id, and is_pre_routing_rejection keys off that flag plus ProxyModelNotFoundError. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark pre-routing rejections with a server-owned sentinel request_data is the client body, so a boolean marker or a litellm_call_id in it can be sent by the client to hide a guardrail block or force a spend row. The failure handler now stores an object() sentinel whenever it finds no logging object, and is_pre_routing_rejection checks identity, which JSON cannot forge. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/proxy_track_cost_callback.py | 40 +++++- litellm/proxy/utils.py | 8 ++ .../hooks/test_proxy_track_cost_callback.py | 115 +++++++++++++++++- .../test_post_call_failure_hook.py | 59 ++++++++- 4 files changed, 211 insertions(+), 11 deletions(-) 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/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/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/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") From 9b7f2c8dead2a2ea5a0b6b4afbcf065d33feb8bd Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 16:44:29 -0700 Subject: [PATCH 27/43] fix(opencode): bill OpenCode Go requests from the reported list_cost (#48254) * fix(opencode): bill OpenCode Go requests from the reported list_cost OpenCode returns usage.opencode.list_cost on every response. Read it into the provider-reported cost the cost calculator checks before the price map, for chat and Responses, streamed and not. Keep the static rows as the fallback, and stop the generic usage.cost path from billing Go requests at 0 (usage.cost is what hit the prepaid balance, 0 inside the plan). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(core): resolve reported cost breakdowns without recursion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(opencode): cover reported list_cost spend through the proxy Replaces the recursive _reported_cost_value with a flat _cost_number helper so the code-quality gate passes, and adds an integration test that runs the proxy against a scripted OpenCode upstream for a non-stream and a stream chat completion, asserting the response cost header and the spend log row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(opencode): bill native /v1/messages from the reported list_cost The native Anthropic Messages route priced OpenCode Go bodies from the price rows twice over: the non-stream body is a TypedDict the response-cost header path priced as-is, and the stream assembled its ModelResponse with the rows cost already stamped. Both now read usage.opencode.list_cost, the non-stream body through billable_result and the stream from the final message_delta. The integration test now covers chat, responses and messages, stream and non-stream, through the proxy against a scripted OpenCode upstream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(opencode): keep the stream cost stamp inside the lint gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(opencode): read list_cost off dict chunks and trust it only from OpenCode on /v1/messages Stream builders receive chunks both as mappings and as ModelResponseStream objects, so the list_cost scan now reads usage through chunk_usage. The native Messages route only trusts a reported cost from opencode_go, every other upstream keeps the price map, which also keeps a zero usage.cost from being billed as free. Served bodies without stop_sequence still price, and the new code stays inside the basedpyright gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 65 +++ litellm/litellm_core_utils/litellm_logging.py | 27 +- .../streaming_chunk_builder_utils.py | 40 +- .../litellm_core_utils/streaming_handler.py | 23 +- litellm/llms/opencode/chat/transformation.py | 57 ++- litellm/llms/opencode/common_utils.py | 11 + .../llms/opencode/responses/transformation.py | 40 +- .../anthropic_passthrough_logging_handler.py | 56 ++- .../test_opencode_go_reported_cost_wire.py | 198 ++++++++ .../test_litellm_logging.py | 276 +++++++----- .../llms/opencode/chat/test_transformation.py | 180 ++++++++ ...t_anthropic_passthrough_logging_handler.py | 422 ++++++++---------- 12 files changed, 1024 insertions(+), 371 deletions(-) create mode 100644 tests/integration/pricing/test_opencode_go_reported_cost_wire.py 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/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/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/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/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/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/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/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 From 205cfd8044269ac43c2a91c6cafac36146e29171 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 24 Oct 2026 16:47:15 -0700 Subject: [PATCH 28/43] fix(proxy): restore Claude gateway model discovery (#48204) (#48221) * fix(proxy): restore Claude gateway model discovery (#48204) * test(proxy): cover gateway discovery after OAuth sign-in (#48204) * refactor(proxy): share model listing route policy (#48204) * test(proxy): isolate SSO settings environment (#48204) --- .../proxy/auth/route_checks.py | 4 +- litellm/proxy/_types.py | 11 +-- .../anthropic_endpoints/gateway_endpoints.py | 8 ++ litellm/proxy/auth/auth_checks.py | 7 +- litellm/proxy/auth/route_checks.py | 4 +- .../test_claude_code_gateway_oauth.py | 27 +++++- .../auth/test_managed_authorization.py | 2 +- .../test_gateway_endpoints.py | 92 +++++++++++++++++++ .../test_proxy_setting_endpoints.py | 8 ++ 9 files changed, 148 insertions(+), 15 deletions(-) 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/proxy/_types.py b/litellm/proxy/_types.py index 3a3306dfdc0..2b01e02bb06 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", @@ -760,8 +760,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..06c53afcf2c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -47,6 +47,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 +911,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 +1028,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", 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/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/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/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 From 387b8b34777d0e2f61b679c7cb90524bf36a3a65 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 24 Oct 2026 17:06:31 -0700 Subject: [PATCH 29/43] fix(anthropic): release upstream messages stream on disconnect (#48197) (#48220) * fix(anthropic): close upstream messages stream on disconnect (#48197) * test(anthropic): prove disconnect cleanup over live proxy (#48197) --- .../adapters/streaming_iterator.py | 21 +- .../test_chat_bridge_disconnect_wire.py | 232 ++++++++++++++++++ .../test_streaming_iterator_sse_stream.py | 67 +++++ 3 files changed, 319 insertions(+), 1 deletion(-) create mode 100644 tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_disconnect_wire.py 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/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/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 From 5367855644b099170460a6ec9161be6effe0cd8f Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Sat, 24 Oct 2026 17:09:55 -0700 Subject: [PATCH 30/43] feat(teams): show caller member spend and budget on Your Teams (#45674) * feat(teams): show caller member spend and budget on Your Teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(teams): use member-linked reset dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(teams): keep type-discipline suppressions on cast lines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(teams): preserve empty page list response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(teams): keep the member reset date inside the spend column Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(teams): match email-only roster members for caller membership Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): show team member budget reset date in UTC Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(teams): address review on caller membership enrichment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): label the caller's team spend line My Spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 52 ++- .../management_endpoints/team_endpoints.py | 241 ++++++++-- .../management_endpoints/team_endpoints.py | 12 + .../management/test_team_list_v2.py | 32 +- .../proxy/auth/test_team_member_budget.py | 55 ++- .../test_team_endpoints.py | 435 ++++++++++++++++-- .../TeamMemberSpendBudgetCell.test.tsx | 90 ++++ .../TeamsPage/TeamMemberSpendBudgetCell.tsx | 70 +++ .../TeamsPage/TeamsTable.integration.test.tsx | 24 +- .../components/TeamsPage/teamTableColumns.tsx | 12 +- .../components/key_team_helpers/key_list.tsx | 2 + .../shared/table_cells/spend_budget_cell.tsx | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 16 + .../src/utils/budgetUtils.test.ts | 23 + ui/litellm-dashboard/src/utils/budgetUtils.ts | 10 + 15 files changed, 987 insertions(+), 91 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.test.tsx create mode 100644 ui/litellm-dashboard/src/components/TeamsPage/TeamMemberSpendBudgetCell.tsx create mode 100644 ui/litellm-dashboard/src/utils/budgetUtils.test.ts diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 06c53afcf2c..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 @@ -1799,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, @@ -7034,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/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 361778ab2ea..95d4d03a7f0 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): ... @@ -5467,16 +5499,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 — @@ -5886,11 +5917,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 +5957,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 +6201,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 +6276,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 +6394,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/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/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/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/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 37028558af3..4feb7fb5bc8 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() 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/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/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b5a86dd2187..e81fa946833 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46319,6 +46319,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. @@ -46355,6 +46370,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 */ 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..a50fecabe83 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/budgetUtils.test.ts @@ -0,0 +1,23 @@ +import { afterEach, describe, expect, it } from "vitest"; + +import { 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"); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/budgetUtils.ts b/ui/litellm-dashboard/src/utils/budgetUtils.ts index 3d3278db88f..dfc57f2fa66 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,10 @@ 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"); +} From 68b9370c7c280426d9d95bb98693b7eec8168e57 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 24 Oct 2026 17:11:46 -0700 Subject: [PATCH 31/43] fix(chatgpt): route uncatalogued models through responses (#48244) --- litellm/responses/bridge_check.py | 2 +- .../test_responses_api_bridge_flag.py | 90 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) 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/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") From 943583722a03db41619d614b47d318277b93de11 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Sat, 24 Oct 2026 17:14:16 -0700 Subject: [PATCH 32/43] fix(ui): show "Never resets" instead of "n/a" for an unset budget reset period (#47065) * fix(ui): show "Never resets" instead of "n/a" for an unset budget reset period Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): label unset reset period as Keep current/No change where blank keeps the stored value Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): show a single Never resets option in create team when no default is set Without a default team budget_duration, the create team Reset Budget dropdown offered both "Not set" and "Never resets", and both created a team that never resets. The explicit Never resets item is now only shown when there is a default to opt out of, otherwise the empty option itself reads "Never resets" Also makes the Reset Budget create tests use chooseSelectOption so they no longer flake on the select popup's pointer-events transition --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../components/AccessGroupBudgetModal.tsx | 1 + ...cessGroupBudgetsPanel.integration.test.tsx | 30 ++++++++++++++ .../users/_components/BulkEditUsers.test.tsx | 13 ++++++ .../users/_components/user_edit_view.test.tsx | 29 +++++++++++++ .../users/_components/user_edit_view.tsx | 9 +++- .../user_info_view.integration.test.tsx | 4 +- .../src/components/Teams.test.tsx | 17 ++++---- ui/litellm-dashboard/src/components/Teams.tsx | 4 +- .../budget_duration_dropdown.test.tsx | 41 +++++++++++++++++++ .../budget_duration_dropdown.tsx | 2 +- .../team/EditMembership.integration.test.tsx | 21 ++++++++++ 11 files changed, 158 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.test.tsx 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/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/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, From 2fa6f22aff24b08030bb756c7b772f5a1678989d Mon Sep 17 00:00:00 2001 From: pfolex <93488901+pfolex@users.noreply.github.com> Date: Sat, 24 Oct 2026 20:16:33 -0400 Subject: [PATCH 33/43] feat(models): add Cerebras Llama 3.1 hyphenated aliases (cerebras/llama-3.1-8b, cerebras/llama-3.1-70b) (#48157) (#48158) --- ...odel_prices_and_context_window_backup.json | 22 ++++ model_prices_and_context_window.json | 22 ++++ .../test_cerebras_llama_model_metadata.py | 115 ++++++++++++++++++ 3 files changed, 159 insertions(+) create mode 100644 tests/unit/test_cerebras_llama_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02eaa906522..98df2fc4204 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15663,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", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 02eaa906522..98df2fc4204 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15663,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", 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 From 5171b3fbd85051969ee4dcf512dec634fd8f472b Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 17:17:09 -0700 Subject: [PATCH 34/43] feat(cost): parse 12-hour cache write rates for batch and ultrafast tiers (#48259) * feat(cost): parse 12-hour cache write rates for batch and ultrafast tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost): bill 12-hour batch cache writes once when only the 12hr batch key exists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 23 +- .../litellm_core_utils/llm_cost_calc/utils.py | 11 +- litellm/types/utils.py | 8 + litellm/utils.py | 12 + .../cost_calculation/cost_tracking_case.py | 4 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 80 ++ tests/unit/test_cost_calculator.py | 106 +- tests/unit/test_utils.py | 910 +++++++++--------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 24 + 9 files changed, 730 insertions(+), 448 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a0da0f6ddca..8b02f93fca4 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, @@ -2969,7 +2970,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 +3020,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 +3043,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/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/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/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/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/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_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/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e81fa946833..262e8b6deae 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35223,14 +35223,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 */ @@ -50253,14 +50261,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 */ @@ -50558,14 +50574,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 */ From 222389501066caa22bb070fa93ec620e345609ee Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 24 Oct 2026 17:29:48 -0700 Subject: [PATCH 35/43] fix(cloudflare): clarify chat and decisions endpoints (#47965) (#48265) --- litellm/llms/cloudflare/chat/transformation.py | 2 +- litellm/llms/cloudflare/decisions/transformation.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) 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: From cb65a21439eab88a6498780591a7f894427f6fe3 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 24 Oct 2026 17:37:38 -0700 Subject: [PATCH 36/43] fix(decisions): preserve missing native usage (#48147) (#48261) --- litellm/cost_calculator.py | 2 ++ .../llms/base_llm/decisions/transformation.py | 10 ++++-- .../fireworks_ai/decisions/transformation.py | 2 ++ .../llms/openai/decisions/transformation.py | 13 ++++---- .../openrouter/decisions/transformation.py | 2 +- litellm/types/decisions.py | 2 +- .../providers/test_decisions_wire.py | 19 +++++++++++ tests/unit/decisions/test_main.py | 32 +++++++++++++++++++ 8 files changed, 71 insertions(+), 11 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8b02f93fca4..8331c68fe24 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1551,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 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/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/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/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/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/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/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) From 84754cca0a560f2488f7f844d43d5f6583f07d32 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 24 Oct 2026 17:52:01 -0700 Subject: [PATCH 37/43] fix(groq): allow streamed reasoning replies in history (#47931) (#48268) --- litellm/llms/groq/chat/transformation.py | 4 +- .../chat/test_groq_chat_transformation.py | 66 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) 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/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." From 24a7bf6af728272eca2ba3277a2f3cdfa7e60863 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 17:57:09 -0700 Subject: [PATCH 38/43] fix(proxy): reconcile team member total_spend with current cycle spend (#48260) * fix(proxy): reconcile team member total_spend with current cycle spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): read /team/list as a list and drain the scripted provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(repositories): compose the membership reset writer and drop SQL text asserts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: type the cascade repair test and flatten memberships without a stacked comprehension Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: keep the reset job test file's existing typing untouched Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/reset_budget_job.py | 4 +- litellm/proxy/db/db_spend_update_writer.py | 4 +- .../management_endpoints/team_endpoints.py | 38 ++- litellm/repositories/prisma_protocols.py | 2 + litellm/repositories/unit_of_work.py | 41 ++- ..._team_member_total_spend_reconciliation.py | 273 ++++++++++++++++++ .../common_utils/test_reset_budget_job.py | 33 +++ .../proxy/db/test_db_spend_update_writer.py | 5 +- .../test_team_endpoints.py | 104 ++++++- tests/unit/repositories/test_unit_of_work.py | 21 +- .../src/components/team/MyUserTab.test.tsx | 25 ++ .../src/components/team/MyUserTab.tsx | 18 +- .../components/team/TeamMemberTab.test.tsx | 42 ++- .../src/components/team/TeamMemberTab.tsx | 19 +- .../src/utils/budgetUtils.test.ts | 25 +- ui/litellm-dashboard/src/utils/budgetUtils.ts | 38 +++ 16 files changed, 655 insertions(+), 37 deletions(-) create mode 100644 tests/integration/spend/test_team_member_total_spend_reconciliation.py 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 24ab09283f1..8ac2597c99b 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/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 95d4d03a7f0..1fa18e9b2b8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1013,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 @@ -4202,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) @@ -4286,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, @@ -5548,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, ) 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/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/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 6071569b806..a42b6404278 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/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 4feb7fb5bc8..f727448f3ac 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -12561,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): """ @@ -16608,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) @@ -16631,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() @@ -16755,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", @@ -17994,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/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/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/utils/budgetUtils.test.ts b/ui/litellm-dashboard/src/utils/budgetUtils.test.ts index a50fecabe83..a3b5f4da00a 100644 --- a/ui/litellm-dashboard/src/utils/budgetUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/budgetUtils.test.ts @@ -1,6 +1,6 @@ import { afterEach, describe, expect, it } from "vitest"; -import { formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils"; +import { describeBudgetResetSchedule, formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils"; const originalTimezone: string | undefined = process.env.TZ; @@ -21,3 +21,26 @@ describe("formatBudgetResetUtc", () => { 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 dfc57f2fa66..171ce5818f7 100644 --- a/ui/litellm-dashboard/src/utils/budgetUtils.ts +++ b/ui/litellm-dashboard/src/utils/budgetUtils.ts @@ -16,3 +16,41 @@ export function formatBudgetResetUtc(iso: string | null | undefined): string | n 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}`, + }; +} From 2e52e0fb2a84b03f000b2c6498578036306710bd Mon Sep 17 00:00:00 2001 From: Hazelhof <110658790+Hazelhof@users.noreply.github.com> Date: Sun, 25 Oct 2026 12:02:46 +1100 Subject: [PATCH 39/43] fix(deepseek): pass reasoning_effort through to the API (#38736) * fix(deepseek): pass reasoning_effort through to the API DeepSeekChatConfig.map_openai_params mapped reasoning_effort to a binary thinking {type: enabled/disabled} toggle and discarded the effort value, so DeepSeek always ran default high effort. High-effort reasoning can run away and exhaust max_tokens during reasoning, returning empty content. Preserve the reasoning_effort value into the request body when thinking is enabled so callers can bound effort (e.g. reasoning_effort="low"). Backward-compatible. See https://api-docs.deepseek.com/guides/thinking_mode * test(deepseek): reasoning_effort passthrough coverage * fix(deepseek): exclude reasoning_effort=none sentinel from passthrough "none" is the thinking-OFF sentinel, not a real effort value. Passing it alongside explicit thinking:enabled would send contradictory instructions to DeepSeek. Only forward real effort values (low/medium/high/xhigh/max). Addresses review feedback on the reasoning_effort passthrough change. * refactor(deepseek): name supported effort values and type mapped params * fix(deepseek): keep mapping types local to the adapter * fix(deepseek): remove cast and type mapped return --------- Co-authored-by: moe-berri --- litellm/llms/deepseek/chat/transformation.py | 11 +++- .../chat/test_deepseek_chat_transformation.py | 56 +++++++++++++++++-- 2 files changed, 60 insertions(+), 7 deletions(-) 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/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 From de5dc5b238b6e972869d7dedc6ba3c4fb8d931fe Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 24 Oct 2026 18:20:36 -0700 Subject: [PATCH 40/43] feat(mistral): add zai-glm-5-3-flash to the cost map (#48270) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 23 +++++++++++++++++++ model_prices_and_context_window.json | 23 +++++++++++++++++++ .../test_model_prices_and_context_window.py | 15 ++++-------- 3 files changed, 50 insertions(+), 11 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 98df2fc4204..9a7d88b8eda 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -42726,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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 98df2fc4204..9a7d88b8eda 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -42726,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, 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 == [] From d3b70c42ed1599727669ea7883c3971008a9113d Mon Sep 17 00:00:00 2001 From: Jeevan Babu Gotru <114808955+jeevan1098@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:15:46 -0500 Subject: [PATCH 41/43] fix(proxy): keep JSON integers outside the 64-bit range exact when parsing request bodies (#48213) Preserve oversized JSON integers in proxy requests and model rewrites Parse large bodies in a worker thread to avoid blocking request handling Fixes #47947 Co-authored-by: Jeevan Babu Gotru <114808955+jeevan1098@users.noreply.github.com> Co-authored-by: moe-berri --- .../proxy/common_utils/http_parsing_utils.py | 67 ++++++++++++-- .../common_utils/test_http_parsing_utils.py | 92 +++++++++++++++++++ 2 files changed, 153 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index e34d03a04c3..9a3aed0b67b 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -1,3 +1,4 @@ +import asyncio import json import re from collections.abc import Collection, Mapping @@ -161,10 +162,54 @@ def coerce_numeric_form_fields( } -def _parse_binary_body(body: bytes) -> dict: +_ORJSON_INT_MIN: Final = -(2**63) +_ORJSON_INT_MAX: Final = 2**64 - 1 +_JSON_PARSE_THREAD_THRESHOLD: Final = 64 * 1024 +_JSON_LONG_DIGITS: Final = re.compile(rb"[0-9]{19}") +_JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) +_JSON_BODY_BYTES: Final = TypeAdapter(bytes) +# Consume unterminated strings through EOF to avoid rescanning every escaped quote in malformed bodies. +_JSON_STRING_OR_NUMBER: Final = re.compile( + rb'"[^"\\]*(?:\\.[^"\\]*)*(?:"|\\?$)|(-?[0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?)', re.DOTALL +) + + +def _has_int_outside_orjson_range(body: bytes) -> bool: + if _JSON_LONG_DIGITS.search(body) is None: + return False + for token in _JSON_STRING_OR_NUMBER.finditer(body): + if ( + (number := token.group(1)) is None + or len(number) < 19 + or any(marker in number for marker in (b".", b"e", b"E")) + ): + continue + if len(number) > 20 or not _ORJSON_INT_MIN <= int(number) <= _ORJSON_INT_MAX: + return True + return False + + +def _loads_json_body(body: bytes) -> Mapping[str, object]: + if not _has_int_outside_orjson_range(body): + return _JSON_REQUEST_BODY.validate_python(orjson.loads(body)) + try: + parsed: Final = _JSON_REQUEST_BODY.validate_python(json.loads(body)) + json.dumps(parsed, ensure_ascii=False, allow_nan=False).encode("utf-8") + except ValueError: + return _JSON_REQUEST_BODY.validate_python(orjson.loads(body)) + return parsed + + +async def _loads_json_body_async(body: bytes) -> Mapping[str, object]: + if len(body) > _JSON_PARSE_THREAD_THRESHOLD: + return await asyncio.to_thread(_loads_json_body, body) + return _loads_json_body(body) + + +async def _parse_binary_body(body: bytes) -> dict: """JSON sent under a binary content type still parses; real binary (protobuf) carries no params -> {}.""" try: - parsed: Final = _DECODED_JSON.validate_python(orjson.loads(body)) + parsed: Final = _DECODED_JSON.validate_python(await _loads_json_body_async(body)) if isinstance(parsed, dict): return parsed except orjson.JSONDecodeError: @@ -236,7 +281,7 @@ async def read_request_body(request: Request | None) -> dict: if _normalize_media_type(content_type) in _BINARY_CONTENT_TYPES: binary_body: Final = await request.body() _mark_body_received(len(binary_body)) - parsed_body = _parse_binary_body(binary_body) + parsed_body = await _parse_binary_body(binary_body) elif _is_form_content_type(content_type): try: form_data: Final = await request.form() @@ -272,7 +317,7 @@ async def read_request_body(request: Request | None) -> dict: parsed_body = {} else: try: - parsed_body = orjson.loads(body) + parsed_body = dict(await _loads_json_body_async(body)) except orjson.JSONDecodeError as e: # The surrogate-repair fallback below runs two full-body re.sub # passes, which block the event loop on multi-MB malformed bodies. @@ -403,6 +448,14 @@ def detach_request_parsed_body( return client_body +def _dumps_json_body(body: Mapping[str, object]) -> bytes: + try: + encoded: Final = _JSON_BODY_BYTES.validate_python(orjson.dumps(body)) + except TypeError: + return json.dumps(body, separators=(",", ":")).encode("utf-8") + return encoded + + def rewrite_request_model( request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader request: Request | None, @@ -415,10 +468,12 @@ def rewrite_request_model( if request is None: return cached_body: Final = _safe_get_request_parsed_body(request=request) - body: Final = {**cached_body, "model": model} if cached_body is not None else request_data + body: Final = ( + {**_JSON_REQUEST_BODY.validate_python(cached_body), "model": model} if cached_body is not None else request_data + ) safe_set_request_parsed_body(request=request, parsed_body=body) request._json = body # pyright: ignore[reportPrivateUsage] # Starlette JSON cache - request._body = orjson.dumps(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache + request._body = _dumps_json_body(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache def safe_get_request_headers(request: Request | None) -> dict: diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index cd5820fa2fd..3ec95dfdcdc 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import orjson import pytest +from pydantic import TypeAdapter from fastapi.testclient import TestClient from starlette.datastructures import FormData from starlette.requests import Request @@ -104,6 +105,97 @@ async def test_read_request_body_marks_body_received_for_binary_and_form_bodies( ] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + b'{"maximum": 12345678901234567890123}', + b'{"minimum": -9223372036854775809}', + b'{"maximum": 18446744073709551616}', + b'{"maximum": 18446744073709551615}', + b'{"minimum": -9223372036854775808}', + b'{"const": 9007199254740993}', + b'{"nested": {"list": [12345678901234567890123, {"deep": -9223372036854775809}]}}', + b'{"maximum": ' + b"9" * 300 + b"}", + b'{"maximum": ' + b"9" * 400 + b"}", + rb'{"text": "escaped \"quote\" and backslash \\", "maximum": 18446744073709551616}', + ], +) +@pytest.mark.parametrize("content_type", ["application/json", "application/x-protobuf"]) +async def test_read_request_body_keeps_large_integers_exact(body: bytes, content_type: str) -> None: + request: Final = _starlette_request(body, content_type) + parsed: Final = await read_request_body(request) + + assert parsed == json.loads(body) + assert json.dumps(parsed).encode() == json.dumps(json.loads(body)).encode() + assert await read_request_body(request) == parsed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("padding", [0, 100_000]) +@pytest.mark.parametrize("content_type", ["application/json", "application/x-protobuf"]) +async def test_model_rewrite_keeps_large_integers_in_every_body_reader(padding: int, content_type: str) -> None: + body: Final = { + "model": "integer-alias", + "input": [12345678901234567890123, -9223372036854775809, {"nested": 10**400}], + "text": "escaped \"quote\" and backslash \\" + "x" * padding, + } + request: Final = _starlette_request(json.dumps(body).encode(), content_type) + parsed: Final = dict(TypeAdapter(Mapping[str, object]).validate_python(await read_request_body(request))) + + http_parsing_utils.rewrite_request_model(parsed, request, "integer-target") + + expected: Final = {**body, "model": "integer-target"} + assert parsed == expected + assert await read_request_body(request) == expected + assert await request.json() == expected + assert json.loads(await request.body()) == expected + + +@pytest.mark.asyncio +async def test_read_request_body_does_not_turn_floats_or_digit_strings_into_integers() -> None: + body: Final = ( + b'{"model": "glm-5.3", "content": "id 12345678901234567890123", ' + b'"score": 1.234567890123456789, "rate": 12345678901234567890e0, "count": 42}' + ) + parsed: Final = await read_request_body(_starlette_request(body, "application/json")) + + assert parsed == { + "model": "glm-5.3", + "content": "id 12345678901234567890123", + "score": 1.234567890123456789, + "rate": 12345678901234567890e0, + "count": 42, + } + assert isinstance(parsed["score"], float) + assert isinstance(parsed["rate"], float) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + b'{"maximum": 18446744073709551616,}', + b'{"maximum": 18446744073709551616, "text": "\\ud800"}', + b'{"maximum": 18446744073709551616, "text": "\\udc00"}', + b'{"text": "' + b'\\"' * 100 + b"18446744073709551616\\", + ], +) +async def test_large_integer_does_not_bypass_invalid_json_rejection(body: bytes) -> None: + with pytest.raises(ProxyException) as error: + await read_request_body(_starlette_request(body, "application/json")) + + assert error.value.code == "400" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("number", [b"NaN", b"Infinity", b"-Infinity", b"1e400"]) +async def test_large_integer_does_not_make_nonfinite_binary_json_valid(number: bytes) -> None: + body: Final = b'{"maximum": 18446744073709551616, "other": ' + number + b"}" + + assert await read_request_body(_starlette_request(body, "application/x-protobuf")) == {} + + @pytest.mark.asyncio async def test_read_raw_json_body_returns_the_bytes_the_parsed_body_came_from(): body = b'{"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]}' From 72f1ac4a2795e160d1044fb963041aec3111ff70 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 24 Oct 2026 20:57:49 -0700 Subject: [PATCH 42/43] revert(proxy): revert large-integer request parsing (#48289) Reverts d3b70c42ed1599727669ea7883c3971008a9113d because the JSON scanner causes excessive memory use and blocks concurrent requests on escape-heavy payloads --- .../proxy/common_utils/http_parsing_utils.py | 64 ++----------- .../common_utils/test_http_parsing_utils.py | 92 ------------------- 2 files changed, 6 insertions(+), 150 deletions(-) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 9a3aed0b67b..a38769ac41b 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -1,4 +1,3 @@ -import asyncio import json import re from collections.abc import Collection, Mapping @@ -32,6 +31,7 @@ _BINARY_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobu _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) _DECODED_JSON: Final = TypeAdapter(object) +_JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) def resolve_inference_model( @@ -162,54 +162,10 @@ def coerce_numeric_form_fields( } -_ORJSON_INT_MIN: Final = -(2**63) -_ORJSON_INT_MAX: Final = 2**64 - 1 -_JSON_PARSE_THREAD_THRESHOLD: Final = 64 * 1024 -_JSON_LONG_DIGITS: Final = re.compile(rb"[0-9]{19}") -_JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) -_JSON_BODY_BYTES: Final = TypeAdapter(bytes) -# Consume unterminated strings through EOF to avoid rescanning every escaped quote in malformed bodies. -_JSON_STRING_OR_NUMBER: Final = re.compile( - rb'"[^"\\]*(?:\\.[^"\\]*)*(?:"|\\?$)|(-?[0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?)', re.DOTALL -) - - -def _has_int_outside_orjson_range(body: bytes) -> bool: - if _JSON_LONG_DIGITS.search(body) is None: - return False - for token in _JSON_STRING_OR_NUMBER.finditer(body): - if ( - (number := token.group(1)) is None - or len(number) < 19 - or any(marker in number for marker in (b".", b"e", b"E")) - ): - continue - if len(number) > 20 or not _ORJSON_INT_MIN <= int(number) <= _ORJSON_INT_MAX: - return True - return False - - -def _loads_json_body(body: bytes) -> Mapping[str, object]: - if not _has_int_outside_orjson_range(body): - return _JSON_REQUEST_BODY.validate_python(orjson.loads(body)) - try: - parsed: Final = _JSON_REQUEST_BODY.validate_python(json.loads(body)) - json.dumps(parsed, ensure_ascii=False, allow_nan=False).encode("utf-8") - except ValueError: - return _JSON_REQUEST_BODY.validate_python(orjson.loads(body)) - return parsed - - -async def _loads_json_body_async(body: bytes) -> Mapping[str, object]: - if len(body) > _JSON_PARSE_THREAD_THRESHOLD: - return await asyncio.to_thread(_loads_json_body, body) - return _loads_json_body(body) - - -async def _parse_binary_body(body: bytes) -> dict: +def _parse_binary_body(body: bytes) -> dict: """JSON sent under a binary content type still parses; real binary (protobuf) carries no params -> {}.""" try: - parsed: Final = _DECODED_JSON.validate_python(await _loads_json_body_async(body)) + parsed: Final = _DECODED_JSON.validate_python(orjson.loads(body)) if isinstance(parsed, dict): return parsed except orjson.JSONDecodeError: @@ -281,7 +237,7 @@ async def read_request_body(request: Request | None) -> dict: if _normalize_media_type(content_type) in _BINARY_CONTENT_TYPES: binary_body: Final = await request.body() _mark_body_received(len(binary_body)) - parsed_body = await _parse_binary_body(binary_body) + parsed_body = _parse_binary_body(binary_body) elif _is_form_content_type(content_type): try: form_data: Final = await request.form() @@ -317,7 +273,7 @@ async def read_request_body(request: Request | None) -> dict: parsed_body = {} else: try: - parsed_body = dict(await _loads_json_body_async(body)) + parsed_body = dict(_JSON_REQUEST_BODY.validate_python(orjson.loads(body))) except orjson.JSONDecodeError as e: # The surrogate-repair fallback below runs two full-body re.sub # passes, which block the event loop on multi-MB malformed bodies. @@ -448,14 +404,6 @@ def detach_request_parsed_body( return client_body -def _dumps_json_body(body: Mapping[str, object]) -> bytes: - try: - encoded: Final = _JSON_BODY_BYTES.validate_python(orjson.dumps(body)) - except TypeError: - return json.dumps(body, separators=(",", ":")).encode("utf-8") - return encoded - - def rewrite_request_model( request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader request: Request | None, @@ -473,7 +421,7 @@ def rewrite_request_model( ) safe_set_request_parsed_body(request=request, parsed_body=body) request._json = body # pyright: ignore[reportPrivateUsage] # Starlette JSON cache - request._body = _dumps_json_body(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache + request._body = orjson.dumps(body) # pyright: ignore[reportPrivateUsage] # Starlette body cache def safe_get_request_headers(request: Request | None) -> dict: diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index 3ec95dfdcdc..cd5820fa2fd 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -7,7 +7,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import orjson import pytest -from pydantic import TypeAdapter from fastapi.testclient import TestClient from starlette.datastructures import FormData from starlette.requests import Request @@ -105,97 +104,6 @@ async def test_read_request_body_marks_body_received_for_binary_and_form_bodies( ] -@pytest.mark.asyncio -@pytest.mark.parametrize( - "body", - [ - b'{"maximum": 12345678901234567890123}', - b'{"minimum": -9223372036854775809}', - b'{"maximum": 18446744073709551616}', - b'{"maximum": 18446744073709551615}', - b'{"minimum": -9223372036854775808}', - b'{"const": 9007199254740993}', - b'{"nested": {"list": [12345678901234567890123, {"deep": -9223372036854775809}]}}', - b'{"maximum": ' + b"9" * 300 + b"}", - b'{"maximum": ' + b"9" * 400 + b"}", - rb'{"text": "escaped \"quote\" and backslash \\", "maximum": 18446744073709551616}', - ], -) -@pytest.mark.parametrize("content_type", ["application/json", "application/x-protobuf"]) -async def test_read_request_body_keeps_large_integers_exact(body: bytes, content_type: str) -> None: - request: Final = _starlette_request(body, content_type) - parsed: Final = await read_request_body(request) - - assert parsed == json.loads(body) - assert json.dumps(parsed).encode() == json.dumps(json.loads(body)).encode() - assert await read_request_body(request) == parsed - - -@pytest.mark.asyncio -@pytest.mark.parametrize("padding", [0, 100_000]) -@pytest.mark.parametrize("content_type", ["application/json", "application/x-protobuf"]) -async def test_model_rewrite_keeps_large_integers_in_every_body_reader(padding: int, content_type: str) -> None: - body: Final = { - "model": "integer-alias", - "input": [12345678901234567890123, -9223372036854775809, {"nested": 10**400}], - "text": "escaped \"quote\" and backslash \\" + "x" * padding, - } - request: Final = _starlette_request(json.dumps(body).encode(), content_type) - parsed: Final = dict(TypeAdapter(Mapping[str, object]).validate_python(await read_request_body(request))) - - http_parsing_utils.rewrite_request_model(parsed, request, "integer-target") - - expected: Final = {**body, "model": "integer-target"} - assert parsed == expected - assert await read_request_body(request) == expected - assert await request.json() == expected - assert json.loads(await request.body()) == expected - - -@pytest.mark.asyncio -async def test_read_request_body_does_not_turn_floats_or_digit_strings_into_integers() -> None: - body: Final = ( - b'{"model": "glm-5.3", "content": "id 12345678901234567890123", ' - b'"score": 1.234567890123456789, "rate": 12345678901234567890e0, "count": 42}' - ) - parsed: Final = await read_request_body(_starlette_request(body, "application/json")) - - assert parsed == { - "model": "glm-5.3", - "content": "id 12345678901234567890123", - "score": 1.234567890123456789, - "rate": 12345678901234567890e0, - "count": 42, - } - assert isinstance(parsed["score"], float) - assert isinstance(parsed["rate"], float) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "body", - [ - b'{"maximum": 18446744073709551616,}', - b'{"maximum": 18446744073709551616, "text": "\\ud800"}', - b'{"maximum": 18446744073709551616, "text": "\\udc00"}', - b'{"text": "' + b'\\"' * 100 + b"18446744073709551616\\", - ], -) -async def test_large_integer_does_not_bypass_invalid_json_rejection(body: bytes) -> None: - with pytest.raises(ProxyException) as error: - await read_request_body(_starlette_request(body, "application/json")) - - assert error.value.code == "400" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("number", [b"NaN", b"Infinity", b"-Infinity", b"1e400"]) -async def test_large_integer_does_not_make_nonfinite_binary_json_valid(number: bytes) -> None: - body: Final = b'{"maximum": 18446744073709551616, "other": ' + number + b"}" - - assert await read_request_body(_starlette_request(body, "application/x-protobuf")) == {} - - @pytest.mark.asyncio async def test_read_raw_json_body_returns_the_bytes_the_parsed_body_came_from(): body = b'{"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]}' From dddbaa803d6e9fcc862ea7a5412f771cb1ae040e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 24 Oct 2026 22:33:45 -0700 Subject: [PATCH 43/43] build(docker): refresh base images and ship only runtime files in component images (#48279) * build(docker): move to newer wolfi-base and nginx base images Bump the pinned wolfi-base digest in every Wolfi Dockerfile to 360896ef, which ships brotli 1.2.1-r1 and ca-certificates-bundle 20260909-r3. The previous digest already carried libbrotli 1.2.0-r6 as an OpenSSL dependency, and `apk add --no-cache` never upgrades a package the base already has Move the UI image to nginx 1.31.8-alpine3.24, whose CA bundle is the 20261020 Mozilla set, and upgrade brotli-libs in its runtime stage since no nginx digest ships 1.2.1 yet * build(docker): copy only runtime files into the gateway, backend and migrations images The component runtime stages copied the builder's whole /app, which is the entire repository plus build output: six lockfiles, helm, terraform, docs, .env.example and the Rust build tree. Measured at runtime with an audit hook on a booted container, the gateway reads only gateway/, litellm/ and enterprise/ from /app, the backend only backend/, litellm/ and enterprise/, and the migrations job only litellm/ and run.py. The pyproject.toml read is prisma-python's optional [tool.prisma] lookup, which ours does not define Copy just those, plus .venv, the entrypoint script and schema.prisma for the gateway and backend. PYTHONPATH=/app and the in-Dockerfile asserts are unchanged --- Dockerfile | 4 ++-- backend/Dockerfile | 11 ++++++++--- docker/Dockerfile.database | 4 ++-- docker/Dockerfile.non_root | 4 ++-- gateway/Dockerfile | 11 ++++++++--- migrations/Dockerfile | 8 +++++--- ui/Dockerfile | 4 +++- 7 files changed, 30 insertions(+), 16 deletions(-) diff --git a/Dockerfile b/Dockerfile index 713d7f8f5ab..4cef63c8d71 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 diff --git a/backend/Dockerfile b/backend/Dockerfile index 46d65c432bb..1c22f914100 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -124,7 +124,12 @@ ENV HOME=/home/nonroot \ PYTHONUNBUFFERED=1 \ PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries -COPY --from=builder --chown=nonroot:nonroot /app /app +COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv +COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm +COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise +COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend +COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh +COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma COPY --from=builder /opt/prisma /opt/prisma RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 69d9c26b2a8..20f810686a3 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 54bf8c807aa..90bcc523d9e 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,8 +1,8 @@ # syntax=docker/dockerfile:1.7 # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 390a5f5666f..69068cfc8a5 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) ARG PGBOUNCER_VERSION=1.26.0 @@ -140,7 +140,12 @@ ENV HOME=/home/nonroot \ PYTHONUNBUFFERED=1 \ PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries -COPY --from=builder --chown=nonroot:nonroot /app /app +COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv +COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm +COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise +COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway +COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh +COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma COPY --from=builder /opt/prisma /opt/prisma COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 84ea67fcc5f..2b1fd029143 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -137,7 +137,9 @@ ENV HOME=/home/nonroot \ PRISMA_CLI_QUERY_ENGINE_TYPE=binary \ PRISMA_OFFLINE_MODE=true -COPY --from=builder --chown=nonroot:nonroot /app /app +COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv +COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm +COPY --from=builder --chown=nonroot:nonroot /app/run.py /app/run.py COPY --from=builder /opt/prisma /opt/prisma RUN chmod -R a+rX /opt/prisma && \ diff --git a/ui/Dockerfile b/ui/Dockerfile index 4b6b0c1581a..d78a155b874 100644 --- a/ui/Dockerfile +++ b/ui/Dockerfile @@ -3,7 +3,7 @@ # UI container — Next.js static export served by nginx. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 -ARG NGINX_VERSION=1.31.5-alpine3.24@sha256:34f40471dea485273c5e2a04dd5e97a682332ceb4a9adecd67de450dcb2fb390 +ARG NGINX_VERSION=1.31.8-alpine3.24@sha256:1ef7ca94d7879df82421428ce13bd1eaa412404953859cf7a5b6512c25e67a60 # ---------- builder ---------- FROM ${UI_BUILD_IMAGE} AS builder @@ -27,6 +27,8 @@ RUN npm run build # ---------- runtime ---------- FROM nginx:${NGINX_VERSION} AS runtime +RUN apk upgrade --no-cache brotli-libs + # Drop the upstream default :80 server; we own the config. RUN rm -f /etc/nginx/conf.d/default.conf