Merge 0c13188b2b into 82ca10cc2d

.github/scripts/auto_update_price_and_context_window_file.py

@@ -2,7 +2,7 @@ import asyncio
 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})
 

litellm/model_prices_and_context_window_backup.json

@@ -49238,10 +49238,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": {
@@ -49582,10 +49579,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": {
@@ -49840,10 +49834,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": {
@@ -49906,10 +49897,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": {

model_prices_and_context_window.json

@@ -49238,10 +49238,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": {
@@ -49582,10 +49579,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": {
@@ -49840,10 +49834,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": {
@@ -49906,10 +49897,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": {

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