diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 0f449f01ec9..1b60f986ca4 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -96,6 +96,7 @@ ARRAY_KEYS: dict[str, JsonSchema] = { "output_cost_per_token": NONNEG_NUMBER, "output_cost_per_reasoning_token": NONNEG_NUMBER, "cache_read_input_token_cost": NONNEG_NUMBER, + "cache_creation_input_token_cost": NONNEG_NUMBER, "input_cost_per_query": NONNEG_NUMBER, }, "additionalProperties": False, diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index e73b887ae0a..baf522c0bb1 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -5,7 +5,7 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.llm_cost_calc.utils import _parse_prompt_tokens_details +from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch from litellm.types.utils import CallTypes, ModelInfo, Usage from litellm.utils import token_counter @@ -101,7 +101,7 @@ def _iter_successful_output_line_stats( continue response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider) usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider) - prompt_details = _parse_prompt_tokens_details(usage) + prompt_details = parse_prompt_tokens_details(usage) raw_model = response_body.get("model") response_model = raw_model if isinstance(raw_model, str) and raw_model else None if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"): diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6b6653c5646..b37ff865c65 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -26,11 +26,11 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _generic_cost_per_character, _get_regional_uplift_multiplier, _get_service_tier_cost_key, - _parse_prompt_tokens_details, calculate_cost_component, generic_cost_per_token, get_billable_input_tokens, get_token_type_cost_breakdown, + parse_prompt_tokens_details, select_cost_metric_for_model, ) from litellm.llms.anthropic.cost_calculation import ( @@ -645,7 +645,11 @@ def cost_per_token( else: model_info: Final = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) - if (model_info.get("input_cost_per_token") or 0.0) > 0 or (model_info.get("output_cost_per_token") or 0.0) > 0: + if ( + (model_info.get("input_cost_per_token") or 0.0) > 0 + or (model_info.get("output_cost_per_token") or 0.0) > 0 + or model_info.get("tiered_pricing") is not None + ): return generic_cost_per_token( model=model, usage=usage_block, @@ -2159,7 +2163,7 @@ def batch_cost_calculator( if input_cost_per_token_batches: total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches elif input_cost_per_token: - details: Final = _parse_prompt_tokens_details(usage) + details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = details["cache_hit_tokens"] cache_creation_tokens: Final = details["cache_creation_tokens"] 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 fb0f130a6cf..9bcc2b1743c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py @@ -1,5 +1,5 @@ """ -Provider-neutral graduated tiered pricing calculation. +Provider-neutral tiered pricing calculation. Shared by provider cost calculators (e.g. Dashscope) and the proxy budget reservation logic so neither has to depend on the other. @@ -25,80 +25,6 @@ def _coerce_cost_per_token(value: float | str | None) -> float: return float(value) -def calculate_tiered_cost( - tokens: int, - tiered_pricing: list[dict], - cost_key: str, - fallback_cost_key: str | None = None, -) -> float: - """ - Calculate cost for a given number of tokens based on a true tiered pricing structure. - - This function iterates through sorted pricing tiers, calculates the cost for the - number of tokens that fall into each tier's range, and sums them up to get the total cost. - - Args: - tokens (int): The total number of tokens to calculate the cost for. - tiered_pricing (List[dict]): A list of dictionaries, where each dictionary - represents a pricing tier. - cost_key (str): The key in the tier dictionary that holds the per-token cost - (e.g., 'input_cost_per_token'). - fallback_cost_key (Optional[str], optional): A fallback key to use if the - primary `cost_key` is not found in a tier. Defaults to None. - - Returns: - float: The total calculated cost for the given tokens. - - Example: - >>> tiered_pricing = [ - ... {"range": [0, 100000], "input_cost_per_token": 0.0001}, - ... {"range": [100000, 500000], "input_cost_per_token": 0.00005}, - ... ] - - Calculating cost for 150,000 tokens: - (100,000 * 0.0001) + (50,000 * 0.00005) = $12.5 - """ - if not tiered_pricing or tokens <= 0: - return 0.0 - - total_cost = 0.0 - tokens_processed = 0 - - sorted_tiers: Final = sorted(tiered_pricing, key=lambda x: x.get("range", [0, 0])[0]) - - for tier in sorted_tiers: - if tokens_processed >= tokens: - break - - tier_range = tier.get("range", []) - if len(tier_range) != 2: - continue - - range_start, range_end = tier_range - - if tokens <= range_start: - continue - - tier_start = max(range_start, tokens_processed) - tier_end = min(range_end, tokens) - - if tier_end > tier_start: - tokens_in_tier = tier_end - tier_start - cost_per_token = tier.get(cost_key) or tier.get(fallback_cost_key, 0) - total_cost += tokens_in_tier * _coerce_cost_per_token(cost_per_token) - tokens_processed = tier_end - - # After loop, check if any tokens remain (i.e., tokens > highest tier's end range) - # and charge them at the last tier's rate. - if tokens_processed < tokens and sorted_tiers: - last_tier: Final = sorted_tiers[-1] - remaining_tokens: Final = tokens - tokens_processed - cost_per_token = last_tier.get(cost_key) or last_tier.get(fallback_cost_key, 0) - total_cost += remaining_tokens * _coerce_cost_per_token(cost_per_token) - - return total_cost - - def select_tier_for_input( tiered_pricing: list[dict], input_tokens: int, @@ -134,6 +60,12 @@ def tier_rate( cost_key: str, fallback_cost_key: str | None = None, ) -> float: - """Read a per-token rate from a tier, coercing YAML string costs to float.""" - raw: Final = tier.get(cost_key) or tier.get(fallback_cost_key, 0) - return _coerce_cost_per_token(raw) + """Read a per-token rate from a tier, coercing YAML string costs to float. + + A rate that is explicitly present wins over the fallback, an explicit zero + included, so a tier can declare a token type free. + """ + primary: Final = tier.get(cost_key) + if primary is not None: + return _coerce_cost_per_token(primary) + 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 38bdf89981e..9d6ad8b6e39 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -8,6 +8,10 @@ from typing import Any, Final, Literal, TypedDict, cast import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( + select_tier_for_input, + tier_rate, +) from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, @@ -95,7 +99,7 @@ def get_billable_input_tokens(usage: Usage) -> int: Returns the number of billable input tokens. Subtracts cached tokens from prompt tokens if applicable. """ - details: Final = _parse_prompt_tokens_details(usage) + details: Final = parse_prompt_tokens_details(usage) return usage.prompt_tokens - details["cache_hit_tokens"] @@ -207,6 +211,57 @@ def _parse_above_token_threshold(key: str) -> float: return float(threshold_str.replace("k", "")) * (1000 if "k" in threshold_str else 1) +def _select_priced_tier(model_info: ModelInfo, usage: Usage) -> dict | None: + tiered_pricing: Final = model_info.get("tiered_pricing") + if not isinstance(tiered_pricing, list) or not tiered_pricing: + return None + + tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=usage.prompt_tokens) + if tier is None or "input_cost_per_token" not in tier: + return None + return tier + + +def _get_tiered_reasoning_rate(model_info: ModelInfo, usage: Usage) -> float | None: + tier: Final = _select_priced_tier(model_info=model_info, usage=usage) + if tier is None: + return None + if "output_cost_per_reasoning_token" not in tier and "output_cost_per_token" not in tier: + return None + 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: + """ + Resolve the base rates from a model's ``tiered_pricing`` table, if it has one. + + Tiered pricing is all-or-nothing: one tier is picked from the request's input tokens + and every token of the request is billed at that tier's rate. Rates the tier does not + declare fall back to the tier's input rate, so a request never mixes tiers. + + An output rate is the exception: a tier table that spells out only input rates would + otherwise serve every completion for free, so the model's own output rate stands in. + """ + tier: Final = _select_priced_tier(model_info=model_info, usage=usage) + 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", "cache_creation_input_token_cost") + or cache_creation_cost, + tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + ) + + def _get_token_base_cost( model_info: ModelInfo, usage: Usage, @@ -226,6 +281,10 @@ def _get_token_base_cost( 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 tiered_base_costs + # Get service tier aware cost keys input_cost_key: Final = _get_service_tier_cost_key("input_cost_per_token", service_tier) output_cost_key: Final = _get_service_tier_cost_key("output_cost_per_token", service_tier) @@ -470,7 +529,7 @@ class PromptTokensDetailsResult(TypedDict): audio_length_seconds: float -def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: +def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: cache_hit_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "cached_tokens", 0)) or 0 cache_creation_tokens: Final = ( cast( @@ -540,7 +599,7 @@ class CompletionTokensDetailsResult(TypedDict): video_tokens: int -def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: +def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: audio_tokens: Final = ( cast( int | None, @@ -777,7 +836,7 @@ def generic_cost_per_token( audio_length_seconds=0.0, ) if usage.prompt_tokens_details: - prompt_tokens_details = _parse_prompt_tokens_details(usage) + prompt_tokens_details = parse_prompt_tokens_details(usage) ## EDGE CASE - text tokens not set or includes cached tokens (double-counting) ## Some providers (like xAI) report text_tokens = prompt_tokens (including cached) @@ -832,7 +891,7 @@ def generic_cost_per_token( video_tokens = 0 is_text_tokens_total = False if usage.completion_tokens_details is not None: - completion_tokens_details: Final = _parse_completion_tokens_details(usage) + completion_tokens_details: Final = parse_completion_tokens_details(usage) audio_tokens = completion_tokens_details["audio_tokens"] text_tokens = completion_tokens_details["text_tokens"] reasoning_tokens = completion_tokens_details["reasoning_tokens"] @@ -869,10 +928,15 @@ def generic_cost_per_token( ## REASONING COST if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0: - _output_cost_per_reasoning_token = _resolve_reasoning_token_cost( - model_info=model_info, - service_tier=service_tier, - completion_base_cost=completion_base_cost, + tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) + _output_cost_per_reasoning_token = ( + tiered_reasoning_rate + if tiered_reasoning_rate is not None + else _resolve_reasoning_token_cost( + model_info=model_info, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + ) ) completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token @@ -952,26 +1016,29 @@ def get_token_type_cost_breakdown( ) reasoning_tokens = ( - _parse_completion_tokens_details(usage)["reasoning_tokens"] - if usage.completion_tokens_details is not None - else 0 + parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 ) if not reasoning_tokens: reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - # Reasoning is billed at the explicit per-reasoning-token rate when the model - # defines one, otherwise at the standard output-token rate - this mirrors how the - # total completion cost is computed, so the breakdown can never diverge from it. - reasoning_rate = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None) - if reasoning_rate is None: - reasoning_rate = completion_base_cost + # Reasoning is billed at the selected tier's reasoning rate for tiered models, + # else at the explicit per-reasoning-token rate when the model defines one, + # otherwise at the standard output-token rate - this mirrors how the total + # completion cost is computed, so the breakdown can never diverge from it. + tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) + flat_reasoning_rate: Final = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None) + reasoning_rate: Final = ( + tiered_reasoning_rate + if tiered_reasoning_rate is not None + else (flat_reasoning_rate if flat_reasoning_rate is not None else completion_base_cost) + ) reasoning_cost = float(reasoning_tokens) * reasoning_rate cache_read_tokens = 0 cache_creation_tokens = 0 cache_creation_token_details: CacheCreationTokenDetails | None = None if usage.prompt_tokens_details is not None: - prompt_tokens_details: Final = _parse_prompt_tokens_details(usage) + prompt_tokens_details: Final = parse_prompt_tokens_details(usage) cache_read_tokens = prompt_tokens_details["cache_hit_tokens"] cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"] cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"] diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index e792f69622c..7bb3e0294f0 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -10,10 +10,10 @@ from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, _get_web_search_requests, - _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, get_provider_specific_geo_multiplier, + parse_prompt_tokens_details, ) if TYPE_CHECKING: @@ -33,7 +33,7 @@ def _compute_cache_only_cost(model_info: "ModelInfo", usage: "Usage", service_ti if usage.prompt_tokens_details is None: return 0.0 - prompt_tokens_details: Final = _parse_prompt_tokens_details(usage) + prompt_tokens_details: Final = parse_prompt_tokens_details(usage) ( _, _, diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 22a0d38d598..771ce140f66 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -1,108 +1,111 @@ """ Cost calculator for Dashscope Chat models. -Handles tiered pricing and prompt caching scenarios. +Alibaba Model Studio tiered pricing is all-or-nothing: the tier is picked from the +total input tokens of a single request, and every token of that request (input, +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 dataclasses import dataclass from typing import Final -from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import calculate_tiered_cost +from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + parse_completion_tokens_details, + parse_prompt_tokens_details, +) from litellm.types.utils import ModelInfo, Usage from litellm.utils import get_model_info -@dataclass +@dataclass(frozen=True, slots=True) class TokenBreakdown: - """Token breakdown for cost calculation.""" - text_tokens: int cached_tokens: int + cache_creation_tokens: int completion_tokens: int reasoning_tokens: int + @property + def total_input_tokens(self) -> int: + return self.text_tokens + self.cached_tokens + self.cache_creation_tokens + def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: - """Extract token counts from usage, handling cached and reasoning tokens.""" - cached_tokens = 0 - if usage.prompt_tokens_details and hasattr(usage.prompt_tokens_details, "cached_tokens"): - cached_tokens = usage.prompt_tokens_details.cached_tokens or 0 + prompt_details: Final = parse_prompt_tokens_details(usage) + cached_tokens: Final = prompt_details["cache_hit_tokens"] + cache_creation_tokens: Final = prompt_details["cache_creation_tokens"] + text_tokens: Final = max(usage.prompt_tokens - cached_tokens - cache_creation_tokens, 0) - text_tokens: Final = usage.prompt_tokens - cached_tokens + reasoning_tokens: Final = parse_completion_tokens_details(usage)["reasoning_tokens"] + completion_tokens: Final = max((usage.completion_tokens or 0) - reasoning_tokens, 0) - reasoning_tokens = 0 - if ( - hasattr(usage, "completion_tokens_details") - and usage.completion_tokens_details - and hasattr(usage.completion_tokens_details, "reasoning_tokens") - ): - reasoning_tokens = usage.completion_tokens_details.reasoning_tokens or 0 + return TokenBreakdown( + text_tokens=text_tokens, + cached_tokens=cached_tokens, + cache_creation_tokens=cache_creation_tokens, + completion_tokens=completion_tokens, + reasoning_tokens=reasoning_tokens, + ) - completion_tokens: Final = (usage.completion_tokens or 0) - reasoning_tokens - return TokenBreakdown(text_tokens, cached_tokens, completion_tokens, reasoning_tokens) +def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) -> float: + value: Final = model_info.get(cost_key) + if value is None: + return float(model_info.get(fallback_cost_key) or 0.0) + return float(value) def _calculate_prompt_cost( breakdown: TokenBreakdown, model_info: ModelInfo, - tiered_pricing: list[dict] | None, + tier: dict | None, ) -> float: - """Calculate total prompt cost including cached tokens.""" - if tiered_pricing: - text_cost: Final = calculate_tiered_cost( - tokens=breakdown.text_tokens, - tiered_pricing=tiered_pricing, - cost_key="input_cost_per_token", + if tier is not None: + return ( + (breakdown.text_tokens * tier_rate(tier, "input_cost_per_token")) + + (breakdown.cached_tokens * tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token")) + + ( + breakdown.cache_creation_tokens + * tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") + ) ) - cache_cost = calculate_tiered_cost( - tokens=breakdown.cached_tokens, - tiered_pricing=tiered_pricing, - cost_key="cache_read_input_token_cost", - fallback_cost_key="input_cost_per_token", - ) - return text_cost + cache_cost input_cost: Final = float(model_info.get("input_cost_per_token") or 0.0) + cache_read_cost: Final = _flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token") + cache_creation_cost: Final = _flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token") - # For cache_cost, first try the specific key, then fall back to input_cost. - cache_cost_val: Final = model_info.get("cache_read_input_token_cost") - if cache_cost_val is None: - cache_cost = input_cost - else: - cache_cost = float(cache_cost_val) - - return (breakdown.text_tokens * input_cost) + (breakdown.cached_tokens * cache_cost) + return ( + (breakdown.text_tokens * input_cost) + + (breakdown.cached_tokens * cache_read_cost) + + (breakdown.cache_creation_tokens * cache_creation_cost) + ) def _calculate_completion_cost( breakdown: TokenBreakdown, model_info: ModelInfo, - tiered_pricing: list[dict] | None, + tier: dict | None, ) -> float: - """Calculate total completion cost including reasoning tokens.""" - if tiered_pricing: - completion_cost: Final = calculate_tiered_cost( - tokens=breakdown.completion_tokens, - tiered_pricing=tiered_pricing, - cost_key="output_cost_per_token", - ) - reasoning_cost = calculate_tiered_cost( - tokens=breakdown.reasoning_tokens, - tiered_pricing=tiered_pricing, - cost_key="output_cost_per_reasoning_token", - fallback_cost_key="output_cost_per_token", - ) - return completion_cost + reasoning_cost - - output_cost: Final = float(model_info.get("output_cost_per_token") or 0.0) - - # For reasoning_cost, first try the specific key, then fall back to output_cost. - reasoning_cost_val: Final = model_info.get("output_cost_per_reasoning_token") - if reasoning_cost_val is None: - reasoning_cost = output_cost - else: - reasoning_cost = float(reasoning_cost_val) + # 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 + tier_declares_output: Final = tier is not None and "output_cost_per_token" in tier + output_cost: Final = ( + tier_rate(tier, "output_cost_per_token") + if tier_declares_output + else float(model_info.get("output_cost_per_token") or 0.0) + ) + tier_declares_reasoning: Final = tier is not None and "output_cost_per_reasoning_token" in tier + model_reasoning_rate: Final = None if tier_declares_output else model_info.get("output_cost_per_reasoning_token") + reasoning_cost: Final = ( + tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") + if tier_declares_reasoning + else float(model_reasoning_rate) + if model_reasoning_rate is not None + else output_cost + ) return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) @@ -122,11 +125,15 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ model_info: Final = get_model_info(model=model, custom_llm_provider="dashscope") breakdown: Final = _extract_token_breakdown(usage) - tiered_pricing = model_info.get("tiered_pricing") if isinstance(model_info.get("tiered_pricing"), list) else None - - prompt_cost = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing) - completion_cost: Final = _calculate_completion_cost( - breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing + raw_tiers: Final = model_info.get("tiered_pricing") + tiered_pricing: Final = raw_tiers if isinstance(raw_tiers, list) else None + tier: Final = ( + select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=breakdown.total_input_tokens) + if tiered_pricing + else None ) + prompt_cost: Final = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier) + completion_cost: Final = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier) + return prompt_cost, completion_cost diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3db94211032..b7f91bfba0d 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -7,6 +7,7 @@ import re import time from collections.abc import Callable, Iterable, Iterator, Mapping from typing import Any, Final, TypedDict +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -43,6 +44,9 @@ from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + transform_openai_input_gemini_embed_content, +) from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, @@ -54,14 +58,28 @@ from litellm.types.llms.openai import ( OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + LlmProviders, + ModelResponse, + Usage, +) from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +_VERTEX_BATCH_KEY_FIELD: Final = "key" +_MANAGED_GCS_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( + ("outputDimensionality", "output_dimensionality"), + ("taskType", "task_type"), + ("title", "title"), +) +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN: Final = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") class _GcsObjectMetadataJson(TypedDict, total=False): @@ -164,8 +182,26 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: objec labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk -def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> str: +def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, object]) -> str: + """ + Resolve the OpenAI `custom_id` for a Vertex batch output row. + + Embedding rows carry it in the top-level `key` field that Vertex echoes back; + `generateContent` rows have no such field, so it is smuggled through request + labels instead (see `_set_litellm_batch_custom_id_labels`). + """ + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is not None: + return unquote(str(key)) + request_data = vertex_output_row.get("request") + labels = request_data.get("labels") if isinstance(request_data, Mapping) else None + return _get_litellm_batch_custom_id_from_labels(labels) + + +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object] | None) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" + if not labels: + return "unknown" raw: Final = labels.get("litellm_custom_id_raw") if raw: raw_chunks: Final = [str(raw)] @@ -182,17 +218,311 @@ def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> st return str(labels.get("litellm_custom_id", "unknown")) -def _openai_batch_jsonl_entry_to_vertex_wrapped_request( +def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool: + """ + Whether a Vertex batch output row came from an `EmbedContentRequest`. + + Successful rows hold the vector under `response.embedding.values`; failed rows only + carry `status`, so they are recognized from the singular `content` that the + embeddings request shape echoes back. + """ + if "request" not in vertex_output_row: + return False + response = vertex_output_row.get("response") + if isinstance(response, dict) and isinstance(response.get("embedding"), dict): + return True + request_data = vertex_output_row.get("request") + return bool(vertex_output_row.get("status")) and isinstance(request_data, dict) and "content" in request_data + + +def _openai_batch_output_row( + custom_id: str, + body: Mapping[str, Any] | None = None, + error_code: str | None = None, + error_message: str = "", +) -> _OpenAIBatchOutputRow: + """ + One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set + `response` to null and populate `error` instead. + """ + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": None + if body is None + else { + "status_code": 200, + "request_id": body.get("id", ""), + "body": body, + }, + "error": None if error_code is None else {"code": error_code, "message": error_message}, + } + + +def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int, int]: + """ + Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch + output row. + + A `/v1/embeddings` entry whose `input` is an array fans out into one Vertex row per + element, tagged `#/` (see + `_vertex_batch_embeddings_key`), so the rows can be reassembled into a single OpenAI + response. + """ + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is None: + return _get_litellm_batch_custom_id(vertex_output_row), 0, 1 + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(str(key)) + if match is None: + return unquote(str(key)), 0, 1 + return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) + + +def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: + """ + Prompt tokens billed for one Vertex Gemini Embedding batch row. + + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. + """ + usage_metadata = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return int(usage_metadata.get("promptTokenCount") or 0) + return int(vertex_response.get("tokenCount") or 0) + + +def _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id: str, + vertex_output_rows: tuple[Mapping[str, Any], ...], + element_indices: tuple[int, ...], + element_count: int, + model: str | None, +) -> _OpenAIBatchOutputRow: + """ + Transforms the Vertex Gemini Embedding batch output rows belonging to one OpenAI + batch entry into an OpenAI batch output row holding an `/v1/embeddings` response. + + Example Vertex jsonl + {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} + + An entry that asked for several embeddings at once maps to several rows here, which + become the indexed elements of a single `data` array. One failed or missing element + fails the whole entry, since an OpenAI batch row is either a response or an error and + a partial `data` array would silently shift the remaining embeddings onto the wrong + input positions. Rows carry no `modelVersion`, so the model comes from the batch they + belong to. + """ + status = next((row["status"] for row in vertex_output_rows if row.get("status")), "") + if status: + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=status, + ) + + if element_indices != tuple(range(element_count)): + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=( + f"Vertex returned embeddings for input positions {list(element_indices)} " + f"of the {element_count} requested" + ), + ) + + responses = tuple(row["response"] for row in vertex_output_rows) + token_count = sum(_embedding_prompt_token_count(response) for response in responses) + body = EmbeddingResponse( + model=model or "", + data=[ + Embedding( + embedding=response["embedding"]["values"], + index=index, + object="embedding", + ) + for index, response in enumerate(responses) + ], + usage=Usage(prompt_tokens=token_count, total_tokens=token_count), + ).model_dump() + return _openai_batch_output_row(custom_id=custom_id, body=body) + + +def _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows: Iterable[Mapping[str, Any]], + model: str | None, +) -> tuple[_OpenAIBatchOutputRow, ...]: + """ + Transforms a whole Vertex Gemini Embedding batch output into OpenAI batch output + rows, one per OpenAI batch entry, in the order the entries first appear. + + Rows are grouped rather than mapped one to one because a single entry can fan out + into several Vertex rows, and Vertex returns them in arbitrary order. + """ + keyed_rows = tuple((_split_vertex_batch_key(row), row) for row in vertex_output_rows) + grouped_rows = { + custom_id: tuple(group) + for custom_id, group in itertools.groupby(sorted(keyed_rows, key=lambda kr: kr[0]), key=lambda kr: kr[0][0]) + } + return tuple( + _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id=custom_id, + vertex_output_rows=tuple(row for _, row in grouped_rows[custom_id]), + element_indices=tuple(index for (_, index, _), _ in grouped_rows[custom_id]), + element_count=max(total for (_, _, total), _ in grouped_rows[custom_id]), + model=model, + ) + for custom_id in dict.fromkeys(custom_id for (custom_id, _, _), _ in keyed_rows) + ) + + +def _model_from_managed_gcs_url(url: str) -> str | None: + """ + Extracts the model from a LiteLLM-managed Vertex batch GCS url. + + Batch inputs and their sibling outputs are stored under + `.../publishers/google/models//...`, which is the only place the model of an + embeddings batch output row can be recovered from; unlike `generateContent` + responses, embedding rows carry no `modelVersion`. + """ + match = _MANAGED_GCS_MODEL_PATH_PATTERN.search(unquote(url)) + return match.group(1) if match else None + + +def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: + """ + Whether an OpenAI batch JSONL line targets the embeddings endpoint. + + OpenAI puts the target route on each line's `url` (e.g. `/v1/embeddings`); Vertex + has no equivalent per-line field, so the route decides which Vertex request shape + the line has to be translated into. + """ + url = openai_entry.get("url") + if not isinstance(url, str): + return False + path = url.split("?")[0].rstrip("/") + return path == "embeddings" or path.endswith("/embeddings") + + +def _openai_embedding_input_elements( + embedding_input: GeminiEmbeddingInput, +) -> tuple[str | list[str], ...]: + """ + Split an OpenAI `input` into the elements that each get their own embedding. + + A string is one embedding, a flat array is one embedding per element, and a nested + array is one combined embedding per inner array, matching the online + `batchEmbedContents` path. + """ + if isinstance(embedding_input, list): + return tuple(embedding_input) + return (embedding_input,) + + +def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: + """ + The top-level `key` Vertex echoes back on an embeddings row. + + An entry asking for several embeddings needs several Vertex rows, so its key also + carries the element index and the group size; `_split_vertex_batch_key` reads them + back out. The `custom_id` is percent-encoded so that a customer one ending in + `#/` cannot be mistaken for that tag, which would merge two entries. + """ + encoded_custom_id = quote(custom_id, safe="") + return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" + + +def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]: + """ + One Vertex Gemini Embedding batch input row. + + The config fields live inside the `EmbedContentRequest` under their snake_case batch + names, and the OpenAI `custom_id` rides along in the top-level `key` that Vertex + echoes back. + """ + request = { + "content": embed_content_request["content"], + **{ + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM + if gemini_param in embed_content_request + }, + } + if key is None: + return {"request": request} + return {_VERTEX_BATCH_KEY_FIELD: key, "request": request} + + +def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( + openai_entry: Mapping[str, Any], +) -> tuple[Mapping[str, Any], ...]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding + batch rows, one per requested embedding. + + Example Vertex jsonl + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + + Note that `content` is singular (an `EmbedContentRequest`, not a + `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level + `key`. An `EmbedContentRequest` returns exactly one vector, so an entry whose `input` + is an array fans out into one row per element and is reassembled on the way back. + The docs put the per-row config in an `embed_content_config` sibling of `request`, + but the API rejects that key outright and fails the whole batch job, so the config + fields go inside the `EmbedContentRequest` itself. + + API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + openai_request_body = openai_entry.get("body") + if not isinstance(openai_request_body, dict): + raise TypeError( + "`body` on /v1/embeddings batch requests must be a JSON object, but was missing or not an object" + ) + embedding_input = openai_request_body.get("input") + if embedding_input is None: + raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") + + elements = _openai_embedding_input_elements(embedding_input) + if not elements: + raise ValueError("`input` on /v1/embeddings batch requests must not be empty") + + embed_content_requests = tuple( + transform_openai_input_gemini_embed_content( + input=element, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + for element in elements + ) + custom_id = openai_entry.get("custom_id") + return tuple( + _vertex_embeddings_row( + key=None + if custom_id is None + else _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), + ), + embed_content_request=embed_content_request, + ) + for index, embed_content_request in enumerate(embed_content_requests) + ) + + +def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: +) -> tuple[Mapping[str, Any], ...]: """ - Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. + Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ + if _is_embeddings_batch_entry(openai_entry): + return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) + openai_request_body: Final = openai_entry.get("body") or {} vertex_request_body: Final = _transform_request_body( messages=openai_request_body.get("messages", []), @@ -209,7 +539,7 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( vertex_request_body["labels"] = {} _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) - return {"request": vertex_request_body} + return ({"request": vertex_request_body},) def _iter_stripped_lines(raw_lines: Iterable[str | bytes]) -> Iterator[str]: @@ -312,10 +642,10 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for entry in _iter_openai_jsonl_entries(self._openai_file_content): - wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, self._map_openai_to_vertex_params) - prefix = b"" if first else b"\n" - first = False - yield prefix + json.dumps(wrapped).encode("utf-8") + for wrapped in _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params): + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(wrapped).encode("utf-8") def iter_bytes(self) -> Iterator[bytes]: return self._iter_vertex_jsonl_chunks() @@ -667,6 +997,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): transformed_content: Final = self._try_transform_vertex_batch_output_to_openai( content=content, logging_obj=logging_obj, + model=_model_from_managed_gcs_url(str(raw_response.request.url)), ) if transformed_content != content: # Create a new response with transformed content and updated Content-Length @@ -688,7 +1019,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return HttpxBinaryResponseContent(response=raw_response) def _try_transform_vertex_batch_output_to_openai( - self, content: bytes, logging_obj: LiteLLMLoggingObj | None = None + self, + content: bytes, + logging_obj: LiteLLMLoggingObj | None = None, + model: str | None = None, ) -> bytes: """ Try to transform Vertex AI batch output to OpenAI format. @@ -730,7 +1064,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row: Final = _parse_vertex_batch_output_row(first_line) - is_vertex_batch_output: Final = ( + is_vertex_batch_output: Final = _is_vertex_embeddings_batch_output_row(first_row) or ( "request" in first_row and "response" in first_row and "processed_time" in first_row @@ -763,11 +1097,23 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) + all_lines = itertools.chain((first_line,), lines) + + # Embedding rows are grouped by `custom_id` rather than transformed one at a + # time, since an entry that asked for several embeddings comes back as + # several rows, in arbitrary order. + if _is_vertex_embeddings_batch_output_row(first_row): + openai_outputs = _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows=(json.loads(line) for line in all_lines), + model=model, + ) + return b"\n".join(json.dumps(openai_output).encode("utf-8") for openai_output in openai_outputs) + # Transform each row straight into the output buffer, so peak memory # stays at ~one row plus the output. If any row fails, return the # original content unchanged. output = bytearray() - for line in itertools.chain([first_line], lines): + for line in all_lines: try: openai_output = self._transform_single_vertex_batch_output_to_openai( vertex_output=_parse_vertex_batch_output_row(line), @@ -798,25 +1144,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ - # Extract custom_id from request labels (prefer raw for OpenAI round-trip) - request_data: Final = vertex_output.get("request", {}) - labels: Final[Mapping[str, object]] = request_data.get("labels", {}) or {} - custom_id: Final = _get_litellm_batch_custom_id_from_labels(labels) + custom_id: Final = _get_litellm_batch_custom_id(vertex_output) # Check if there's an error status: Final = vertex_output.get("status", "") has_error: Final = bool(status) if has_error: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { - "code": "vertex_ai_error", - "message": status, - }, - } + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=status, + ) # Transform successful response using existing transformation vertex_response: Final = vertex_output.get("response", {}) @@ -842,24 +1181,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): response_dict: Final = transformed_response.model_dump() # Return in OpenAI batch format - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": response_dict.get("id", ""), - "body": response_dict, - }, - "error": None, - } + return _openai_batch_output_row(custom_id=custom_id, body=response_dict) except Exception as e: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { - "code": "transformation_error", - "message": f"Failed to transform response: {e}", - }, - } + return _openai_batch_output_row( + custom_id=custom_id, + error_code="transformation_error", + error_message=f"Failed to transform response: {e}", + ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 912e18150b3..7fbbaf84422 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -342,16 +342,25 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: ) -# The six mirrored pricing fields plus the three remaining fields +# The mirrored per-token pricing fields plus the three remaining fields # Router._inherit_builtin_cache_pricing back-fills from the public cost map. An unset field is # what that back-fill targets, so a field left out here is one a PTU deployment still bills. -_PTU_ZEROED_PRICING_FIELDS: Final = SPECIAL_MODEL_INFO_PARAMS + ( +# tiered_pricing is the one mirrored field that is a table of ranges, not a rate, so it is stored +# empty (see _PTU_EMPTIED_PRICING_FIELDS): its tiers outrank the zeros written beside them, so +# dropping it would leave the cost map's tiers billing the traffic the reserved capacity covers. +_PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in SPECIAL_MODEL_INFO_PARAMS if f != "tiered_pricing") + ( "cache_creation_input_token_cost_above_1hr", "cache_creation_input_token_cost_above_200k_tokens", "cache_read_input_token_cost_above_200k_tokens", ) -_PTU_ZEROED_PRICING: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)) -_NO_PRICING_OVERRIDE: Final[Mapping[str, float]] = MappingProxyType({}) +_PTU_EMPTIED_PRICING_FIELDS: Final = frozenset({"tiered_pricing"}) +_PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()]]] = MappingProxyType( + { + **dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0), + **dict.fromkeys(_PTU_EMPTIED_PRICING_FIELDS, ()), + } +) +_NO_PRICING_OVERRIDE: Final[Mapping[str, float | tuple[()]]] = MappingProxyType({}) _EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE # Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges # (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of @@ -364,6 +373,8 @@ def _is_nonzero_price(value: object) -> bool: def _is_zero_price(value: object) -> bool: + if isinstance(value, (list, tuple)): + return not value return isinstance(value, (int, float)) and not isinstance(value, bool) and value == 0 @@ -378,7 +389,12 @@ def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supp return if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None: return - priced: Final = tuple(sorted(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field)))) + priced: Final = tuple( + sorted( + tuple(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field))) + + tuple(field for field in _PTU_EMPTIED_PRICING_FIELDS if supplied.get(field)) + ) + ) if not priced: return raise HTTPException( @@ -395,7 +411,7 @@ def _ptu_zeroed_pricing( model_info: Mapping[str, object], litellm_params: Mapping[str, object], supplied: Mapping[str, object], -) -> Mapping[str, float]: +) -> Mapping[str, float | tuple[()]]: """The pricing a PTU deployment must carry, empty unless one is being stored. Reserved capacity is already billed by the flat cost the rollup writes, so charging the @@ -432,7 +448,7 @@ def _ptu_pricing_delta( model_info: Mapping[str, object], litellm_params: Mapping[str, object], patch: updateDeployment, -) -> tuple[Mapping[str, float], frozenset[str]]: +) -> tuple[Mapping[str, float | tuple[()]], frozenset[str]]: """The pricing a patch must write into both blobs, and the pricing it must drop from them. A patch that takes the deployment off PTU takes the zeroed pricing with it, since the zeros @@ -454,7 +470,7 @@ def _ptu_pricing_delta( return _NO_PRICING_OVERRIDE, frozenset() return _NO_PRICING_OVERRIDE, frozenset( field - for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS) + for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS, _PTU_EMPTIED_PRICING_FIELDS) if _is_zero_price(model_info.get(field)) or _is_zero_price(litellm_params.get(field)) ) @@ -466,11 +482,16 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: override: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=litellm_params) if not override: return model_params + # model_copy validates nothing, so the emptied tier table has to arrive as the list the field + # declares or Pydantic warns on every later dump of it + stored: Final = MappingProxyType( + {key: [] if isinstance(value, tuple) else value for key, value in override.items()} + ) return model_params.model_copy( update=MappingProxyType( { - "litellm_params": model_params.litellm_params.model_copy(update=override), - "model_info": model_params.model_info.model_copy(update=override), + "litellm_params": model_params.litellm_params.model_copy(update=stored), + "model_info": model_params.model_info.model_copy(update=stored), } ) ) diff --git a/litellm/router.py b/litellm/router.py index fb2af41dcf2..4a450cf3c7c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7635,6 +7635,39 @@ class Router: if backend_value is not None: model_info[field] = backend_value + @staticmethod + def _inherit_builtin_tiered_output_rate( + model_info: dict, backend_model: str, custom_llm_provider: str | None + ) -> None: + """Fill a missing entry-level output rate on a deployment entry whose tier + table omits one, from the backend model's built-in cost map entry. + + A deployment's custom pricing is registered as its own standalone + ``litellm.model_cost`` entry holding only the supplied fields, and the + tiered-cost output fallback reads that same entry, so a tier table that + spells out only input-side rates would bill every completion at 0. + + A user-specified ``output_cost_per_token`` always wins. No-op without a + tier table, when every tier declares its own output rate, or when the + backend model has no canonical entry or no flat output rate: + ``get_model_info`` synthesizes a zero for tiered-only backends, and + storing that zero would mark the deployment as explicitly priced free. + """ + tiers: Final = model_info.get("tiered_pricing") + if not isinstance(tiers, list) or not tiers: + return + if model_info.get("output_cost_per_token") is not None: + return + if all(isinstance(tier, dict) and "output_cost_per_token" in tier for tier in tiers): + return + try: + backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model + return + backend_rate: Final = backend_info.get("output_cost_per_token") + if backend_rate: + model_info["output_cost_per_token"] = backend_rate + def _create_deployment( self, deployment_info: dict, @@ -7670,6 +7703,11 @@ class Router: backend_model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) + Router._inherit_builtin_tiered_output_rate( + model_info=_model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP Router._register_deployment_in_model_cost( @@ -8368,6 +8406,11 @@ class Router: backend_model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) + Router._inherit_builtin_tiered_output_rate( + model_info=_model_info_dict, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) # Register custom pricing in litellm.model_cost. # Mirrors _create_deployment() logic to ensure dynamically-added deployments @@ -8598,6 +8641,11 @@ class Router: backend_model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) + Router._inherit_builtin_tiered_output_rate( + model_info=model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) return model_info @staticmethod diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 220826ccbca..9baadc36f6b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3262,6 +3262,7 @@ class MirroredPricingParams(BaseModel): output_cost_per_character: float | None = None cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None + tiered_pricing: list[dict[str, Any]] | None = None class CustomPricingLiteLLMParams(MirroredPricingParams): @@ -3329,7 +3330,6 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_audio_per_second: float | None = None search_context_cost_per_query: dict[str, Any] | None = None citation_cost_per_token: float | None = None - tiered_pricing: list[dict[str, Any]] | None = None cache_read_input_token_cost_above_272k_tokens: float | None = None cache_read_input_token_cost_above_512k_tokens: float | None = None input_cost_per_image_token: float | None = None diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 56400e0666b..4c54822736c 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -731,6 +731,10 @@ "type": "number", "minimum": 0 }, + "cache_creation_input_token_cost": { + "type": "number", + "minimum": 0 + }, "input_cost_per_query": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index d22a139ba79..4d157e74482 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -587,6 +587,246 @@ def test_generic_cost_per_token_honors_non_standard_above_threshold(): litellm.model_cost.pop(model, None) +def test_generic_cost_per_token_tiered_pricing_charges_cache_creation_at_tier_rate(): + """Regression for LIT-4375: a tier's cache_creation_input_token_cost must be billed + on the generic (provider-agnostic) path, not silently dropped.""" + model = "litellm-test-tiered-cache-creation" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 256000], + "input_cost_per_token": 3.25e-07, + "output_cost_per_token": 1.95e-06, + "cache_creation_input_token_cost": 4.063e-07, + "cache_read_input_token_cost": 3.25e-08, + }, + { + "range": [256000, 1000000], + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 3.9e-06, + "cache_creation_input_token_cost": 8.125e-07, + "cache_read_input_token_cost": 6.5e-08, + }, + ], + } + } + ) + + try: + usage = Usage( + prompt_tokens=300000, # 200k new + 60k cache creation + 40k cache read + completion_tokens=1000, + total_tokens=301000, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=40000, cache_creation_tokens=60000 + ), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + expected_prompt = ( + (200000 * 6.5e-07) + (60000 * 8.125e-07) + (40000 * 6.5e-08) + ) + assert round(prompt_cost, 10) == round(expected_prompt, 10) + assert round(completion_cost, 10) == round(1000 * 3.9e-06, 10) + finally: + litellm.model_cost.pop(model, None) + + +def test_generic_cost_per_token_tiered_pricing_is_all_or_nothing(): + """Tiered pricing bills the whole request at the tier picked from its input tokens, + for any provider, and falls back to flat pricing when no tier matches.""" + model = "litellm-test-tiered-all-or-nothing" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "tiered_pricing": [ + { + "range": [0, 32000], + "input_cost_per_token": 4.6e-07, + "output_cost_per_token": 2.3e-06, + }, + { + "range": [32000, 128000], + "input_cost_per_token": 7e-07, + "output_cost_per_token": 3.5e-06, + }, + ], + } + } + ) + + try: + usage = Usage(prompt_tokens=40000, completion_tokens=1000, total_tokens=41000) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + assert round(prompt_cost, 10) == round(40000 * 7e-07, 10) + assert round(completion_cost, 10) == round(1000 * 3.5e-06, 10) + + boundary_usage = Usage(prompt_tokens=32000, completion_tokens=10, total_tokens=32010) + boundary_prompt_cost, _ = generic_cost_per_token( + model=model, + usage=boundary_usage, + custom_llm_provider=custom_llm_provider, + ) + assert round(boundary_prompt_cost, 10) == round(32000 * 4.6e-07, 10) + + empty_prompt_usage = Usage(prompt_tokens=0, completion_tokens=100, total_tokens=100) + empty_prompt_cost, empty_completion_cost = generic_cost_per_token( + model=model, + usage=empty_prompt_usage, + custom_llm_provider=custom_llm_provider, + ) + assert empty_prompt_cost == 0.0 + assert round(empty_completion_cost, 10) == round(100 * 2e-06, 10) + finally: + litellm.model_cost.pop(model, None) + + +def test_generic_cost_per_token_tier_without_an_output_rate_bills_the_model_rate(): + """Regression: a tier table that spells out only input rates served every completion for + free, since a tier's missing output rate has no tier-level fallback to stand in for it.""" + model = "litellm-test-tiered-input-only" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "output_cost_per_token": 2e-06, + "output_cost_per_reasoning_token": 5e-06, + "tiered_pricing": [{"range": [0, 128000], "input_cost_per_token": 1e-03}], + } + } + ) + + try: + usage = Usage( + prompt_tokens=13, + completion_tokens=182, + total_tokens=195, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=100), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + assert round(prompt_cost, 12) == round(13 * 1e-03, 12) + assert round(completion_cost, 12) == round((82 * 2e-06) + (100 * 5e-06), 12) + finally: + litellm.model_cost.pop(model, None) + + +def test_router_deployment_with_input_only_tiers_bills_completions_at_the_backend_rate(): + """Regression: the router registers a deployment's custom pricing as a standalone + model_cost entry holding only the supplied fields, so an input-only tier table left + the output-rate fallback nothing to read and billed every completion at 0.""" + from litellm import Router + + model_id = "litellm-test-router-tiered-input-only" + backend_model = "anthropic/claude-haiku-4-5" + backend_output_rate = litellm.get_model_info(backend_model)["output_cost_per_token"] + Router( + model_list=[ + { + "model_name": "tiered-input-only", + "litellm_params": { + "model": backend_model, + "api_key": "sk-test", + "tiered_pricing": [ + {"range": [0, 3000], "input_cost_per_token": 3.25e-07}, + {"range": [3000, 128000], "input_cost_per_token": 8.125e-07}, + ], + }, + "model_info": {"id": model_id}, + } + ] + ) + + try: + usage = Usage(prompt_tokens=21, completion_tokens=4, total_tokens=25) + prompt_cost, completion_cost = generic_cost_per_token( + model=model_id, + usage=usage, + custom_llm_provider="anthropic", + ) + assert round(prompt_cost, 12) == round(21 * 3.25e-07, 12) + assert round(completion_cost, 12) == round(4 * backend_output_rate, 12) + assert backend_output_rate > 0 + finally: + litellm.model_cost.pop(model_id, None) + + +def test_generic_cost_per_token_tiered_pricing_bills_reasoning_at_tier_rate(): + """Regression: a tier's output_cost_per_reasoning_token must price reasoning tokens + on the generic path and in the logged breakdown, not the tier's plain output rate.""" + model = "litellm-test-tiered-reasoning" + custom_llm_provider = "openrouter" + litellm.register_model( + { + model: { + "litellm_provider": custom_llm_provider, + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 256000], + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 4e-06, + }, + { + "range": [256000, 1000000], + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 3.6e-06, + "output_cost_per_reasoning_token": 1.2e-05, + }, + ], + } + } + ) + + try: + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=400), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + assert round(prompt_cost, 12) == round(1000 * 4e-07, 12) + assert round(completion_cost, 12) == round((100 * 1.2e-06) + (400 * 4e-06), 12) + + breakdown = get_token_type_cost_breakdown( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + ) + assert round(breakdown.reasoning_cost, 12) == round(400 * 4e-06, 12) + finally: + litellm.model_cost.pop(model, None) + + def test_generic_cost_per_token_gpt55(): """gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input.""" model = "gpt-5.5" diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index 6041a8c8377..6f5aaabae06 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -2,13 +2,12 @@ Test suite for Dashscope cost calculation functionality. Tests the cost calculation for Dashscope models including: -- Correctly calculates graduated tiered pricing. +- All-or-nothing tiered pricing, selected by the request's total input tokens. - Falls back to flat-rate pricing for non-tiered models. -- Handles interactions with cached tokens. -- Correctly calculates costs for token counts exceeding the highest defined tier. +- Handles cache read and cache creation tokens. +- Correctly prices requests exceeding the highest defined tier. """ -import json import math import os import sys @@ -22,7 +21,11 @@ import litellm from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, ) -from litellm.types.utils import Usage, PromptTokensDetailsWrapper +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, + Usage, +) class TestDashscopeCostCalculator: @@ -41,7 +44,6 @@ class TestDashscopeCostCalculator: """ usage = Usage(prompt_tokens=1000, completion_tokens=500) - # We call the specific calculator for dashscope prompt_cost, completion_cost = dashscope_cost_per_token( model="qwen-max", usage=usage ) @@ -55,7 +57,7 @@ class TestDashscopeCostCalculator: def test_dashscope_tiered_pricing_within_first_tier(self): """ - Tests the dashscope tiered pricing when token count is entirely within the first tier. + Tests the dashscope tiered pricing when the request's input falls in the first tier. Uses 'dashscope/qwen-flash' as a real-world example. """ # Tier 1 for qwen-flash is [0, 256,000] tokens @@ -73,10 +75,10 @@ class TestDashscopeCostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - def test_dashscope_tiered_pricing_spanning_multiple_tiers(self): + def test_dashscope_tiered_pricing_bills_whole_request_at_selected_tier(self): """ - Tests the dashscope tiered pricing with the corrected graduated calculation logic. - This is the most important test for validating the fix. + Regression: Model Studio tiered pricing is all-or-nothing, not graduated. An input + above the first tier's range must bill every token at the higher tier's rate. """ # Tiering for qwen-flash: Tier 1: [0, 256k], Tier 2: [256k, 1M] usage = Usage(prompt_tokens=300000, completion_tokens=300000) @@ -88,23 +90,54 @@ class TestDashscopeCostCalculator: tier_1 = model_info["tiered_pricing"][0] tier_2 = model_info["tiered_pricing"][1] - # Expected prompt cost: (256,000 tokens * tier_1_price) + (44,000 tokens * tier_2_price) - expected_prompt_cost = (256000 * tier_1["input_cost_per_token"]) + ( - 44000 * tier_2["input_cost_per_token"] - ) - - # Expected completion cost: (256,000 tokens * tier_1_price) + (44,000 tokens * tier_2_price) - expected_completion_cost = (256000 * tier_1["output_cost_per_token"]) + ( - 44000 * tier_2["output_cost_per_token"] - ) + expected_prompt_cost = 300000 * tier_2["input_cost_per_token"] + expected_completion_cost = 300000 * tier_2["output_cost_per_token"] assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + graduated_prompt_cost = (256000 * tier_1["input_cost_per_token"]) + ( + 44000 * tier_2["input_cost_per_token"] + ) + assert prompt_cost > graduated_prompt_cost + + def test_dashscope_tiered_pricing_boundary_stays_in_lower_tier(self): + """ + A request of exactly range_end tokens stays in the lower tier, matching the + official `0 < Token <= 256K` phrasing. + """ + usage = Usage(prompt_tokens=256000, completion_tokens=1000) + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-flash", usage=usage + ) + + tier_1 = litellm.get_model_info("dashscope/qwen-flash")["tiered_pricing"][0] + + assert math.isclose( + prompt_cost, 256000 * tier_1["input_cost_per_token"], rel_tol=1e-10 + ) + assert math.isclose( + completion_cost, 1000 * tier_1["output_cost_per_token"], rel_tol=1e-10 + ) + + def test_dashscope_tiered_pricing_output_uses_input_selected_tier(self): + """ + The tier is chosen by input volume only: a small input with a huge output stays + on the first tier's output rate. + """ + usage = Usage(prompt_tokens=1000, completion_tokens=400000) + _, completion_cost = dashscope_cost_per_token(model="qwen-flash", usage=usage) + + tier_1 = litellm.get_model_info("dashscope/qwen-flash")["tiered_pricing"][0] + + assert math.isclose( + completion_cost, 400000 * tier_1["output_cost_per_token"], rel_tol=1e-10 + ) + def test_dashscope_tiered_pricing_with_caching(self): """ - Tests tiered pricing with cached tokens. This replaces the old, incorrect test. - Uses qwen3-coder-plus, which has cache-specific pricing defined. + Tests tiered pricing with cached tokens: the tier is selected from the total + input (text + cached), and cache reads bill at that tier's cache rate. """ usage = Usage( prompt_tokens=50000, # 10k cached + 40k new @@ -115,28 +148,43 @@ class TestDashscopeCostCalculator: prompt_cost, _ = dashscope_cost_per_token(model="qwen3-coder-plus", usage=usage) - model_info = litellm.get_model_info("dashscope/qwen3-coder-plus") - tier_1 = model_info["tiered_pricing"][0] - tier_2 = model_info["tiered_pricing"][1] + # 50k total input falls in qwen3-coder-plus tier 2 ([32k, 128k]) + tier_2 = litellm.get_model_info("dashscope/qwen3-coder-plus")["tiered_pricing"][1] - # 10k cached tokens are all in the first tier - expected_cache_cost = 10000 * tier_1["cache_read_input_token_cost"] - - # 40k new tokens: 32k in tier 1, and the remaining 8k in tier 2 - expected_text_cost = (32000 * tier_1["input_cost_per_token"]) + ( - 8000 * tier_2["input_cost_per_token"] + expected_prompt_cost = (40000 * tier_2["input_cost_per_token"]) + ( + 10000 * tier_2["cache_read_input_token_cost"] ) - expected_total_prompt_cost = expected_cache_cost + expected_text_cost + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) - assert math.isclose(prompt_cost, expected_total_prompt_cost, rel_tol=1e-10) + def test_dashscope_tiered_pricing_exceeding_highest_tier(self): + """ + Requests above the highest declared range bill entirely at the last tier's rate. + """ + usage = Usage( + prompt_tokens=1200000, completion_tokens=1000 + ) # Max defined range for qwen-flash is 1M - def _register_string_valued_tiered_model(self, model_key: str) -> None: - """Register a model whose tier costs are strings, mimicking YAML config parsing.""" + prompt_cost, _ = dashscope_cost_per_token(model="qwen-flash", usage=usage) + + tier_2 = litellm.get_model_info("dashscope/qwen-flash")["tiered_pricing"][1] + + assert math.isclose( + prompt_cost, 1200000 * tier_2["input_cost_per_token"], rel_tol=1e-10 + ) + + def _register_tiered_model(self, model_key: str, tiered_pricing: list[dict]) -> None: litellm.model_cost[model_key] = { "litellm_provider": "dashscope", "mode": "chat", - "tiered_pricing": [ + "tiered_pricing": tiered_pricing, + } + + def _register_string_valued_tiered_model(self, model_key: str) -> None: + """Register a model whose tier costs are strings, mimicking YAML config parsing.""" + self._register_tiered_model( + model_key, + [ { "range": [0, 1000], "input_cost_per_token": "4e-07", @@ -148,12 +196,12 @@ class TestDashscopeCostCalculator: "output_cost_per_token": "3.2e-06", }, ], - } + ) def test_dashscope_tiered_pricing_string_costs_within_tier(self): """ - Regression: YAML-parsed tier costs can be strings (e.g. "4e-07"). Costs that - fall entirely within a single tier must still be computed as floats. + Regression: YAML-parsed tier costs can be strings (e.g. "4e-07") and must still + be computed as floats. """ self._register_string_valued_tiered_model("dashscope/qwen-str-tier-test") @@ -162,18 +210,13 @@ class TestDashscopeCostCalculator: model="qwen-str-tier-test", usage=usage ) - expected_prompt_cost = 500 * float("4e-07") - expected_completion_cost = 200 * float("1.6e-06") - - assert prompt_cost > 0 - assert completion_cost > 0 - assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) - assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + assert math.isclose(prompt_cost, 500 * float("4e-07"), rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * float("1.6e-06"), rel_tol=1e-10) def test_dashscope_tiered_pricing_string_costs_exceeding_highest_tier(self): """ - Regression: string-valued tier costs must also be coerced in the - remaining-tokens path that charges tokens above the highest tier. + Regression: string-valued tier costs must also be coerced on the last-tier + fallback path used by requests above the highest range. """ self._register_string_valued_tiered_model("dashscope/qwen-str-tier-test") @@ -182,43 +225,265 @@ class TestDashscopeCostCalculator: model="qwen-str-tier-test", usage=usage ) - # prompt: 1000 @ tier1 + 1000 @ tier2 + 500 remaining @ tier2 rate - expected_prompt_cost = ( - (1000 * float("4e-07")) + (1000 * float("8e-07")) + (500 * float("8e-07")) - ) - # completion: 1000 @ tier1 + 1000 @ tier2 + 1000 remaining @ tier2 rate - expected_completion_cost = ( - (1000 * float("1.6e-06")) + (1000 * float("3.2e-06")) + (1000 * float("3.2e-06")) + assert math.isclose(prompt_cost, 2500 * float("8e-07"), rel_tol=1e-10) + assert math.isclose(completion_cost, 3000 * float("3.2e-06"), rel_tol=1e-10) + + def test_dashscope_tiered_cache_creation_tokens_use_tier_rate(self): + """ + Regression (tiered cache creation): cache-creation tokens must bill at the + selected tier's cache_creation_input_token_cost, not the input rate. + """ + self._register_tiered_model( + "dashscope/qwen-cache-write-test", + [ + { + "range": [0, 256000], + "input_cost_per_token": 3.25e-07, + "output_cost_per_token": 1.95e-06, + "cache_creation_input_token_cost": 4.063e-07, + "cache_read_input_token_cost": 3.25e-08, + }, + { + "range": [256000, 1000000], + "input_cost_per_token": 6.5e-07, + "output_cost_per_token": 3.9e-06, + "cache_creation_input_token_cost": 8.125e-07, + "cache_read_input_token_cost": 6.5e-08, + }, + ], ) - assert prompt_cost > 0 - assert completion_cost > 0 - assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) - assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - - def test_dashscope_tiered_pricing_exceeding_highest_tier(self): - """ - Tests tiered pricing when token count exceeds the highest defined tier range. - This replaces the old, incorrect test and validates the new fallback logic. - """ usage = Usage( - prompt_tokens=1200000, completion_tokens=1000 - ) # Max defined range for qwen-flash is 1M + prompt_tokens=300000, # 200k new + 60k cache creation + 40k cache read + completion_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=40000, cache_creation_tokens=60000 + ), + ) - prompt_cost, _ = dashscope_cost_per_token(model="qwen-flash", usage=usage) - - model_info = litellm.get_model_info("dashscope/qwen-flash") - tier_1 = model_info["tiered_pricing"][0] - tier_2 = model_info["tiered_pricing"][1] - - # Expected cost: (tier_1_tokens * tier_1_price) + (tokens_up_to_max_range_in_tier_2 * tier_2_price) + (remaining_tokens * tier_2_price) - tokens_in_tier_2_range = 1000000 - 256000 - remaining_tokens_over_max = 1200000 - 1000000 + prompt_cost, _ = dashscope_cost_per_token( + model="qwen-cache-write-test", usage=usage + ) expected_prompt_cost = ( - (256000 * tier_1["input_cost_per_token"]) - + (tokens_in_tier_2_range * tier_2["input_cost_per_token"]) - + (remaining_tokens_over_max * tier_2["input_cost_per_token"]) + (200000 * 6.5e-07) + (60000 * 8.125e-07) + (40000 * 6.5e-08) ) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + + def test_dashscope_tiered_cache_creation_falls_back_to_tier_input_rate(self): + """ + Tiers without a cache_creation_input_token_cost bill cache-creation tokens at + that tier's input rate. + """ + self._register_tiered_model( + "dashscope/qwen-no-cache-write-test", + [ + { + "range": [0, 256000], + "input_cost_per_token": 3.25e-07, + "output_cost_per_token": 1.95e-06, + } + ], + ) + + usage = Usage( + prompt_tokens=10000, + completion_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_creation_tokens=4000), + ) + + prompt_cost, _ = dashscope_cost_per_token( + model="qwen-no-cache-write-test", usage=usage + ) + + assert math.isclose(prompt_cost, 10000 * 3.25e-07, rel_tol=1e-10) + + def test_dashscope_flat_cache_creation_tokens_use_flat_rate(self): + """Flat-priced models bill cache-creation tokens at their cache-creation rate.""" + litellm.model_cost["dashscope/qwen-flat-cache-write-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "input_cost_per_token": 3.25e-07, + "output_cost_per_token": 1.95e-06, + "cache_creation_input_token_cost": 4.063e-07, + "cache_read_input_token_cost": 3.25e-08, + } + + usage = Usage( + prompt_tokens=10000, + completion_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=2000, cache_creation_tokens=3000 + ), + ) + + prompt_cost, _ = dashscope_cost_per_token( + model="qwen-flat-cache-write-test", usage=usage + ) + + expected_prompt_cost = ( + (5000 * 3.25e-07) + (3000 * 4.063e-07) + (2000 * 3.25e-08) + ) + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + + def test_dashscope_tier_without_an_output_rate_bills_the_model_rate(self): + """ + Regression: a tier declaring only an input rate served every completion for free, + since a missing tier output rate had no tier-level fallback to stand in for it. + """ + litellm.model_cost["dashscope/qwen-input-only-tier-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "tiered_pricing": [{"range": [0, 1000], "input_cost_per_token": 4e-07}], + } + + usage = Usage(prompt_tokens=500, completion_tokens=200) + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-input-only-tier-test", usage=usage + ) + + assert math.isclose(prompt_cost, 500 * 4e-07, rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 1.6e-06, rel_tol=1e-10) + + def test_dashscope_tier_without_an_output_rate_bills_the_model_reasoning_rate(self): + """ + Regression: a tier declaring only an input rate billed reasoning tokens at the model's + plain output rate, ignoring the model's dedicated reasoning rate. + """ + litellm.model_cost["dashscope/qwen-input-only-reasoning-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "output_cost_per_reasoning_token": 4e-06, + "tiered_pricing": [{"range": [0, 1000], "input_cost_per_token": 4e-07}], + } + + usage = Usage( + prompt_tokens=500, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=150), + ) + _, completion_cost = dashscope_cost_per_token( + model="qwen-input-only-reasoning-test", usage=usage + ) + + assert math.isclose( + completion_cost, (50 * 1.6e-06) + (150 * 4e-06), rel_tol=1e-10 + ) + + def test_dashscope_tier_output_rate_wins_over_the_model_reasoning_rate(self): + """ + A tier declaring its own output rate keeps reasoning tokens on that tier rather than + mixing in a model-level reasoning rate. + """ + litellm.model_cost["dashscope/qwen-tier-output-reasoning-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "tiered_pricing": [ + { + "range": [0, 1000], + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + } + ], + } + + usage = Usage( + prompt_tokens=500, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=150), + ) + _, completion_cost = dashscope_cost_per_token( + model="qwen-tier-output-reasoning-test", usage=usage + ) + + assert math.isclose(completion_cost, 200 * 1.6e-06, rel_tol=1e-10) + + def test_dashscope_model_zero_reasoning_rate_bills_reasoning_free(self): + """ + Regression: a model declaring an explicit zero reasoning rate had it treated as + missing, billing reasoning tokens at the plain output rate instead of free. + """ + litellm.model_cost["dashscope/qwen-zero-reasoning-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "output_cost_per_reasoning_token": 0, + } + + usage = Usage( + prompt_tokens=500, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=150), + ) + _, completion_cost = dashscope_cost_per_token( + model="qwen-zero-reasoning-test", usage=usage + ) + + assert math.isclose(completion_cost, 50 * 1.6e-06, rel_tol=1e-10) + + def test_dashscope_tier_zero_reasoning_rate_bills_reasoning_free(self): + """ + Regression: a tier declaring an explicit zero reasoning rate had it treated as + missing, billing reasoning tokens at the tier's output rate instead of free. + """ + litellm.model_cost["dashscope/qwen-tier-zero-reasoning-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 1000], + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "output_cost_per_reasoning_token": 0, + } + ], + } + + usage = Usage( + prompt_tokens=500, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=150), + ) + _, completion_cost = dashscope_cost_per_token( + model="qwen-tier-zero-reasoning-test", usage=usage + ) + + assert math.isclose(completion_cost, 50 * 1.6e-06, rel_tol=1e-10) + + def test_dashscope_tiered_pricing_zero_input_falls_back_to_flat_rates(self): + """ + No tier can be selected without input tokens, so an empty-prompt request must + not be charged at the most expensive tier. + """ + litellm.model_cost["dashscope/qwen-zero-input-test"] = { + "litellm_provider": "dashscope", + "mode": "chat", + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "tiered_pricing": [ + { + "range": [0, 1000], + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + }, + { + "range": [1000, 2000], + "input_cost_per_token": 8e-07, + "output_cost_per_token": 3.2e-06, + }, + ], + } + + usage = Usage(prompt_tokens=0, completion_tokens=500) + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-zero-input-test", usage=usage + ) + + assert prompt_cost == 0.0 + assert math.isclose(completion_cost, 500 * 1.6e-06, rel_tol=1e-10) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 2e3280c0ed1..957fc7dbcf4 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -37,7 +37,7 @@ from litellm.llms.vertex_ai.files.transformation import ( _get_litellm_batch_custom_id_from_labels, _iter_openai_jsonl_entries, _iter_openai_jsonl_lines, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, ) from litellm.types.llms.openai import CreateFileRequest @@ -84,8 +84,9 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st transform, so the streaming path can be checked against it for parity.""" entries = [json.loads(line) for line in content.splitlines() if line.strip()] return "\n".join( - json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params)) + json.dumps(row) for entry in entries + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8c5305ee67b..3c2d56997b7 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent @@ -32,40 +32,26 @@ class TestParseGcsUri: def test_should_parse_standard_gs_uri(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/object.jsonl" - bucket, encoded = config._parse_gcs_uri( - file_id, litellm_params={"gcs_bucket_name": "my-bucket"} - ) + bucket, encoded = config._parse_gcs_uri(file_id, litellm_params={"gcs_bucket_name": "my-bucket"}) assert bucket == "my-bucket" - assert encoded == urllib.parse.quote( - "litellm-vertex-files/path/to/object.jsonl", safe="" - ) + assert encoded == urllib.parse.quote("litellm-vertex-files/path/to/object.jsonl", safe="") def test_should_parse_uri_with_nested_publisher_path(self, config): uri = "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" - bucket, encoded = config._parse_gcs_uri( - uri, litellm_params={"gcs_bucket_name": "litellm-local"} - ) + bucket, encoded = config._parse_gcs_uri(uri, litellm_params={"gcs_bucket_name": "litellm-local"}) assert bucket == "litellm-local" - expected_path = ( - "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" - ) + expected_path = "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" assert encoded == urllib.parse.quote(expected_path, safe="") def test_should_handle_url_encoded_input(self, config): - encoded_uri = urllib.parse.quote( - "gs://my-bucket/litellm-vertex-files/some/path", safe="" - ) - bucket, encoded = config._parse_gcs_uri( - encoded_uri, litellm_params={"gcs_bucket_name": "my-bucket"} - ) + encoded_uri = urllib.parse.quote("gs://my-bucket/litellm-vertex-files/some/path", safe="") + bucket, encoded = config._parse_gcs_uri(encoded_uri, litellm_params={"gcs_bucket_name": "my-bucket"}) assert bucket == "my-bucket" assert encoded == urllib.parse.quote("litellm-vertex-files/some/path", safe="") def test_should_reject_bucket_only(self, config): with pytest.raises(ValueError, match="object name"): - config._parse_gcs_uri( - "gs://my-bucket", litellm_params={"gcs_bucket_name": "my-bucket"} - ) + config._parse_gcs_uri("gs://my-bucket", litellm_params={"gcs_bucket_name": "my-bucket"}) def test_should_reject_no_gs_prefix(self, config): with pytest.raises(ValueError, match="gs://"): @@ -110,9 +96,7 @@ class TestParseGcsUri: "gs://my-bucket/private/object.txt", litellm_params={ "gcs_bucket_name": "my-bucket", - "_litellm_internal_model_credentials": { - "allow_legacy_cloud_file_ids": True - }, + "_litellm_internal_model_credentials": {"allow_legacy_cloud_file_ids": True}, }, ) @@ -176,7 +160,6 @@ class TestCreateFileUrl: class TestTransformRetrieveFile: - def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_retrieve_file_request( @@ -184,13 +167,8 @@ class TestTransformRetrieveFile: optional_params={}, litellm_params={"gcs_bucket_name": "my-bucket"}, ) - expected_encoded = urllib.parse.quote( - "litellm-vertex-files/path/to/file.jsonl", safe="" - ) - assert ( - url - == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{expected_encoded}" - ) + expected_encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") + assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{expected_encoded}" assert params == {} def test_should_return_openai_file_object_from_gcs_response(self, config): @@ -237,7 +215,6 @@ class TestTransformRetrieveFile: class TestTransformFileContent: - def test_should_build_gcs_media_download_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_file_content_request( @@ -246,10 +223,7 @@ class TestTransformFileContent: litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") - assert ( - url - == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}?alt=media" - ) + assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}?alt=media" assert params == {} def test_should_return_binary_response_content(self, config): @@ -269,9 +243,7 @@ class TestTransformFileContent: assert isinstance(result, HttpxBinaryResponseContent) assert result.response.content == b'{"line": 1}\n{"line": 2}\n' - def test_should_not_mutate_caller_logging_obj_for_batch_output_transform( - self, config, monkeypatch - ): + def test_should_not_mutate_caller_logging_obj_for_batch_output_transform(self, config, monkeypatch): original_model = "vertex_ai/original-model" original_start_time = 123.456 original_optional_params = {"temperature": 0.1} @@ -283,9 +255,7 @@ class TestTransformFileContent: "processed_time": "2024-11-01T18:13:16.826+00:00", "request": {"labels": {"litellm_custom_id": "request-1"}}, "response": { - "candidates": [ - {"content": {"parts": [{"text": "ok"}], "role": "model"}} - ], + "candidates": [{"content": {"parts": [{"text": "ok"}], "role": "model"}}], "modelVersion": "gemini-2.0-flash-001@default", }, } @@ -308,9 +278,7 @@ class TestTransformFileContent: captured["logging_obj"] = logging_obj logging_obj.model = "gemini-2.0-flash-001" logging_obj.start_time = 789.0 - return { - "custom_id": vertex_output["request"]["labels"]["litellm_custom_id"] - } + return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]} monkeypatch.setattr( config, @@ -330,9 +298,7 @@ class TestTransformFileContent: assert logging_obj.optional_params == original_optional_params assert result.response is not raw_response - def test_should_skip_batch_output_transformation_when_opt_out_flag_set( - self, config, monkeypatch - ): + def test_should_skip_batch_output_transformation_when_opt_out_flag_set(self, config, monkeypatch): """When `litellm.disable_vertex_batch_output_transformation` is True the Vertex predictions.jsonl content must be returned untouched, so callers that parse raw `candidates`/`modelVersion` keep working.""" @@ -344,9 +310,7 @@ class TestTransformFileContent: "processed_time": "2024-11-01T18:13:16.826+00:00", "request": {"labels": {"litellm_custom_id": "request-1"}}, "response": { - "candidates": [ - {"content": {"parts": [{"text": "ok"}], "role": "model"}} - ], + "candidates": [{"content": {"parts": [{"text": "ok"}], "role": "model"}}], "modelVersion": "gemini-2.0-flash-001@default", }, } @@ -358,9 +322,7 @@ class TestTransformFileContent: request=httpx.Request("GET", "https://example.com"), ) - monkeypatch.setattr( - litellm, "disable_vertex_batch_output_transformation", True, raising=False - ) + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) result = config.transform_file_content_response( raw_response=raw_response, @@ -381,9 +343,7 @@ class TestTransformDeleteFile: litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") - assert ( - url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}" - ) + assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}" assert params == {} def test_should_return_file_deleted_with_reconstructed_id(self, config): @@ -393,9 +353,7 @@ class TestTransformDeleteFile: "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc", safe="", ) - mock_request.url = ( - f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_name}" - ) + mock_request.url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_name}" raw_response.request = mock_request result = config.transform_delete_file_response( @@ -407,10 +365,7 @@ class TestTransformDeleteFile: assert isinstance(result, FileDeleted) assert result.deleted is True assert result.object == "file" - assert ( - result.id - == "gs://my-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" - ) + assert result.id == "gs://my-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" def test_should_fallback_to_deleted_id_when_no_request(self, config): raw_response = MagicMock(spec=httpx.Response) @@ -435,9 +390,7 @@ class TestTransformDeleteFile: raw_response = MagicMock(spec=httpx.Response) mock_request = MagicMock() encoded_object = urllib.parse.quote("path/to/file.jsonl", safe="") - mock_request.url = ( - f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_object}" - ) + mock_request.url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_object}" raw_response.request = mock_request result = config.transform_delete_file_response( @@ -466,8 +419,7 @@ class TestTransformDeleteFile: ) assert result.id == ( - "gs://prod-bucket/litellm-vertex-files/publishers/google/" - "models/gemini-2.0-flash-001/abc-123" + "gs://prod-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" ) @@ -504,9 +456,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) result = json.loads(transformed_content.decode("utf-8")) # Verify OpenAI format @@ -548,9 +498,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) result = json.loads(transformed_content.decode("utf-8")) # Per OpenAI Batch output spec, error entries set response to null @@ -584,9 +532,7 @@ class TestVertexBatchOutputTransformation: } class _RaisingGeminiConfig(VertexGeminiConfig): - def _transform_google_generate_content_to_openai_model_response( - self, *args, **kwargs - ): + def _transform_google_generate_content_to_openai_model_response(self, *args, **kwargs): raise ValueError("simulated transform failure") mock_response = httpx.Response( @@ -637,9 +583,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) result = json.loads(transformed_content.decode("utf-8")) assert result["custom_id"] == "myrequest-1" @@ -651,9 +595,7 @@ class TestVertexBatchOutputTransformation: "status": "", "processed_time": "2024-11-01T18:13:16.826+00:00", "request": { - "contents": [ - {"role": "user", "parts": [{"text": "First request"}]} - ], + "contents": [{"role": "user", "parts": [{"text": "First request"}]}], "labels": {"litellm_custom_id": "request-1"}, }, "response": { @@ -678,9 +620,7 @@ class TestVertexBatchOutputTransformation: "status": "", "processed_time": "2024-11-01T18:13:17.826+00:00", "request": { - "contents": [ - {"role": "user", "parts": [{"text": "Second request"}]} - ], + "contents": [{"role": "user", "parts": [{"text": "Second request"}]}], "labels": {"litellm_custom_id": "request-2"}, }, "response": { @@ -703,12 +643,8 @@ class TestVertexBatchOutputTransformation: }, ] - content = "\n".join(json.dumps(output) for output in vertex_outputs).encode( - "utf-8" - ) - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + content = "\n".join(json.dumps(output) for output in vertex_outputs).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) lines = transformed_content.decode("utf-8").strip().split("\n") assert len(lines) == 2 @@ -718,14 +654,12 @@ class TestVertexBatchOutputTransformation: assert "id" in result assert "response" in result assert result["response"]["status_code"] == 200 - assert result["custom_id"] == f"request-{i+1}" + assert result["custom_id"] == f"request-{i + 1}" body = result["response"]["body"] assert "choices" in body assert len(body["choices"]) > 0 - def test_transform_vertex_batch_output_with_first_line_prompt_feedback( - self, config, monkeypatch - ): + def test_transform_vertex_batch_output_with_first_line_prompt_feedback(self, config, monkeypatch): """Test that promptFeedback-only first lines are detected as Vertex batch output.""" vertex_outputs = [ { @@ -751,9 +685,7 @@ class TestVertexBatchOutputTransformation: logging_obj, mock_httpx_response, ): - return { - "custom_id": vertex_output["request"]["labels"]["litellm_custom_id"] - } + return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]} monkeypatch.setattr( config, @@ -761,15 +693,9 @@ class TestVertexBatchOutputTransformation: mock_transform_single, ) - content = "\n".join(json.dumps(output) for output in vertex_outputs).encode( - "utf-8" - ) - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) - results = [ - json.loads(line) for line in transformed_content.decode("utf-8").split("\n") - ] + content = "\n".join(json.dumps(output) for output in vertex_outputs).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + results = [json.loads(line) for line in transformed_content.decode("utf-8").split("\n")] assert [result["custom_id"] for result in results] == [ "blocked-request", @@ -786,9 +712,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(non_batch_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) assert transformed_content == content @@ -818,9 +742,7 @@ class TestVertexBatchOutputTransformation: id(mock_httpx_response), ) ) - return { - "custom_id": vertex_output["request"]["labels"]["litellm_custom_id"] - } + return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]} monkeypatch.setattr( config, @@ -828,12 +750,8 @@ class TestVertexBatchOutputTransformation: mock_transform_single, ) - content = "\n".join(json.dumps(output) for output in vertex_outputs).encode( - "utf-8" - ) - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + content = "\n".join(json.dumps(output) for output in vertex_outputs).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) assert len(transformed_content.decode("utf-8").strip().split("\n")) == 2 assert len(set(helper_ids)) == 1 @@ -841,17 +759,13 @@ class TestVertexBatchOutputTransformation: def test_non_batch_output_passthrough(self, config): """Test that non-batch output is returned as-is""" regular_content = b"This is just a regular file content" - transformed_content = config._try_transform_vertex_batch_output_to_openai( - regular_content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(regular_content) assert transformed_content == regular_content def test_invalid_json_passthrough(self, config): """Test that invalid JSON is returned as-is""" invalid_content = b'{"invalid": json content}' - transformed_content = config._try_transform_vertex_batch_output_to_openai( - invalid_content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(invalid_content) assert transformed_content == invalid_content def test_binary_content_passthrough(self, config): @@ -903,9 +817,7 @@ class TestVertexBatchOutputTransformation: }, } - content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode( - "utf-8" - ) + content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode("utf-8") def list_pipeline() -> bytes: gemini_config = VertexGeminiConfig() @@ -944,9 +856,7 @@ class TestVertexBatchOutputTransformation: finally: tracemalloc.stop() - streaming_peak = peak_of( - lambda: config._try_transform_vertex_batch_output_to_openai(content) - ) + streaming_peak = peak_of(lambda: config._try_transform_vertex_batch_output_to_openai(content)) list_peak = peak_of(list_pipeline) assert streaming_peak < list_peak * 0.75, ( @@ -999,9 +909,9 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: logging_obj=logging_obj, ) - assert ( - logging_obj.model == sentinel_model - ), "logging_obj.model was mutated by _try_transform_vertex_batch_output_to_openai" + assert logging_obj.model == sentinel_model, ( + "logging_obj.model was mutated by _try_transform_vertex_batch_output_to_openai" + ) def test_should_not_overwrite_start_time_on_caller_logging_obj(self, config): sentinel_start = 1234567890.0 @@ -1014,9 +924,9 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: logging_obj=logging_obj, ) - assert ( - logging_obj.start_time == sentinel_start - ), "logging_obj.start_time was mutated by _try_transform_vertex_batch_output_to_openai" + assert logging_obj.start_time == sentinel_start, ( + "logging_obj.start_time was mutated by _try_transform_vertex_batch_output_to_openai" + ) def test_should_not_overwrite_optional_params_on_caller_logging_obj(self, config): sentinel_params = {"temperature": 0.5, "top_p": 0.9} @@ -1028,9 +938,9 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: logging_obj=logging_obj, ) - assert ( - logging_obj.optional_params is sentinel_params - ), "logging_obj.optional_params was replaced by _try_transform_vertex_batch_output_to_openai" + assert logging_obj.optional_params is sentinel_params, ( + "logging_obj.optional_params was replaced by _try_transform_vertex_batch_output_to_openai" + ) assert logging_obj.optional_params == { "temperature": 0.5, "top_p": 0.9, @@ -1054,14 +964,13 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: def _wrap_entries(openai_jsonl_content): - """Vertex-wrapped requests for a list of OpenAI batch entries, built via the - live single-entry transform that the streaming upload path uses.""" + """Vertex rows for a list of OpenAI batch entries, built via the live + single-entry transform that the streaming upload path uses.""" cfg = VertexAIFilesConfig() return [ - _openai_batch_jsonl_entry_to_vertex_wrapped_request( - entry, cfg._map_openai_to_vertex_params - ) + row for entry in openai_jsonl_content + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ] @@ -1122,9 +1031,7 @@ class TestVertexBatchCustomIdLabels: assert "litellm_custom_id_raw_1" in labels_a assert "litellm_custom_id_raw_1" in labels_b assert labels_a["litellm_custom_id_raw"] == labels_b["litellm_custom_id_raw"] - assert ( - labels_a["litellm_custom_id_raw_1"] != labels_b["litellm_custom_id_raw_1"] - ) + assert labels_a["litellm_custom_id_raw_1"] != labels_b["litellm_custom_id_raw_1"] assert _get_litellm_batch_custom_id_from_labels(labels_a) == custom_id_a assert _get_litellm_batch_custom_id_from_labels(labels_b) == custom_id_b @@ -1133,12 +1040,12 @@ class TestVertexBatchCustomIdLabels: openai_jsonl_content = [ { - "custom_id": f"request-{i+1}", + "custom_id": f"request-{i + 1}", "method": "POST", "url": "/v1/chat/completions", "body": { "model": "gemini-1.5-flash-001", - "messages": [{"role": "user", "content": f"Question {i+1}"}], + "messages": [{"role": "user", "content": f"Question {i + 1}"}], }, } for i in range(3) @@ -1149,11 +1056,8 @@ class TestVertexBatchCustomIdLabels: assert len(vertex_jsonl_content) == 3 for i, vertex_request in enumerate(vertex_jsonl_content): - expected_custom_id = f"request-{i+1}" - assert ( - vertex_request["request"]["labels"]["litellm_custom_id"] - == expected_custom_id - ) + expected_custom_id = f"request-{i + 1}" + assert vertex_request["request"]["labels"]["litellm_custom_id"] == expected_custom_id raw_label = vertex_request["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != expected_custom_id assert _sanitize_gcp_label_value(raw_label) == raw_label @@ -1200,9 +1104,7 @@ class TestVertexBatchCustomIdLabels: vertex_input = _wrap_entries(openai_input) # Verify both labels are GCP-safe and encoded raw preserves round-trip. - assert ( - vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" - ) + assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != "MyRequest-1" assert _sanitize_gcp_label_value(raw_label) == raw_label @@ -1230,9 +1132,7 @@ class TestVertexBatchCustomIdLabels: # Step 3: Transform Vertex AI output back to OpenAI format content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) openai_output = json.loads(transformed_content.decode("utf-8")) # Step 4: Verify custom_id was preserved (original casing, not sanitized label) @@ -1268,9 +1168,7 @@ class TestVertexBatchCustomIdLabels: vertex_input = _wrap_entries(openai_input) # Verify both labels are safe for GCP labels. - assert ( - vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" - ) + assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != "MyRequest-1" assert _sanitize_gcp_label_value(raw_label) == raw_label @@ -1279,26 +1177,15 @@ class TestVertexBatchCustomIdLabels: class TestConfiguredBucketNameResolution: def test_should_resolve_new_gcs_bucket_name_key(self, config, monkeypatch): monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - assert ( - config._get_configured_bucket_name({"gcs_bucket_name": "my-new-bucket"}) - == "my-new-bucket" - ) + assert config._get_configured_bucket_name({"gcs_bucket_name": "my-new-bucket"}) == "my-new-bucket" def test_should_resolve_legacy_bucket_name_key(self, config, monkeypatch): monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - assert ( - config._get_configured_bucket_name({"bucket_name": "my-legacy-bucket"}) - == "my-legacy-bucket" - ) + assert config._get_configured_bucket_name({"bucket_name": "my-legacy-bucket"}) == "my-legacy-bucket" def test_should_prefer_new_key_over_legacy(self, config, monkeypatch): monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - assert ( - config._get_configured_bucket_name( - {"gcs_bucket_name": "new", "bucket_name": "legacy"} - ) - == "new" - ) + assert config._get_configured_bucket_name({"gcs_bucket_name": "new", "bucket_name": "legacy"}) == "new" def test_should_fall_back_to_env(self, config, monkeypatch): monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket") @@ -1318,3 +1205,558 @@ class TestConfiguredBucketNameResolution: assert "bucket_name" in OPTIONAL_KWARGS_KEYS params = get_litellm_params(bucket_name="my-legacy-bucket") assert params.get("bucket_name") == "my-legacy-bucket" + + +def _embeddings_entry(**overrides): + entry = { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "gemini-embedding-2", "input": "hello world"}, + } + entry.update(overrides) + return entry + + +class TestVertexEmbeddingsBatchInputTranslation: + """ + /v1/embeddings batch lines must be translated to Vertex's Gemini Embedding batch + shape, not the generateContent shape. + + Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + + def test_should_emit_embed_content_request_shape(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert row["request"] == {"content": {"parts": [{"text": "hello world"}]}} + assert "contents" not in row["request"] + assert "labels" not in row["request"] + + def test_should_round_trip_custom_id_through_top_level_key(self): + (row,) = _wrap_entries([_embeddings_entry(custom_id="MyRequest-1")]) + + assert row["key"] == "MyRequest-1" + + def test_should_omit_key_when_no_custom_id(self): + entry = _embeddings_entry() + del entry["custom_id"] + + (row,) = _wrap_entries([entry]) + + assert "key" not in row + + def test_should_map_openai_params_into_the_embed_content_request(self): + """ + The docs put these in an `embed_content_config` sibling of `request`, but Vertex + rejects that key and fails the whole job, so they belong inside the request. + """ + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": "hello world", + "dimensions": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + ) + ] + ) + + assert row == { + "key": "request-1", + "request": { + "content": {"parts": [{"text": "hello world"}]}, + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + }, + } + + def test_should_omit_config_fields_when_no_params_given(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert set(row["request"]) == {"content"} + + def test_should_translate_multimodal_gcs_input(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + ) + ] + ) + + assert row["request"]["content"]["parts"] == [ + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + } + ] + + @pytest.mark.parametrize("url", ["/v1/embeddings", "v1/embeddings", "/v1/embeddings/"]) + def test_should_detect_embeddings_route_variants(self, url): + (row,) = _wrap_entries([_embeddings_entry(url=url)]) + + assert "content" in row["request"] + + def test_should_raise_when_input_missing(self): + with pytest.raises(ValueError, match="`input` is required"): + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2"})]) + + def test_should_raise_when_input_empty(self): + with pytest.raises(ValueError, match="must not be empty"): + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2", "input": []})]) + + def test_should_fan_an_input_array_out_into_one_row_per_element(self): + """ + An `EmbedContentRequest` returns exactly one vector, so an OpenAI entry asking + for several embeddings needs several Vertex rows. + """ + rows = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": ["first", "second"], + "dimensions": 768, + } + ) + ] + ) + + assert rows == [ + { + "key": "request-1#0/2", + "request": { + "content": {"parts": [{"text": "first"}]}, + "output_dimensionality": 768, + }, + }, + { + "key": "request-1#1/2", + "request": { + "content": {"parts": [{"text": "second"}]}, + "output_dimensionality": 768, + }, + }, + ] + + def test_should_keep_the_bare_custom_id_for_single_element_arrays(self): + (row,) = _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2", "input": ["only one"]})]) + + assert row["key"] == "request-1" + + def test_should_encode_a_custom_id_that_looks_like_a_fan_out_tag(self): + """A customer custom_id ending in `#/` must not read back as fan-out metadata.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "hello world"}, + ) + ] + ) + + assert row["key"] == "request-1%230%2F2" + + def test_should_combine_a_nested_input_into_one_multipart_row(self): + """Nested arrays are the combined-embedding shape, as on the online path.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": [ + [ + "a caption", + "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + ] + ], + } + ) + ] + ) + + assert row["key"] == "request-1" + assert row["request"]["content"]["parts"] == [ + {"text": "a caption"}, + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + }, + ] + + def test_should_keep_chat_completions_lines_on_generate_content_path(self): + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + assert row["request"]["labels"]["litellm_custom_id"] == "request-1" + assert "key" not in row + + def test_should_keep_lines_without_a_url_on_generate_content_path(self): + """`url` is optional on a batch line, and chat is the shape LiteLLM has always assumed.""" + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + + def test_should_translate_each_line_by_its_own_url(self): + chat_row, embeddings_row = _wrap_entries( + [ + { + "custom_id": "chat-1", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + }, + _embeddings_entry(custom_id="embed-1"), + ] + ) + + assert "contents" in chat_row["request"] + assert "content" in embeddings_row["request"] + + +class TestVertexEmbeddingsBatchOutputTranslation: + """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" + + def _vertex_embeddings_output_row(self, **overrides): + row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": { + "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, + }, + } + row.update(overrides) + return row + + def _transform(self, config, rows, url="https://example.com"): + content = "\n".join(json.dumps(row) for row in rows).encode("utf-8") + result = config.transform_file_content_response( + raw_response=httpx.Response( + status_code=200, + content=content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", url), + ), + logging_obj=MagicMock(), + litellm_params={}, + ) + return [json.loads(line) for line in result.response.content.decode("utf-8").split("\n")] + + def test_should_transform_embeddings_output_to_openai_batch_row(self, config): + (result,) = self._transform(config, [self._vertex_embeddings_output_row()]) + + assert result["custom_id"] == "request-1" + assert result["error"] is None + assert result["response"]["status_code"] == 200 + body = result["response"]["body"] + assert body["object"] == "list" + assert body["data"] == [{"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"}] + assert body["usage"]["prompt_tokens"] == 2 + assert body["usage"]["total_tokens"] == 2 + + def test_should_fall_back_to_documented_token_count_field(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + response={ + "embedding": {"values": [-0.015, 0.024]}, + "tokenCount": "2", + } + ) + ], + ) + + assert result["response"]["body"]["usage"]["prompt_tokens"] == 2 + + def test_should_resolve_model_from_managed_gcs_object_path(self, config): + object_path = urllib.parse.quote( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/" + "prediction-model-2026-07-29T05:55:52Z/predictions.jsonl", + safe="", + ) + url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{object_path}?alt=media" + + (result,) = self._transform(config, [self._vertex_embeddings_output_row()], url=url) + + assert result["response"]["body"]["model"] == "gemini-embedding-2" + + def test_should_surface_failed_embeddings_row_as_error(self, config): + (result,) = self._transform( + config, + [self._vertex_embeddings_output_row(status="Failed to parse JSON into proto", response={})], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert "Failed to parse JSON into proto" in result["error"]["message"] + + def test_should_transform_every_row_of_a_multi_row_file(self, config): + results = self._transform( + config, + [self._vertex_embeddings_output_row(key=f"request-{index}") for index in range(3)], + ) + + assert [result["custom_id"] for result in results] == [ + "request-0", + "request-1", + "request-2", + ] + + def test_should_reassemble_a_fanned_out_input_array_into_one_row(self, config): + """Vertex returns the rows of one entry in arbitrary order.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + key="request-1#1/2", + response={ + "embedding": {"values": [0.3, 0.4]}, + "usageMetadata": {"promptTokenCount": 5}, + }, + ), + self._vertex_embeddings_output_row( + key="request-1#0/2", + response={ + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": {"promptTokenCount": 3}, + }, + ), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"]["body"]["data"] == [ + {"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}, + {"embedding": [0.3, 0.4], "index": 1, "object": "embedding"}, + ] + assert result["response"]["body"]["usage"]["prompt_tokens"] == 8 + + def test_should_keep_fanned_out_entries_apart_and_in_file_order(self, config): + results = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-2#0/2"), + self._vertex_embeddings_output_row(key="request-1"), + self._vertex_embeddings_output_row(key="request-2#1/2"), + ], + ) + + assert [result["custom_id"] for result in results] == ["request-2", "request-1"] + assert len(results[0]["response"]["body"]["data"]) == 2 + assert len(results[1]["response"]["body"]["data"]) == 1 + + def test_should_not_merge_an_entry_whose_custom_id_looks_like_a_fan_out_tag(self, config): + """`request-1#0/2` is a legal custom_id, and a distinct entry from `request-1`.""" + lookalike_row, plain_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "lookalike"}, + ), + _embeddings_entry( + custom_id="request-1", + body={"model": "gemini-embedding-2", "input": "plain"}, + ), + ] + ) + + results = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in ((lookalike_row, [0.1]), (plain_row, [0.2])) + ], + ) + + assert [result["custom_id"] for result in results] == [ + "request-1#0/2", + "request-1", + ] + assert [result["response"]["body"]["data"][0]["embedding"] for result in results] == [[0.1], [0.2]] + + def test_should_round_trip_a_fan_out_of_a_custom_id_holding_the_separator(self, config): + rows = _wrap_entries( + [ + _embeddings_entry( + custom_id="request#1/1", + body={ + "model": "gemini-embedding-2", + "input": ["first", "second"], + }, + ) + ] + ) + + assert [row["key"] for row in rows] == [ + "request%231%2F1#0/2", + "request%231%2F1#1/2", + ] + + (result,) = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in zip(reversed(rows), ([0.3], [0.1])) + ], + ) + + assert result["custom_id"] == "request#1/1" + assert [embedding["embedding"] for embedding in result["response"]["body"]["data"]] == [[0.1], [0.3]] + + def test_should_fail_the_whole_entry_when_one_of_its_rows_failed(self, config): + """An OpenAI batch row is either a response or an error, never both.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-1#0/2"), + self._vertex_embeddings_output_row(key="request-1#1/2", status="Quota exceeded", response={}), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["message"] == "Quota exceeded" + + def test_should_fail_the_whole_entry_when_a_fanned_out_row_is_missing(self, config): + """A partial `data` array would shift embeddings onto the wrong input positions.""" + (result,) = self._transform( + config, + [self._vertex_embeddings_output_row(key="request-1#1/2")], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert result["error"]["message"] == ("Vertex returned embeddings for input positions [1] of the 2 requested") + + def test_should_fail_the_whole_entry_when_a_fanned_out_row_is_duplicated(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-1#0/2"), + self._vertex_embeddings_output_row(key="request-1#0/2"), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert result["error"]["message"] == ( + "Vertex returned embeddings for input positions [0, 0] of the 2 requested" + ) + + def test_should_end_to_end_round_trip_a_fanned_out_embeddings_batch(self, config): + first_row, second_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": ["hello world", "goodbye world"], + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **row, + "status": "", + "response": {"embedding": {"values": values}}, + } + for row, values in ((second_row, [0.3]), (first_row, [0.1])) + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert [embedding["embedding"] for embedding in result["response"]["body"]["data"]] == [[0.1], [0.3]] + + def test_should_end_to_end_round_trip_openai_embeddings_batch(self, config): + (vertex_row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": "hello world", + "dimensions": 2, + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **vertex_row, + "status": "", + "processed_time": "2026-07-29T05:55:52.379528Z", + "response": { + "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, + }, + } + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert result["response"]["body"]["data"][0]["embedding"] == [-0.015, 0.024] + + def test_should_leave_legacy_predict_embeddings_output_untouched(self, config): + legacy_row = { + "instance": {"content": "hello world"}, + "predictions": [ + { + "embeddings": { + "statistics": {"token_count": 2, "truncated": False}, + "values": [0.2], + } + } + ], + "status": "", + } + content = json.dumps(legacy_row).encode("utf-8") + + assert config._try_transform_vertex_batch_output_to_openai(content) == content diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index 6e670e48b6a..d3aec8010f8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -15,6 +15,7 @@ from litellm.proxy._types import ( ReconcileOutcome, UserAPIKeyAuth, ) +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.proxy.auth.auth_checks import _is_model_cost_zero from litellm.proxy.management_endpoints.model_management_endpoints import ( _PTU_ZEROED_PRICING_FIELDS, @@ -37,6 +38,7 @@ from litellm.types.router import ( updateDeployment, updateLiteLLMParams, ) +from litellm.types.utils import Usage def test_model_info_accepts_valid_ptu_fields(): @@ -762,7 +764,10 @@ class TestPtuDeploymentsAreNotBilledPerToken: assert self._zeroed(model_info={"ptu_count": 15}) == {} def test_every_field_the_cost_map_could_fill_is_zeroed(self): - assert self._zeroed(model_info=self.PTU) == dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0) + assert self._zeroed(model_info=self.PTU) == { + **dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0), + "tiered_pricing": (), + } def test_nothing_is_zeroed_while_the_feature_is_disabled(self, monkeypatch): monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) @@ -777,6 +782,54 @@ class TestPtuDeploymentsAreNotBilledPerToken: assert exc.value.status_code == 400 assert field in str(exc.value.detail) + def test_a_tiered_price_the_caller_supplies_is_refused(self): + """Tier rates bill the traffic per token just as surely as a flat rate does.""" + with pytest.raises(HTTPException) as exc: + self._zeroed(model_info=self.PTU, supplied={"tiered_pricing": [{"range": [0, 100], "input_cost_per_token": 1e-06}]}) + assert exc.value.status_code == 400 + assert "tiered_pricing" in str(exc.value.detail) + + def test_tiered_pricing_already_on_the_row_is_emptied_not_zeroed(self): + """tiered_pricing is a table of ranges, so the zero the other fields store would not even + validate. Dropping it instead would fall back to the cost map's tiers, whose rates outrank + the zeros written beside them, so it is stored empty.""" + tiers = [{"range": [0, 128000], "input_cost_per_token": 3e-06}] + priced = _ptu_priced_deployment( + Deployment( + model_name="tiered", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo( + id="dep-tiered", + team_id="t", + tiered_pricing=tiers, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + ) + assert priced.litellm_params.tiered_pricing == [] + assert priced.model_info.tiered_pricing == [] + + written = update_db_model( + db_model=Deployment( + model_name="tiered", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", tiered_pricing=tiers), + model_info=ModelInfo(id="dep-tiered", team_id="t"), + ), + updated_patch=updateDeployment( + model_info=ModelInfo( + id="dep-tiered", + team_id="t", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ) + ), + ) + for blob in ("model_info", "litellm_params"): + stored = json.loads(written[blob]) + assert stored["tiered_pricing"] == [], blob + assert stored["input_cost_per_token"] == 0, blob + def test_a_price_the_caller_supplies_as_zero_is_accepted(self): assert self._zeroed(model_info={**self.PTU, "input_cost_per_token": 0}, supplied={"input_cost_per_token": 0})[ "input_cost_per_token" @@ -910,6 +963,32 @@ class TestPtuDeploymentsAreNotBilledPerToken: charged = {k: v for k, v in registered.items() if "cost" in k and k != "cost_per_ptu_per_hour" and v} assert charged == {} + def test_the_cost_map_tiers_contribute_no_price_to_a_priced_ptu_deployment(self): + """A tier table outranks the zeroed flat rates wherever cost is read, so leaving the + deployment's own table unset bills the reserved capacity's traffic at the map's tiers.""" + priced = _ptu_priced_deployment( + Deployment( + model_name="ptu-deployment", + litellm_params=LiteLLM_Params(model="dashscope/qwen-flash", api_key="fake-key"), + model_info=ModelInfo( + id="dep-ptu", + team_id="team-1", + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + **self.PTU, + ), + ) + ) + router = Router(model_list=[priced.to_json(exclude_none=True)]) + registered = router.get_deployment_model_info(model_id="dep-ptu", model_name="dashscope/qwen-flash") + assert registered is not None + assert registered["tiered_pricing"] == [] + assert generic_cost_per_token( + model="dashscope/qwen-flash", + usage=Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100), + custom_llm_provider="dashscope", + model_info=registered, + ) == (0.0, 0.0) + def test_the_zeroed_pricing_does_not_waive_budget_enforcement(self): """A zero price otherwise tells auth the model is free and skips every budget check.""" priced = _ptu_priced_deployment( @@ -1098,8 +1177,10 @@ class TestPtuDeploymentsAreNotBilledPerToken: ) written = add_team_model_to_db.call_args.kwargs["model_params"] - assert all(getattr(written.model_info, field, None) == 0 for field in SPECIAL_MODEL_INFO_PARAMS) + assert all(getattr(written.model_info, field, None) == 0 for field in SPECIAL_MODEL_INFO_PARAMS if field != "tiered_pricing") + assert written.model_info.tiered_pricing == [] assert all(written.litellm_params.get(field) == 0 for field in _PTU_ZEROED_PRICING_FIELDS) + assert written.litellm_params.tiered_pricing == [] @pytest.mark.asyncio async def test_model_new_refuses_a_priced_ptu_deployment(self): diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 26ba485d796..a51f4e733b6 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -42,6 +42,26 @@ def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch): assert prompt_usd + completion_usd > 0 +def test_cost_per_token_tiered_only_model_bills_at_tier_rate(monkeypatch): + """ + Regression: models that publish only tiered_pricing (no top-level per-token rates), + e.g. volcengine doubao-seed-2.0, must reach the generic tiered path instead of + recording zero spend. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + prompt_usd, completion_usd = cost_per_token( + model="volcengine/doubao-seed-2-0-pro-260215", + prompt_tokens=40000, + completion_tokens=500, + custom_llm_provider="volcengine", + ) + + assert prompt_usd == pytest.approx(40000 * 7e-07) + assert completion_usd == pytest.approx(500 * 3.5e-06) + + def test_cost_per_token_non_string_model_does_not_hang(): """ The provider-prefix dedup loop must not spin forever when `model` is a diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index ccb0541d318..cb7023e6c12 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -100,6 +100,24 @@ def test_schema_accepts_minimal_and_unknown_optional_fields(committed_schema: di assert validator.is_valid({"some-model": {"litellm_provider": "openai", "brand_new_field": {"nested": True}}}) +def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_schema: dict): + validator = build_validator(committed_schema) + entry = { + "litellm_provider": "dashscope", + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 256000], + "input_cost_per_token": 3.25e-07, + "output_cost_per_token": 1.95e-06, + "cache_creation_input_token_cost": 4.063e-07, + "cache_read_input_token_cost": 3.25e-08, + } + ], + } + assert validator.is_valid({"some-model": entry}) + + DATED_VARIANT = re.compile(r"^(.*?)-(\d{4}-\d{2}-\d{2})$") SERVICE_TIER_SUFFIXES = ("_flex", "_priority") diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 3fb4511e52c..dfe46d54ab8 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -1529,3 +1529,58 @@ def test_strategy_router_alias_pricing_never_enters_model_cost(monkeypatch): finally: litellm.model_cost = saved_model_cost _invalidate_model_cost_lowercase_map() + + +def test_inherit_builtin_tiered_output_rate_fills_the_backend_flat_rate(): + """ + A deployment entry whose custom tiers publish only input rates would bill + completions at 0, so the backend model's flat output rate is copied in at + registration. + """ + model_info = {"tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}]} + + Router._inherit_builtin_tiered_output_rate( + model_info=model_info, + backend_model="claude-haiku-4-5", + custom_llm_provider="anthropic", + ) + + backend_rate = litellm.get_model_info(model="claude-haiku-4-5", custom_llm_provider="anthropic")[ + "output_cost_per_token" + ] + assert backend_rate > 0 + assert model_info["output_cost_per_token"] == backend_rate + + +def test_inherit_builtin_tiered_output_rate_never_stores_a_synthesized_zero(): + """ + Regression: get_model_info reports output_cost_per_token 0 for a backend that + only publishes tiered rates (e.g. dashscope/qwen-flash), and storing that zero + would mark the deployment as explicitly priced free. + """ + backend_info = litellm.get_model_info(model="qwen-flash", custom_llm_provider="dashscope") + assert backend_info["output_cost_per_token"] == 0 + + model_info = {"tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}]} + Router._inherit_builtin_tiered_output_rate( + model_info=model_info, + backend_model="qwen-flash", + custom_llm_provider="dashscope", + ) + + assert "output_cost_per_token" not in model_info + + +def test_inherit_builtin_tiered_output_rate_leaves_a_user_rate_alone(): + model_info = { + "tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}], + "output_cost_per_token": 9e-07, + } + + Router._inherit_builtin_tiered_output_rate( + model_info=model_info, + backend_model="claude-haiku-4-5", + custom_llm_provider="anthropic", + ) + + assert model_info["output_cost_per_token"] == 9e-07 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e9e6167fb9..01e7e5c7ffd 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1021,6 +1021,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "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", diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index 1b66863a82f..5ce5eca4954 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -46,12 +46,30 @@ def test_custom_pricing_params_keeps_every_field_it_had(): @pytest.mark.parametrize("field", SPECIAL_MODEL_INFO_PARAMS) def test_deployment_mirrors_pricing_from_litellm_params_onto_model_info(field): + value = [{"range": [0, 128000], "input_cost_per_token": 3e-06}] if field == "tiered_pricing" else 3e-06 deployment = Deployment( model_name="my-model", - litellm_params=LiteLLM_Params(model="gpt-4o", **{field: 3e-06}), + litellm_params=LiteLLM_Params(model="gpt-4o", **{field: value}), ) - assert getattr(deployment.model_info, field) == 3e-06 - assert deployment.model_info.model_dump(exclude_none=True)[field] == 3e-06 + assert getattr(deployment.model_info, field) == value + assert deployment.model_info.model_dump(exclude_none=True)[field] == value + + +def test_deployment_mirrors_tiered_pricing_onto_model_info(): + """ + Regression: tiered_pricing set under a deployment's litellm_params was silently + ignored at cost time because the Deployment mirror excluded it, so the logging + path never flagged the deployment as custom-priced. + """ + tiers = [ + {"range": [0, 3000], "input_cost_per_token": 3.25e-07, "output_cost_per_token": 1.95e-06}, + {"range": [3000, 128000], "input_cost_per_token": 6.5e-07, "output_cost_per_token": 3.9e-06}, + ] + deployment = Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="anthropic/claude-haiku-4-5", tiered_pricing=tiers), + ) + assert deployment.model_info.tiered_pricing == tiers def test_unset_pricing_is_still_absent_from_dumps(): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 603ee0c5396..3b4b96bdd4e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35830,6 +35830,10 @@ export interface components { team_public_model_name?: string | null; /** Tier */ tier?: ("free" | "paid") | null; + /** Tiered Pricing */ + tiered_pricing?: { + [key: string]: unknown; + }[] | null; /** Updated At */ updated_at?: string | null; /** Updated By */