mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
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:
parent
3448bf79f8
commit
3d1f7f3544
2 changed files with 155 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue