From 3d1f7f354409e7ca2fc841309e4270424ae5750f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 7 Jun 2026 03:17:21 +0000 Subject: [PATCH] fix(ci): handle Vercel AI Gateway models without input/output token pricing The weekly auto_update_price_and_context_window job crashes with KeyError: 'output' when the Vercel AI Gateway /v1/models response includes models whose pricing lacks the 'output' key. The current upstream payload has ~88 such rows: embedding models (input-only), image generators, video generators, and a handful with empty pricing. Skip rows that have neither input nor output token pricing (image and video models do not map to token-based cost fields anyway) and fall back to 0.0 for the standard chat/embedding rows so embeddings get included with input-only pricing instead of aborting the whole sync. Add a regression test covering the embedding-only, image-only, video-only, empty-pricing, and standard chat-with-cache shapes. Co-authored-by: Krrish Dholakia --- ...to_update_price_and_context_window_file.py | 23 +-- ...to_update_price_and_context_window_file.py | 142 ++++++++++++++++++ 2 files changed, 155 insertions(+), 10 deletions(-) create mode 100644 tests/test_litellm/test_auto_update_price_and_context_window_file.py diff --git a/.github/workflows/auto_update_price_and_context_window_file.py b/.github/workflows/auto_update_price_and_context_window_file.py index 461d8d347d9..5278c46dd61 100644 --- a/.github/workflows/auto_update_price_and_context_window_file.py +++ b/.github/workflows/auto_update_price_and_context_window_file.py @@ -85,24 +85,27 @@ def transform_openrouter_data(data): def transform_vercel_ai_gateway_data(data): transformed = {} for row in data: + pricing = row.get("pricing") or {} + + if "input" not in pricing and "output" not in pricing: + continue + obj = { "max_tokens": row["context_window"], - "input_cost_per_token": float(row["pricing"]["input"]), - "output_cost_per_token": float(row["pricing"]["output"]), + "input_cost_per_token": float(pricing.get("input", 0)), + "output_cost_per_token": float(pricing.get("output", 0)), 'max_output_tokens': row['max_tokens'], 'max_input_tokens': row["context_window"], } - # Handle cache pricing if available - if "pricing" in row: - if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None: - obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}") - - if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None: - obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}") + if pricing.get("input_cache_read") is not None: + obj['cache_read_input_token_cost'] = float(f"{float(pricing['input_cache_read']):e}") + + if pricing.get("input_cache_write") is not None: + obj['cache_creation_input_token_cost'] = float(f"{float(pricing['input_cache_write']):e}") mode = "embedding" if "embedding" in row["id"].lower() else "chat" - + obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode}) transformed[f'vercel_ai_gateway/{row["id"]}'] = obj diff --git a/tests/test_litellm/test_auto_update_price_and_context_window_file.py b/tests/test_litellm/test_auto_update_price_and_context_window_file.py new file mode 100644 index 00000000000..a9499ed6cd3 --- /dev/null +++ b/tests/test_litellm/test_auto_update_price_and_context_window_file.py @@ -0,0 +1,142 @@ +import importlib.util +import os +import sys + +import pytest + +SCRIPT_PATH = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + ".github", + "workflows", + "auto_update_price_and_context_window_file.py", + ) +) + + +@pytest.fixture(scope="module") +def transform_vercel_ai_gateway_data(): + spec = importlib.util.spec_from_file_location( + "auto_update_price_and_context_window_file", SCRIPT_PATH + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module.transform_vercel_ai_gateway_data + + +def test_transform_skips_image_video_and_empty_pricing_rows( + transform_vercel_ai_gateway_data, +): + data = [ + { + "id": "vendor/image-model", + "context_window": 0, + "max_tokens": 0, + "pricing": {"image": "0.05"}, + }, + { + "id": "vendor/video-model", + "context_window": 0, + "max_tokens": 0, + "pricing": {"video_duration_pricing": "0.10"}, + }, + { + "id": "vendor/no-pricing-model", + "context_window": 0, + "max_tokens": 0, + "pricing": {}, + }, + { + "id": "vendor/missing-pricing-model", + "context_window": 0, + "max_tokens": 0, + }, + ] + + assert transform_vercel_ai_gateway_data(data) == {} + + +def test_transform_includes_embedding_rows_with_input_only( + transform_vercel_ai_gateway_data, +): + data = [ + { + "id": "vendor/qwen3-embedding-0.6b", + "context_window": 32768, + "max_tokens": 32768, + "pricing": {"input": "0.00000001"}, + } + ] + + result = transform_vercel_ai_gateway_data(data) + + key = "vercel_ai_gateway/vendor/qwen3-embedding-0.6b" + assert key in result + entry = result[key] + assert entry["mode"] == "embedding" + assert entry["input_cost_per_token"] == pytest.approx(1e-8) + assert entry["output_cost_per_token"] == 0.0 + assert entry["max_tokens"] == 32768 + assert entry["max_input_tokens"] == 32768 + assert entry["max_output_tokens"] == 32768 + assert entry["litellm_provider"] == "vercel_ai_gateway" + + +def test_transform_standard_chat_row_with_cache_pricing( + transform_vercel_ai_gateway_data, +): + data = [ + { + "id": "vendor/chat-model", + "context_window": 128000, + "max_tokens": 16384, + "pricing": { + "input": "0.000001", + "output": "0.000002", + "input_cache_read": "0.00000025", + "input_cache_write": "0.0000005", + }, + } + ] + + result = transform_vercel_ai_gateway_data(data) + + entry = result["vercel_ai_gateway/vendor/chat-model"] + assert entry["mode"] == "chat" + assert entry["input_cost_per_token"] == pytest.approx(1e-6) + assert entry["output_cost_per_token"] == pytest.approx(2e-6) + assert entry["cache_read_input_token_cost"] == pytest.approx(2.5e-7) + assert entry["cache_creation_input_token_cost"] == pytest.approx(5e-7) + + +def test_transform_mixed_payload_does_not_raise(transform_vercel_ai_gateway_data): + data = [ + { + "id": "vendor/image-only", + "context_window": 0, + "max_tokens": 0, + "pricing": {"image": "0.05"}, + }, + { + "id": "vendor/chat-model", + "context_window": 8192, + "max_tokens": 4096, + "pricing": {"input": "0.000001", "output": "0.000002"}, + }, + { + "id": "vendor/embedding-model", + "context_window": 8192, + "max_tokens": 8192, + "pricing": {"input": "0.00000001"}, + }, + ] + + result = transform_vercel_ai_gateway_data(data) + + assert set(result.keys()) == { + "vercel_ai_gateway/vendor/chat-model", + "vercel_ai_gateway/vendor/embedding-model", + }