From 6cfcb6cd839c1d23c7a59f247b1900346b9b2cab Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 14:08:18 +0000 Subject: [PATCH 01/19] fix(vertex_ai): translate /v1/embeddings batch rows to Gemini embedding shape Vertex batch files sent every jsonl line through the generateContent transform, so embeddings rows went out as {"request": {"contents": [...]}} and Vertex rejected each one with "no such field: 'contents'"; the OpenAI "input" was dropped along the way too. Route lines by their own url: embeddings lines now emit the EmbedContentRequest shape (singular content, embed_content_config sibling, custom_id round-tripping through the top-level key), and matching output rows come back as OpenAI embeddings responses. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 253 +++++++++++++--- .../test_vertex_ai_files_transformation.py | 277 ++++++++++++++++++ 2 files changed, 491 insertions(+), 39 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index dd877b52eb8..516a3ca7184 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,6 +5,7 @@ import json import os import re import time +from collections.abc import Mapping from typing import ( Any, Callable, @@ -16,6 +17,7 @@ from typing import ( Tuple, Union, ) +from urllib.parse import unquote import httpx from httpx import Headers, Response @@ -51,6 +53,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.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -62,13 +67,26 @@ from litellm.types.llms.openai import ( ) from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import LlmProviders, ModelResponse +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 = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" +_VERTEX_BATCH_KEY_FIELD = "key" +_MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM = { + "outputDimensionality": "output_dimensionality", + "taskType": "task_type", + "title": "title", +} def _sanitize_gcp_label_value(value: str) -> str: @@ -131,6 +149,21 @@ def _set_litellm_batch_custom_id_labels(labels: Dict[str, str], custom_id: Any) labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk +def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> 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 str(key) + request_data = vertex_output_row.get("request") or {} + return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) + + def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw = labels.get("litellm_custom_id_raw") @@ -149,10 +182,156 @@ def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) +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: Mapping[str, str] | None = None, +) -> Mapping[str, Any]: + """ + 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": error, + } + + +def _transform_vertex_embeddings_batch_output_row_to_openai( + vertex_output_row: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + """ + Transforms one Vertex Gemini Embedding batch output row into an OpenAI batch + output row holding an `/v1/embeddings` response body. + + Example Vertex jsonl + {"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}}} + + `tokenCount` is serialized as a string by Vertex (int64 proto field), and the row + carries no `modelVersion`, so the model comes from the batch the row belongs to. + """ + custom_id = _get_litellm_batch_custom_id(vertex_output_row) + status = vertex_output_row.get("status", "") + if status: + return _openai_batch_output_row( + custom_id=custom_id, + error={"code": "vertex_ai_error", "message": status}, + ) + + vertex_response = vertex_output_row.get("response") or {} + token_count = int(vertex_response.get("tokenCount") or 0) + body = EmbeddingResponse( + model=model or "", + data=[ + Embedding( + embedding=vertex_response["embedding"]["values"], + index=0, + object="embedding", + ) + ], + usage=Usage(prompt_tokens=token_count, total_tokens=token_count), + ).model_dump() + return _openai_batch_output_row(custom_id=custom_id, body=body) + + +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_batch_jsonl_entry_to_vertex_embeddings_row( + openai_entry: Mapping[str, Any], +) -> Mapping[str, Any]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into a Vertex Gemini + Embedding batch row. + + Example Vertex jsonl + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}}, "embed_content_config": {"output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + + Note that `content` is singular (an `EmbedContentRequest`, not a + `GenerateContentRequest`), the per-row config is a sibling of `request` rather than + part of it, and the `custom_id` round-trips through the top-level `key`. + + API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + openai_request_body = openai_entry.get("body") or {} + 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") + + embed_content_request = transform_openai_input_gemini_embed_content( + input=embedding_input, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + embed_content_config = { + config_field: embed_content_request[gemini_param] + for gemini_param, config_field in _EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM.items() + if gemini_param in embed_content_request + } + + custom_id = openai_entry.get("custom_id") + return { + **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), + "request": {"content": embed_content_request["content"]}, + **({"embed_content_config": embed_content_config} if embed_content_config else {}), + } + + def _openai_batch_jsonl_entry_to_vertex_wrapped_request( openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> Dict[str, Any]: +) -> Mapping[str, Any]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -160,6 +339,9 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_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_row(openai_entry) + openai_request_body = openai_entry.get("body") or {} vertex_request_body = _transform_request_body( messages=openai_request_body.get("messages", []), @@ -629,6 +811,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): transformed_content = 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 @@ -650,7 +833,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return HttpxBinaryResponseContent(response=raw_response) def _try_transform_vertex_batch_output_to_openai( - self, content: bytes, logging_obj: Optional[LiteLLMLoggingObj] = None + self, + content: bytes, + logging_obj: LiteLLMLoggingObj | None = None, + model: str | None = None, ) -> bytes: """ Try to transform Vertex AI batch output to OpenAI format. @@ -692,7 +878,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 = json.loads(first_line) - is_vertex_batch_output = ( + is_vertex_batch_output = _is_vertex_embeddings_batch_output_row(first_row) or ( "request" in first_row and "response" in first_row and "processed_time" in first_row @@ -731,11 +917,19 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): output = bytearray() for line in itertools.chain([first_line], lines): try: - openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, + vertex_output_row = json.loads(line) + openai_output = ( + _transform_vertex_embeddings_batch_output_row_to_openai( + vertex_output_row=vertex_output_row, + model=model, + ) + if _is_vertex_embeddings_batch_output_row(vertex_output_row) + else self._transform_single_vertex_batch_output_to_openai( + vertex_output=vertex_output_row, + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, + ) ) except Exception: return content @@ -755,30 +949,22 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> Dict[str, Any]: + ) -> Mapping[str, Any]: """ 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 = vertex_output.get("request", {}) - labels = request_data.get("labels", {}) or {} - custom_id = _get_litellm_batch_custom_id_from_labels(labels) + custom_id = _get_litellm_batch_custom_id(vertex_output) # Check if there's an error status = vertex_output.get("status", "") has_error = 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", "message": status}, + ) # Transform successful response using existing transformation vertex_response = vertex_output.get("response", {}) @@ -804,24 +990,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): response_dict = 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": { + return _openai_batch_output_row( + custom_id=custom_id, + error={ "code": "transformation_error", "message": f"Failed to transform response: {str(e)}", }, - } + ) 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..636a3106617 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 @@ -1318,3 +1318,280 @@ 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_to_embed_content_config_sibling(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": "hello world", + "dimensions": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + ) + ] + ) + + assert row["embed_content_config"] == { + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + assert "output_dimensionality" not in row["request"] + + def test_should_omit_embed_content_config_when_no_params_given(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert "embed_content_config" not in row + + 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_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_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": { + "tokenCount": "2", + "embedding": {"values": [-0.015, 0.024]}, + }, + } + 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_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_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": { + "tokenCount": "2", + "embedding": {"values": [-0.015, 0.024]}, + }, + } + ], + ) + + 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 From e96614a39f45d8b6ffde5a3e2a05e63f6d73541d Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 14:44:25 +0000 Subject: [PATCH 02/19] fix(vertex_ai): put embed config inside the request and read live usage A live Vertex batch run showed the documented "embed_content_config" sibling of "request" is rejected by the API ("unsupported type"), failing the whole job rather than the row; the same fields inside the EmbedContentRequest succeed and honor output_dimensionality. Real output rows also report usage under response.usageMetadata.promptTokenCount, not the documented response.tokenCount, so every row came back with zero tokens. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 29 +++++++------ .../test_vertex_ai_files_transformation.py | 42 ++++++++++++++----- 2 files changed, 48 insertions(+), 23 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 516a3ca7184..1a5d40dabb2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -82,7 +82,7 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" _VERTEX_BATCH_KEY_FIELD = "key" _MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") -_EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM = { +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "outputDimensionality": "output_dimensionality", "taskType": "task_type", "title": "title", @@ -231,10 +231,11 @@ def _transform_vertex_embeddings_batch_output_row_to_openai( output row holding an `/v1/embeddings` response body. Example Vertex jsonl - {"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}}} + {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} - `tokenCount` is serialized as a string by Vertex (int64 proto field), and the row - carries no `modelVersion`, so the model comes from the batch the row belongs to. + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. The row carries no `modelVersion`, so the model comes from the batch it + belongs to. """ custom_id = _get_litellm_batch_custom_id(vertex_output_row) status = vertex_output_row.get("status", "") @@ -245,7 +246,8 @@ def _transform_vertex_embeddings_batch_output_row_to_openai( ) vertex_response = vertex_output_row.get("response") or {} - token_count = int(vertex_response.get("tokenCount") or 0) + usage_metadata = vertex_response.get("usageMetadata") or {} + token_count = int(usage_metadata.get("promptTokenCount") or vertex_response.get("tokenCount") or 0) body = EmbeddingResponse( model=model or "", data=[ @@ -296,11 +298,13 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( Embedding batch row. Example Vertex jsonl - {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}}, "embed_content_config": {"output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + {"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`), the per-row config is a sibling of `request` rather than - part of it, and the `custom_id` round-trips through the top-level `key`. + `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level + `key`. 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 """ @@ -314,17 +318,16 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( model=openai_request_body.get("model", ""), optional_params=openai_request_body, ) - embed_content_config = { - config_field: embed_content_request[gemini_param] - for gemini_param, config_field in _EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM.items() + embed_request_fields = { + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() if gemini_param in embed_content_request } custom_id = openai_entry.get("custom_id") return { **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), - "request": {"content": embed_content_request["content"]}, - **({"embed_content_config": embed_content_config} if embed_content_config else {}), + "request": {"content": embed_content_request["content"], **embed_request_fields}, } 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 636a3106617..3eff3083220 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 @@ -1359,7 +1359,11 @@ class TestVertexEmbeddingsBatchInputTranslation: assert "key" not in row - def test_should_map_openai_params_to_embed_content_config_sibling(self): + 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( @@ -1374,17 +1378,20 @@ class TestVertexEmbeddingsBatchInputTranslation: ] ) - assert row["embed_content_config"] == { - "output_dimensionality": 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", + }, } - assert "output_dimensionality" not in row["request"] - def test_should_omit_embed_content_config_when_no_params_given(self): + def test_should_omit_config_fields_when_no_params_given(self): (row,) = _wrap_entries([_embeddings_entry()]) - assert "embed_content_config" not in row + assert set(row["request"]) == {"content"} def test_should_translate_multimodal_gcs_input(self): (row,) = _wrap_entries( @@ -1465,8 +1472,8 @@ class TestVertexEmbeddingsBatchOutputTranslation: "key": "request-1", "request": {"content": {"parts": [{"text": "hello world"}]}}, "response": { - "tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, }, } row.update(overrides) @@ -1503,6 +1510,21 @@ class TestVertexEmbeddingsBatchOutputTranslation: 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/" @@ -1569,8 +1591,8 @@ class TestVertexEmbeddingsBatchOutputTranslation: "status": "", "processed_time": "2026-07-29T05:55:52.379528Z", "response": { - "tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, }, } ], From 627d2755da56d01a6f5838bec967005027b2f35d Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 15:06:58 +0000 Subject: [PATCH 03/19] fix(vertex_ai): fan array embeddings input out into one vertex row per element Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 221 +++++++++++++----- .../files/test_vertex_ai_files_streaming.py | 5 +- .../test_vertex_ai_files_transformation.py | 180 +++++++++++++- 3 files changed, 339 insertions(+), 67 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 1a5d40dabb2..78a16b002e7 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -66,7 +66,7 @@ from litellm.types.llms.openai import ( PathLike, ) from litellm.types.files import StreamingMediaUploadConfig -from litellm.types.llms.vertex_ai import GcsBucketResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput from litellm.types.utils import ( Embedding, EmbeddingResponse, @@ -87,6 +87,7 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "taskType": "task_type", "title": "title", } +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P.*)#(?P\d+)/(?P\d+)") def _sanitize_gcp_label_value(value: str) -> str: @@ -222,46 +223,93 @@ def _openai_batch_output_row( } -def _transform_vertex_embeddings_batch_output_row_to_openai( - vertex_output_row: Mapping[str, Any], +def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int]: + """ + Resolve `(custom_id, index within that custom_id)` 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 = _get_litellm_batch_custom_id(vertex_output_row) + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(key) + if match is None or int(match["total"]) < 2: + return key, 0 + return match["custom_id"], int(match["index"]) + + +def _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id: str, + vertex_output_rows: tuple[Mapping[str, Any], ...], model: str | None, ) -> Mapping[str, Any]: """ - Transforms one Vertex Gemini Embedding batch output row into an OpenAI batch - output row holding an `/v1/embeddings` response body. + 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}}} - Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as - a fallback. The row carries no `modelVersion`, so the model comes from the batch it - belongs to. + 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 element fails the + whole entry, since an OpenAI batch row is either a response or an error. Live rows + report usage under `usageMetadata`; the documented `tokenCount` is kept as a + fallback. Rows carry no `modelVersion`, so the model comes from the batch they + belong to. """ - custom_id = _get_litellm_batch_custom_id(vertex_output_row) - status = vertex_output_row.get("status", "") + 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", "message": status}, ) - vertex_response = vertex_output_row.get("response") or {} - usage_metadata = vertex_response.get("usageMetadata") or {} - token_count = int(usage_metadata.get("promptTokenCount") or vertex_response.get("tokenCount") or 0) + responses = tuple(row.get("response") or {} for row in vertex_output_rows) + token_count = sum( + int((response.get("usageMetadata") or {}).get("promptTokenCount") or response.get("tokenCount") or 0) + for response in responses + ) body = EmbeddingResponse( model=model or "", data=[ Embedding( - embedding=vertex_response["embedding"]["values"], - index=0, + 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[Mapping[str, Any], ...]: + """ + 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(row for _, row in 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=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. @@ -290,21 +338,49 @@ def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: return path == "embeddings" or path.endswith("/embeddings") -def _openai_batch_jsonl_entry_to_vertex_embeddings_row( - openai_entry: Mapping[str, Any], -) -> Mapping[str, Any]: +def _openai_embedding_input_elements( + embedding_input: GeminiEmbeddingInput, +) -> tuple[Union[str, List[str]], ...]: """ - Transforms a single OpenAI `/v1/embeddings` batch entry into a Vertex Gemini - Embedding batch row. + 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. Entries asking for a single embedding keep their bare `custom_id`. + """ + return custom_id if total < 2 else f"{custom_id}#{index}/{total}" + + +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`. 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. + `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 """ @@ -313,37 +389,58 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( if embedding_input is None: raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") - embed_content_request = transform_openai_input_gemini_embed_content( - input=embedding_input, - model=openai_request_body.get("model", ""), - optional_params=openai_request_body, + 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 ) - embed_request_fields = { - request_field: embed_content_request[gemini_param] - for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() - if gemini_param in embed_content_request - } - custom_id = openai_entry.get("custom_id") - return { - **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), - "request": {"content": embed_content_request["content"], **embed_request_fields}, - } + return tuple( + { + **( + {} + if custom_id is None + else { + _VERTEX_BATCH_KEY_FIELD: _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), + ) + } + ), + "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.items() + if gemini_param in embed_content_request + }, + }, + } + for index, embed_content_request in enumerate(embed_content_requests) + ) -def _openai_batch_jsonl_entry_to_vertex_wrapped_request( +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]], -) -> Mapping[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_row(openai_entry) + return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) openai_request_body = openai_entry.get("body") or {} vertex_request_body = _transform_request_body( @@ -361,7 +458,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[Union[str, bytes]]) -> Iterator[str]: @@ -459,10 +556,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() @@ -914,25 +1011,29 @@ 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: - vertex_output_row = json.loads(line) - openai_output = ( - _transform_vertex_embeddings_batch_output_row_to_openai( - vertex_output_row=vertex_output_row, - model=model, - ) - if _is_vertex_embeddings_batch_output_row(vertex_output_row) - else self._transform_single_vertex_batch_output_to_openai( - vertex_output=vertex_output_row, - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, - ) + openai_output = self._transform_single_vertex_batch_output_to_openai( + vertex_output=json.loads(line), + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, ) except Exception: return content 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 3eff3083220..73d1d6eeb5a 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 @@ -1054,14 +1054,15 @@ 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( + row + for entry in openai_jsonl_content + for row in _openai_batch_jsonl_entry_to_vertex_rows( entry, cfg._map_openai_to_vertex_params ) - for entry in openai_jsonl_content ] @@ -1424,6 +1425,86 @@ class TestVertexEmbeddingsBatchInputTranslation: 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_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( [ @@ -1569,6 +1650,95 @@ class TestVertexEmbeddingsBatchOutputTranslation: "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_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_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( [ From 3c979f0b471214929b38de0fe61573fc864b7906 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 15:23:14 +0000 Subject: [PATCH 04/19] test(vertex_ai): cover batch lines without a url staying on the chat path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_files_transformation.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 73d1d6eeb5a..aef7f8b4684 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 @@ -1526,6 +1526,24 @@ class TestVertexEmbeddingsBatchInputTranslation: 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( [ From bf723fa9c167f48731f68ebe6b3bcab7351f5a83 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 20:50:27 +0000 Subject: [PATCH 05/19] fix(vertex_ai): percent-encode the custom_id in fanned-out vertex batch keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 29 ++++--- .../test_vertex_ai_files_transformation.py | 75 +++++++++++++++++++ 2 files changed, 92 insertions(+), 12 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 78a16b002e7..90fc5fae082 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -17,7 +17,7 @@ from typing import ( Tuple, Union, ) -from urllib.parse import unquote +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -87,7 +87,7 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "taskType": "task_type", "title": "title", } -_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P.*)#(?P\d+)/(?P\d+)") +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") def _sanitize_gcp_label_value(value: str) -> str: @@ -160,7 +160,7 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: """ key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: - return str(key) + return unquote(str(key)) request_data = vertex_output_row.get("request") or {} return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) @@ -228,14 +228,17 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, Resolve `(custom_id, index within that custom_id)` 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. + element, tagged `#/` (see + `_vertex_batch_embeddings_key`), so the rows can be reassembled into a single OpenAI + response. """ - key = _get_litellm_batch_custom_id(vertex_output_row) - match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(key) - if match is None or int(match["total"]) < 2: - return key, 0 - return match["custom_id"], int(match["index"]) + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is None: + return _get_litellm_batch_custom_id(vertex_output_row), 0 + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(str(key)) + if match is None: + return unquote(str(key)), 0 + return unquote(match["custom_id"]), int(match["index"]) def _vertex_embeddings_rows_to_openai_batch_output_row( @@ -359,9 +362,11 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: 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. Entries asking for a single embedding keep their bare `custom_id`. + 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. """ - return custom_id if total < 2 else f"{custom_id}#{index}/{total}" + encoded_custom_id = quote(custom_id, safe="") + return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( 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 aef7f8b4684..f95a63e4421 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 @@ -1476,6 +1476,19 @@ class TestVertexEmbeddingsBatchInputTranslation: 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( @@ -1711,6 +1724,68 @@ class TestVertexEmbeddingsBatchOutputTranslation: 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( From ef614b7b5bcdf94472b876bea64ad17e5dfd4282 Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 14:35:47 +0000 Subject: [PATCH 06/19] refactor(vertex_ai): keep the batch embeddings translation within the LIT002 ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 118 +++++++++++------- 1 file changed, 71 insertions(+), 47 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 9363540fe1b..bbb97a1edc2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -74,11 +74,11 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" _VERTEX_BATCH_KEY_FIELD = "key" _MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") -_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { - "outputDimensionality": "output_dimensionality", - "taskType": "task_type", - "title": "title", -} +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = ( + ("outputDimensionality", "output_dimensionality"), + ("taskType", "task_type"), + ("title", "title"), +) _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") @@ -153,12 +153,15 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: 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") or {} - return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) + 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: dict[str, Any]) -> str: +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, Any] | None) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" + if not labels: + return "unknown" raw = labels.get("litellm_custom_id_raw") if raw: raw_chunks = [str(raw)] @@ -195,7 +198,8 @@ def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) def _openai_batch_output_row( custom_id: str, body: Mapping[str, Any] | None = None, - error: Mapping[str, str] | None = None, + error_code: str | None = None, + error_message: str = "", ) -> Mapping[str, Any]: """ One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set @@ -211,7 +215,7 @@ def _openai_batch_output_row( "request_id": body.get("id", ""), "body": body, }, - "error": error, + "error": None if error_code is None else {"code": error_code, "message": error_message}, } @@ -233,6 +237,19 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, return unquote(match["custom_id"]), int(match["index"]) +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], ...], @@ -247,23 +264,19 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( 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 element fails the - whole entry, since an OpenAI batch row is either a response or an error. Live rows - report usage under `usageMetadata`; the documented `tokenCount` is kept as a - fallback. Rows carry no `modelVersion`, so the model comes from the batch they - belong to. + whole entry, since an OpenAI batch row is either a response or an error. 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", "message": status}, + error_code="vertex_ai_error", + error_message=status, ) - responses = tuple(row.get("response") or {} for row in vertex_output_rows) - token_count = sum( - int((response.get("usageMetadata") or {}).get("promptTokenCount") or response.get("tokenCount") or 0) - for response in responses - ) + 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=[ @@ -361,6 +374,27 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: 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], ...]: @@ -381,7 +415,9 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ - openai_request_body = openai_entry.get("body") or {} + openai_request_body = openai_entry.get("body") + if not isinstance(openai_request_body, dict): + raise ValueError("`body` is required on /v1/embeddings batch requests, but was not provided") 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") @@ -400,27 +436,16 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( ) custom_id = openai_entry.get("custom_id") return tuple( - { - **( - {} - if custom_id is None - else { - _VERTEX_BATCH_KEY_FIELD: _vertex_batch_embeddings_key( - custom_id=str(custom_id), - index=index, - total=len(embed_content_requests), - ) - } + _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), ), - "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.items() - if gemini_param in embed_content_request - }, - }, - } + embed_content_request=embed_content_request, + ) for index, embed_content_request in enumerate(embed_content_requests) ) @@ -1008,7 +1033,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) - all_lines = itertools.chain([first_line], lines) + 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 @@ -1064,7 +1089,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if has_error: return _openai_batch_output_row( custom_id=custom_id, - error={"code": "vertex_ai_error", "message": status}, + error_code="vertex_ai_error", + error_message=status, ) # Transform successful response using existing transformation @@ -1096,8 +1122,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): except Exception as e: return _openai_batch_output_row( custom_id=custom_id, - error={ - "code": "transformation_error", - "message": f"Failed to transform response: {e!s}", - }, + error_code="transformation_error", + error_message=f"Failed to transform response: {e!s}", ) From 6517c1dc067172457bf002c8ef1667af2c1c1c6f Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 02:18:10 +0000 Subject: [PATCH 07/19] fix(cost): support cache_creation_input_token_cost in tiered pricing and make tier selection all-or-nothing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/generate_model_prices_schema.py | 1 + .../llm_cost_calc/tiered_pricing.py | 76 +---- .../litellm_core_utils/llm_cost_calc/utils.py | 35 +++ litellm/llms/dashscope/cost_calculator.py | 130 ++++---- model_prices_and_context_window.schema.json | 4 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 112 +++++++ .../test_dashscope_cost_calculator.py | 287 +++++++++++++----- .../test_litellm/test_model_prices_schema.py | 18 ++ tests/test_litellm/test_utils.py | 1 + 9 files changed, 445 insertions(+), 219 deletions(-) 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/litellm_core_utils/llm_cost_calc/tiered_pricing.py b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py index fb0f130a6cf..d4ce6abfcc4 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, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index b94851794f0..f0ec734bbc6 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, @@ -207,6 +211,33 @@ def _parse_above_token_threshold(key: str) -> float: return float(threshold_str.replace("k", "")) * (1000 if "k" in threshold_str else 1) +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. + """ + 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 + + cache_creation_cost: Final = tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") + return ( + tier_rate(tier, "input_cost_per_token"), + tier_rate(tier, "output_cost_per_token"), + 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 +257,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) diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 22a0d38d598..ea6d50f5b00 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -1,108 +1,100 @@ """ 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", + if tier is not None: + return (breakdown.completion_tokens * tier_rate(tier, "output_cost_per_token")) + ( + breakdown.reasoning_tokens * tier_rate(tier, "output_cost_per_reasoning_token", "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) + reasoning_cost: Final = _flat_rate(model_info, "output_cost_per_reasoning_token", "output_cost_per_token") return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) @@ -122,11 +114,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/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 158cdb45f6b..8531291d42d 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 @@ -523,6 +523,118 @@ 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_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..20549fbd0fb 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 @@ -41,7 +40,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 +53,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 +71,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 +86,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 +144,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 +192,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 +206,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 +221,137 @@ 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_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/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_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", From 2f6f5c49616745ab58bd6d2a905b8f8300d6aa1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:04:48 -0700 Subject: [PATCH 08/19] fix(cost): reach tiered pricing for models without top-level per-token rates --- litellm/cost_calculator.py | 6 +++++- tests/test_litellm/test_cost_calculator.py | 20 ++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6b6653c5646..eb17a53a46e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 3f024e2fd03..8378375c09c 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 From c9685b2a6d4b80b042f67bf26f660a13ae406835 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:21:15 -0700 Subject: [PATCH 09/19] fix(router): honor tiered_pricing set in a deployment's litellm_params --- litellm/types/utils.py | 2 +- tests/test_litellm/types/test_router.py | 24 ++++++++++++++++--- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++++ 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 354857f8d72..34be2afbdd7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3260,6 +3260,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): @@ -3327,7 +3328,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/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 1e46ae9c577..443401eca5d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35483,6 +35483,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 */ From 688575e5bf1aaea22eb088733ca3a054631742b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:26:29 -0700 Subject: [PATCH 10/19] fix(cost): bill reasoning tokens at the selected tier's reasoning rate --- .../litellm_core_utils/llm_cost_calc/utils.py | 56 +++++++++++++------ .../llm_cost_calc/test_llm_cost_calc_utils.py | 53 ++++++++++++++++++ 2 files changed, 93 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f0ec734bbc6..41ca8270b6c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -211,6 +211,24 @@ 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 + 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. @@ -219,12 +237,8 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, 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. """ - 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: + 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") @@ -887,10 +901,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 @@ -977,12 +996,17 @@ def get_token_type_cost_breakdown( 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 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 8531291d42d..4eea8a88843 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 @@ -635,6 +635,59 @@ def test_generic_cost_per_token_tiered_pricing_is_all_or_nothing(): litellm.model_cost.pop(model, 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" From 0d0c712df7872aff1fde85b78940dadd4296e50a Mon Sep 17 00:00:00 2001 From: milan Date: Sat, 15 Aug 2026 00:29:35 +0000 Subject: [PATCH 11/19] fix(vertex_ai): fail an embeddings batch entry whose fan-out came back incomplete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 39 +++++++++++++------ .../test_vertex_ai_files_transformation.py | 14 +++++++ 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index d5363a08a92..b7f91bfba0d 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -259,9 +259,10 @@ def _openai_batch_output_row( } -def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int]: +def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int, int]: """ - Resolve `(custom_id, index within that custom_id)` for a Vertex batch output row. + 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 @@ -270,11 +271,11 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, """ key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is None: - return _get_litellm_batch_custom_id(vertex_output_row), 0 + 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 - return unquote(match["custom_id"]), int(match["index"]) + 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: @@ -293,6 +294,8 @@ def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: 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: """ @@ -303,9 +306,11 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( {"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 element fails the - whole entry, since an OpenAI batch row is either a response or an error. Rows carry - no `modelVersion`, so the model comes from the batch they belong to. + 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: @@ -315,6 +320,16 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( 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( @@ -345,16 +360,18 @@ def _transform_vertex_embeddings_batch_output_to_openai( """ keyed_rows = tuple((_split_vertex_batch_key(row), row) for row in vertex_output_rows) grouped_rows = { - custom_id: tuple(row for _, row in group) + 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=grouped_rows[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) + for custom_id in dict.fromkeys(custom_id for (custom_id, _, _), _ in keyed_rows) ) 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 f95a63e4421..c280e44ff4d 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 @@ -1802,6 +1802,20 @@ class TestVertexEmbeddingsBatchOutputTranslation: 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_end_to_end_round_trip_a_fanned_out_embeddings_batch(self, config): first_row, second_row = _wrap_entries( [ From de77711cf953af010b38c53be5a862d40b2e99d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:37:07 -0700 Subject: [PATCH 12/19] test(vertex_ai): cover duplicated fan-out rows in embeddings batch reassembly Also ruff-formats the batch transformation test file, which the formatter gate flags once the file is touched. --- .../test_vertex_ai_files_transformation.py | 320 +++++------------- 1 file changed, 93 insertions(+), 227 deletions(-) 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 c280e44ff4d..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 @@ -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, @@ -1060,9 +970,7 @@ def _wrap_entries(openai_jsonl_content): return [ row for entry in openai_jsonl_content - for row in _openai_batch_jsonl_entry_to_vertex_rows( - entry, cfg._map_openai_to_vertex_params - ) + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ] @@ -1123,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 @@ -1134,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) @@ -1150,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 @@ -1201,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 @@ -1231,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) @@ -1269,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 @@ -1280,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") @@ -1427,9 +1313,7 @@ class TestVertexEmbeddingsBatchInputTranslation: 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": []})] - ) + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2", "input": []})]) def test_should_fan_an_input_array_out_into_one_row_per_element(self): """ @@ -1466,13 +1350,7 @@ class TestVertexEmbeddingsBatchInputTranslation: ] 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"]} - ) - ] - ) + (row,) = _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2", "input": ["only one"]})]) assert row["key"] == "request-1" @@ -1533,9 +1411,7 @@ class TestVertexEmbeddingsBatchInputTranslation: ] ) - assert row["request"]["contents"] == [ - {"role": "user", "parts": [{"text": "Hello"}]} - ] + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] assert row["request"]["labels"]["litellm_custom_id"] == "request-1" assert "key" not in row @@ -1553,9 +1429,7 @@ class TestVertexEmbeddingsBatchInputTranslation: ] ) - assert row["request"]["contents"] == [ - {"role": "user", "parts": [{"text": "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( @@ -1603,10 +1477,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: logging_obj=MagicMock(), litellm_params={}, ) - return [ - json.loads(line) - for line in result.response.content.decode("utf-8").split("\n") - ] + 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()]) @@ -1616,9 +1487,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: 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["data"] == [{"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"}] assert body["usage"]["prompt_tokens"] == 2 assert body["usage"]["total_tokens"] == 2 @@ -1645,20 +1514,14 @@ class TestVertexEmbeddingsBatchOutputTranslation: ) 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 - ) + (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={} - ) - ], + [self._vertex_embeddings_output_row(status="Failed to parse JSON into proto", response={})], ) assert result["custom_id"] == "request-1" @@ -1669,10 +1532,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: 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) - ], + [self._vertex_embeddings_output_row(key=f"request-{index}") for index in range(3)], ) assert [result["custom_id"] for result in results] == [ @@ -1751,9 +1611,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: "request-1#0/2", "request-1", ] - assert [ - result["response"]["body"]["data"][0]["embedding"] for result in results - ] == [[0.1], [0.2]] + 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( @@ -1782,9 +1640,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: ) assert result["custom_id"] == "request#1/1" - assert [ - embedding["embedding"] for embedding in result["response"]["body"]["data"] - ] == [[0.1], [0.3]] + 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.""" @@ -1792,9 +1648,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: 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={} - ), + self._vertex_embeddings_output_row(key="request-1#1/2", status="Quota exceeded", response={}), ], ) @@ -1809,11 +1663,25 @@ class TestVertexEmbeddingsBatchOutputTranslation: [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 [1] of the 2 requested" + "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): @@ -1842,9 +1710,7 @@ class TestVertexEmbeddingsBatchOutputTranslation: ) assert result["custom_id"] == "MyRequest-1" - assert [ - embedding["embedding"] for embedding in result["response"]["body"]["data"] - ] == [[0.1], [0.3]] + 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( From 0176e4b3f6141acd241ba15db0e1928373018c4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:04:01 -0700 Subject: [PATCH 13/19] refactor(cost): make the shared token-details parsers public parse_prompt_tokens_details and parse_completion_tokens_details are imported by four modules, so the leading underscore made every import a reportPrivateUsage violation --- litellm/batches/batch_utils.py | 4 ++-- litellm/cost_calculator.py | 4 ++-- .../litellm_core_utils/llm_cost_calc/utils.py | 16 +++++++--------- litellm/llms/anthropic/cost_calculation.py | 4 ++-- litellm/llms/dashscope/cost_calculator.py | 8 ++++---- 5 files changed, 17 insertions(+), 19 deletions(-) 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 eb17a53a46e..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 ( @@ -2163,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/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 440447809de..4a33424e88f 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -99,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"] @@ -519,7 +519,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( @@ -589,7 +589,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, @@ -826,7 +826,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) @@ -881,7 +881,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"] @@ -1006,9 +1006,7 @@ 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)) @@ -1030,7 +1028,7 @@ def get_token_type_cost_breakdown( 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 ea6d50f5b00..e22d3e06be1 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -12,8 +12,8 @@ from typing import Final 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, + parse_completion_tokens_details, + parse_prompt_tokens_details, ) from litellm.types.utils import ModelInfo, Usage from litellm.utils import get_model_info @@ -33,12 +33,12 @@ class TokenBreakdown: def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: - prompt_details: Final = _parse_prompt_tokens_details(usage) + 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) - reasoning_tokens: Final = _parse_completion_tokens_details(usage)["reasoning_tokens"] + reasoning_tokens: Final = parse_completion_tokens_details(usage)["reasoning_tokens"] completion_tokens: Final = max((usage.completion_tokens or 0) - reasoning_tokens, 0) return TokenBreakdown( From 5970754a8583dea139325906faacd72c9e398515 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:04:01 -0700 Subject: [PATCH 14/19] fix(ptu): clear a PTU deployment's tiered_pricing instead of zeroing it tiered_pricing is a list, so the 0.0 the flat-rate zeroing stores does not even validate. Supplying tiers alongside PTU config gets the same 400 as a flat rate; tiers already stored are dropped from both blobs --- .../model_management_endpoints.py | 22 +++++--- .../test_ptu_model_settings.py | 51 ++++++++++++++++++- 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 912e18150b3..45a15d0d1f1 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -342,14 +342,17 @@ 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 list, not a rate, so it is dropped from a +# PTU deployment (see _PTU_CLEARED_PRICING_FIELDS) rather than stored as zero. +_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_CLEARED_PRICING_FIELDS: Final = frozenset({"tiered_pricing"}) _PTU_ZEROED_PRICING: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)) _NO_PRICING_OVERRIDE: Final[Mapping[str, float]] = MappingProxyType({}) _EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE @@ -378,7 +381,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_CLEARED_PRICING_FIELDS if supplied.get(field)) + ) + ) if not priced: return raise HTTPException( @@ -448,7 +456,7 @@ def _ptu_pricing_delta( supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied) if zeroed: - return zeroed, frozenset() + return zeroed, _PTU_CLEARED_PRICING_FIELDS was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR) if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR: return _NO_PRICING_OVERRIDE, frozenset() @@ -466,11 +474,13 @@ 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 + cleared: Final = dict.fromkeys(_PTU_CLEARED_PRICING_FIELDS, None) + pricing_update: Final = MappingProxyType(dict(override, **cleared)) 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=pricing_update), + "model_info": model_params.model_info.model_copy(update=pricing_update), } ) ) 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..5af880038eb 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 @@ -777,6 +777,53 @@ 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_cleared_not_zeroed(self): + """tiered_pricing is a list, so the zero the other fields store would not even validate. + Left in place it would keep billing per token at the tier rates.""" + 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 is None + assert priced.model_info.tiered_pricing is None + + 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 "tiered_pricing" not in stored, 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" @@ -1098,8 +1145,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 is None assert all(written.litellm_params.get(field) == 0 for field in _PTU_ZEROED_PRICING_FIELDS) + assert written.litellm_params.tiered_pricing is None @pytest.mark.asyncio async def test_model_new_refuses_a_priced_ptu_deployment(self): From 3f64cbe41b08ef8058e3cd761ee107dfa3b0e298 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 15 Aug 2026 01:14:14 +0000 Subject: [PATCH 15/19] fix(ptu): empty a PTU deployment's tiered_pricing instead of dropping it Dropping it falls back to the public cost map's tier table, whose rates outrank the zeros written beside them, so a PTU deployment on a tiered model keeps billing its traffic per token. Stored empty, the tiers no longer apply and the zeros win --- .../model_management_endpoints.py | 39 +++++++++------ .../test_ptu_model_settings.py | 50 +++++++++++++++---- 2 files changed, 66 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 45a15d0d1f1..7fbbaf84422 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -345,16 +345,22 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: # 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. -# tiered_pricing is the one mirrored field that is a list, not a rate, so it is dropped from a -# PTU deployment (see _PTU_CLEARED_PRICING_FIELDS) rather than stored as zero. +# 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_CLEARED_PRICING_FIELDS: Final = frozenset({"tiered_pricing"}) -_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 @@ -367,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 @@ -384,7 +392,7 @@ def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supp 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_CLEARED_PRICING_FIELDS if supplied.get(field)) + + tuple(field for field in _PTU_EMPTIED_PRICING_FIELDS if supplied.get(field)) ) ) if not priced: @@ -403,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 @@ -440,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 @@ -456,13 +464,13 @@ def _ptu_pricing_delta( supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied) if zeroed: - return zeroed, _PTU_CLEARED_PRICING_FIELDS + return zeroed, frozenset() was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR) if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR: 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)) ) @@ -474,13 +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 - cleared: Final = dict.fromkeys(_PTU_CLEARED_PRICING_FIELDS, None) - pricing_update: Final = MappingProxyType(dict(override, **cleared)) + # 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=pricing_update), - "model_info": model_params.model_info.model_copy(update=pricing_update), + "litellm_params": model_params.litellm_params.model_copy(update=stored), + "model_info": model_params.model_info.model_copy(update=stored), } ) ) 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 5af880038eb..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) @@ -784,9 +789,10 @@ class TestPtuDeploymentsAreNotBilledPerToken: assert exc.value.status_code == 400 assert "tiered_pricing" in str(exc.value.detail) - def test_tiered_pricing_already_on_the_row_is_cleared_not_zeroed(self): - """tiered_pricing is a list, so the zero the other fields store would not even validate. - Left in place it would keep billing per token at the tier rates.""" + 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( @@ -801,8 +807,8 @@ class TestPtuDeploymentsAreNotBilledPerToken: ), ) ) - assert priced.litellm_params.tiered_pricing is None - assert priced.model_info.tiered_pricing is None + assert priced.litellm_params.tiered_pricing == [] + assert priced.model_info.tiered_pricing == [] written = update_db_model( db_model=Deployment( @@ -821,7 +827,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: ) for blob in ("model_info", "litellm_params"): stored = json.loads(written[blob]) - assert "tiered_pricing" not in stored, 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): @@ -957,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( @@ -1146,9 +1178,9 @@ 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 if field != "tiered_pricing") - assert written.model_info.tiered_pricing is None + 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 is None + assert written.litellm_params.tiered_pricing == [] @pytest.mark.asyncio async def test_model_new_refuses_a_priced_ptu_deployment(self): From c3e38a0b528239119256732a4858de0e8fdeeaa2 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 15 Aug 2026 03:05:13 +0000 Subject: [PATCH 16/19] fix(cost): fall back to the model output rate when a tier omits one Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 12 ++++++- litellm/llms/dashscope/cost_calculator.py | 12 +++++-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 35 +++++++++++++++++++ .../test_dashscope_cost_calculator.py | 20 +++++++++++ 4 files changed, 75 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 4a33424e88f..9d6ad8b6e39 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -226,6 +226,8 @@ def _get_tiered_reasoning_rate(model_info: ModelInfo, usage: Usage) -> float | N 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") @@ -236,15 +238,23 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, 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"), - tier_rate(tier, "output_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, diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index e22d3e06be1..0c42f77e8ac 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -88,12 +88,18 @@ def _calculate_completion_cost( model_info: ModelInfo, tier: dict | None, ) -> float: + # A tier declaring no output rate falls back to the model's own, since a table spelling out + # only input rates would otherwise serve every completion for free + output_cost: Final = ( + tier_rate(tier, "output_cost_per_token") + if tier is not None and "output_cost_per_token" in tier + else float(model_info.get("output_cost_per_token") or 0.0) + ) if tier is not None: - return (breakdown.completion_tokens * tier_rate(tier, "output_cost_per_token")) + ( - breakdown.reasoning_tokens * tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") + return (breakdown.completion_tokens * output_cost) + ( + breakdown.reasoning_tokens * (tier_rate(tier, "output_cost_per_reasoning_token") or output_cost) ) - output_cost: Final = float(model_info.get("output_cost_per_token") or 0.0) reasoning_cost: Final = _flat_rate(model_info, "output_cost_per_reasoning_token", "output_cost_per_token") return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) 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 bb70d09681e..821e9d23e89 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 @@ -699,6 +699,41 @@ def test_generic_cost_per_token_tiered_pricing_is_all_or_nothing(): 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_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.""" 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 20549fbd0fb..42577ec44e3 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -324,6 +324,26 @@ class TestDashscopeCostCalculator: 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_tiered_pricing_zero_input_falls_back_to_flat_rates(self): """ No tier can be selected without input tokens, so an empty-prompt request must From 34918d34f9367251d7e631e37909080a94bfa4ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:37:18 -0700 Subject: [PATCH 17/19] fix(cost): inherit the backend output rate when a deployment's tiers omit one --- litellm/router.py | 46 +++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 40 ++++++++++++++++ 2 files changed, 86 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index fb2af41dcf2..cc64d4992a8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7635,6 +7635,37 @@ 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. + """ + 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 is not None: + model_info["output_cost_per_token"] = backend_rate + def _create_deployment( self, deployment_info: dict, @@ -7670,6 +7701,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 +8404,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 +8639,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/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 821e9d23e89..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 @@ -734,6 +734,46 @@ def test_generic_cost_per_token_tier_without_an_output_rate_bills_the_model_rate 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.""" From e46f2ca62f8ba627dc3689f17ab63b10dc7b63ca Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 15 Aug 2026 03:42:23 +0000 Subject: [PATCH 18/19] fix(dashscope): honor the model reasoning rate when a tier omits output rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/dashscope/cost_calculator.py | 19 +++--- .../test_dashscope_cost_calculator.py | 61 ++++++++++++++++++- 2 files changed, 70 insertions(+), 10 deletions(-) diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 0c42f77e8ac..3b3328d02a7 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -88,19 +88,20 @@ def _calculate_completion_cost( model_info: ModelInfo, tier: dict | None, ) -> float: - # A tier declaring no output rate falls back to the model's own, since a table spelling out - # only input rates would otherwise serve every completion for free + # 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 is not None and "output_cost_per_token" in tier + if tier_declares_output else float(model_info.get("output_cost_per_token") or 0.0) ) - if tier is not None: - return (breakdown.completion_tokens * output_cost) + ( - breakdown.reasoning_tokens * (tier_rate(tier, "output_cost_per_reasoning_token") or output_cost) - ) - - reasoning_cost: Final = _flat_rate(model_info, "output_cost_per_reasoning_token", "output_cost_per_token") + tier_reasoning_cost: Final = tier_rate(tier, "output_cost_per_reasoning_token") if tier is not None else 0.0 + model_reasoning_cost: Final = ( + 0.0 if tier_declares_output else float(model_info.get("output_cost_per_reasoning_token") or 0.0) + ) + reasoning_cost: Final = tier_reasoning_cost or model_reasoning_cost or output_cost return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) 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 42577ec44e3..4188cb6a4f3 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -21,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: @@ -344,6 +348,61 @@ class TestDashscopeCostCalculator: 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_tiered_pricing_zero_input_falls_back_to_flat_rates(self): """ No tier can be selected without input tokens, so an empty-prompt request must From 18752c860c3b3e4b13547f092f8c295acff3613f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:05:22 -0700 Subject: [PATCH 19/19] fix(cost): honor explicit zero tier rates and skip synthesized backend output rates --- .../llm_cost_calc/tiered_pricing.py | 12 +++- litellm/llms/dashscope/cost_calculator.py | 12 ++-- litellm/router.py | 6 +- .../test_dashscope_cost_calculator.py | 53 ++++++++++++++++++ .../test_router_model_cost_isolation.py | 55 +++++++++++++++++++ 5 files changed, 129 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py index d4ce6abfcc4..9bcc2b1743c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py @@ -60,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/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 3b3328d02a7..771ce140f66 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -97,11 +97,15 @@ def _calculate_completion_cost( if tier_declares_output else float(model_info.get("output_cost_per_token") or 0.0) ) - tier_reasoning_cost: Final = tier_rate(tier, "output_cost_per_reasoning_token") if tier is not None else 0.0 - model_reasoning_cost: Final = ( - 0.0 if tier_declares_output else float(model_info.get("output_cost_per_reasoning_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 ) - reasoning_cost: Final = tier_reasoning_cost or model_reasoning_cost or output_cost return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) diff --git a/litellm/router.py b/litellm/router.py index cc64d4992a8..4a450cf3c7c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7649,7 +7649,9 @@ class Router: 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. + 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: @@ -7663,7 +7665,7 @@ class Router: 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 is not None: + if backend_rate: model_info["output_cost_per_token"] = backend_rate def _create_deployment( 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 4188cb6a4f3..6f5aaabae06 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -403,6 +403,59 @@ class TestDashscopeCostCalculator: 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 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