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 <krrish-berri-2@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-06-07 03:17:21 +00:00
parent 3448bf79f8
commit 3d1f7f3544
No known key found for this signature in database
2 changed files with 155 additions and 10 deletions

View file

@ -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

View file

@ -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",
}