diff --git a/.github/scripts/auto_update_price_and_context_window_file.py b/.github/scripts/auto_update_price_and_context_window_file.py index 4c9a05bd4f04..e42d6bd198db 100644 --- a/.github/scripts/auto_update_price_and_context_window_file.py +++ b/.github/scripts/auto_update_price_and_context_window_file.py @@ -2,7 +2,7 @@ import aiohttp import json import math -from typing import Any +from typing import Any, Final # Asynchronously fetch data from a given URL async def fetch_data(url): @@ -242,12 +242,15 @@ def transform_vercel_ai_gateway_data(data): row.get("pricing", {}).get(k) is None for k in ("input", "output") ): continue - obj = { + limits: Final = { "max_tokens": row["context_window"], + "max_output_tokens": row["max_tokens"], + "max_input_tokens": row["context_window"], + } + obj = { "input_cost_per_token": float(row["pricing"]["input"]), "output_cost_per_token": float(row["pricing"]["output"]), - 'max_output_tokens': row['max_tokens'], - 'max_input_tokens': row["context_window"], + **{key: value for key, value in limits.items() if value > 0}, } # Handle cache pricing if available @@ -258,7 +261,8 @@ def transform_vercel_ai_gateway_data(data): 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}") - mode = "embedding" if "embedding" in row["id"].lower() else "chat" + is_embedding: Final = row.get("type") == "embedding" or "embed" in row["id"].lower() + mode = "embedding" if is_embedding else "chat" obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode}) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f34606ce1272..30f828b0ea24 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48389,10 +48389,7 @@ "vercel_ai_gateway/amazon/titan-embed-text-v2": { "input_cost_per_token": 2e-08, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/anthropic/claude-3-haiku": { @@ -48733,10 +48730,7 @@ "vercel_ai_gateway/cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/deepseek/deepseek-r1": { @@ -48991,10 +48985,7 @@ "vercel_ai_gateway/mistral/codestral-embed": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/mistral/devstral-small": { @@ -49057,10 +49048,7 @@ "vercel_ai_gateway/mistral/mistral-embed": { "input_cost_per_token": 1e-07, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/mistral/mistral-large": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f34606ce1272..30f828b0ea24 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48389,10 +48389,7 @@ "vercel_ai_gateway/amazon/titan-embed-text-v2": { "input_cost_per_token": 2e-08, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/anthropic/claude-3-haiku": { @@ -48733,10 +48730,7 @@ "vercel_ai_gateway/cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/deepseek/deepseek-r1": { @@ -48991,10 +48985,7 @@ "vercel_ai_gateway/mistral/codestral-embed": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/mistral/devstral-small": { @@ -49057,10 +49048,7 @@ "vercel_ai_gateway/mistral/mistral-embed": { "input_cost_per_token": 1e-07, "litellm_provider": "vercel_ai_gateway", - "max_input_tokens": 0, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0.0 }, "vercel_ai_gateway/mistral/mistral-large": { diff --git a/tests/unit/test_auto_update_price_and_context_window_file.py b/tests/unit/test_auto_update_price_and_context_window_file.py index 435747e9a09e..347044988bdf 100644 --- a/tests/unit/test_auto_update_price_and_context_window_file.py +++ b/tests/unit/test_auto_update_price_and_context_window_file.py @@ -4,8 +4,11 @@ from __future__ import annotations import importlib.util +import json import sys from pathlib import Path +from types import ModuleType +from typing import Final import pytest @@ -220,3 +223,74 @@ def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module): assert entry["supports_pdf_input"] is True assert entry["supports_assistant_prefill"] is True assert entry["input_cost_per_token"] == 1.5e-07 + + +@pytest.mark.parametrize( + "row_id,row_type,expected_mode", + [ + ("cohere/embed-v4.0", "embedding", "embedding"), + ("mistral/mistral-embed", "embedding", "embedding"), + ("some/vendor-vectors", "embedding", "embedding"), + ("amazon/titan-embed-text-v2", None, "embedding"), + ("openai/gpt-4o", "language", "chat"), + ], +) +def test_vercel_transform_marks_embedding_models_by_declared_type( + sync_module: ModuleType, row_id: str, row_type: str | None, expected_mode: str +) -> None: + row: Final[dict[str, object]] = { + "id": row_id, + "context_window": 8192, + "max_tokens": 8192, + "pricing": {"input": "0.0000001", "output": "0"}, + **({"type": row_type} if row_type is not None else {}), + } + transformed: Final = sync_module.transform_vercel_ai_gateway_data([row]) + assert transformed[f"vercel_ai_gateway/{row_id}"]["mode"] == expected_mode + + +def test_vercel_transform_omits_zero_token_limits(sync_module: ModuleType) -> None: + row: Final[dict[str, object]] = { + "id": "cohere/embed-v4.0", + "type": "embedding", + "context_window": 0, + "max_tokens": 0, + "pricing": {"input": "0.0000001", "output": "0"}, + } + transformed: Final = sync_module.transform_vercel_ai_gateway_data([row]) + entry: Final = transformed["vercel_ai_gateway/cohere/embed-v4.0"] + assert entry["mode"] == "embedding" + assert entry["input_cost_per_token"] == 1e-07 + assert {"max_tokens", "max_input_tokens", "max_output_tokens"}.isdisjoint(entry) + + +def test_vercel_transform_keeps_positive_token_limits(sync_module: ModuleType) -> None: + row: Final[dict[str, object]] = { + "id": "openai/gpt-4o", + "type": "language", + "context_window": 128000, + "max_tokens": 16384, + "pricing": {"input": "0.0000025", "output": "0.00001"}, + } + transformed: Final = sync_module.transform_vercel_ai_gateway_data([row]) + entry: Final = transformed["vercel_ai_gateway/openai/gpt-4o"] + assert entry["max_tokens"] == 128000 + assert entry["max_input_tokens"] == 128000 + assert entry["max_output_tokens"] == 16384 + + +def test_vercel_embedding_entries_in_cost_map_are_not_labeled_chat() -> None: + root: Final = Path(__file__).resolve().parents[2] + for name in ( + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", + ): + cost_map: Final = json.loads((root / name).read_text()) + mislabeled: Final = [ + key + for key, entry in cost_map.items() + if key.startswith("vercel_ai_gateway/") + and "embed" in key.lower() + and entry.get("mode") != "embedding" + ] + assert mislabeled == [], f"{name}: {mislabeled}"