Skip to content

fix(vercel_ai_gateway): label embedding models as embedding, not chat #45358

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
thefrojex wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(vercel_ai_gateway): label embedding models as embedding, not chat #45358

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions .github/scripts/auto_update_price_and_context_window_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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})

Expand Down
20 changes: 4 additions & 16 deletions litellm/model_prices_and_context_window_backup.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down
20 changes: 4 additions & 16 deletions model_prices_and_context_window.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Comment thread
thefrojex marked this conversation as resolved.
"output_cost_per_token": 0.0
},
"vercel_ai_gateway/anthropic/claude-3-haiku": {
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down
74 changes: 74 additions & 0 deletions tests/unit/test_auto_update_price_and_context_window_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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}"
Loading