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

[thefrojex]

TLDR

Problem this solves:

• 4 Vercel AI Gateway embedding models sit in the cost map as "mode": "chat"
• Their max_*_tokens are a bogus 0
• Health checks for them fall back to the cost map mode, so they resolve as chat

How it solves it:

• The sync script now uses the gateway's declared type == "embedding"
• An "embed" id match stays as a fallback
• The sync only writes a token limit when it is above 0, so zeros can't come back
• The 4 cost-map entries are corrected to embedding, with the bogus 0 limits removed

User Flow

Before: an operator proxies a Vercel embedding model without setting model_info.mode, and the proxy health check resolves it as a chat model

1. They add vercel_ai_gateway/cohere/embed-v4.0 to the proxy config with no model_info.mode
2. The proxy health check resolves this model's mode to chat
3. The model's info reports mode: chat and max_input_tokens: 0

After: the same deployment is recognised as an embedding model

1. They add vercel_ai_gateway/cohere/embed-v4.0 to the proxy config with no model_info.mode
2. The proxy health check resolves this model's mode to embedding
3. The model's info reports mode: embedding and no 0 token limit

Problem

Four Vercel AI Gateway embedding models are in the cost map as "mode": "chat" with max_input_tokens: 0:

• vercel_ai_gateway/cohere/embed-v4.0
• vercel_ai_gateway/mistral/mistral-embed
• vercel_ai_gateway/mistral/codestral-embed
• vercel_ai_gateway/amazon/titan-embed-text-v2

Root cause: transform_vercel_ai_gateway_data in .github/scripts/auto_update_price_and_context_window_file.py picks the mode with "embedding" in row["id"].lower(). None of these ids contain the word "embedding" (they say embed), so they were classified as chat. The gateway's /v1/models response already declares "type": "embedding" for them, but that field was ignored

Impact: anything that reads mode from the cost map treats these as chat models. resolve_health_check_mode in litellm/proxy/health_check.py falls back to get_model_info(...)["mode"], so a proxy deployment of one of these models without an explicit model_info.mode is health-checked with a chat request. The 0 token limits are also wrong and could trip context-window checks

Fix

• Sync script: use the gateway's declared type == "embedding", and keep an "embed" id fallback
• Sync script: only set max_tokens, max_input_tokens and max_output_tokens when the value is above 0, so a sync can't write zero limits back into the catalog
• Cost maps (model_prices_and_context_window.json and the backup): set mode: "embedding" on the 4 entries and remove the bogus 0 max_*_tokens values rather than guessing limits

Tests

Added to tests/unit/test_auto_update_price_and_context_window_file.py:

• Parametrized test that the transform labels embedding rows by declared type, and keeps language models as chat
• Test that an embedding row with context_window 0 and max_tokens 0 produces no limit fields
• Test that positive limits are still written
• Guard test that no vercel_ai_gateway/*embed* entry in either shipped cost map has a non-embedding mode

Without the fix: 6 failed, 15 passed. With the fix: 21 passed

Pre-Submission checklist

• I have added meaningful tests
• The test file covering my change passes locally: uv run pytest tests/unit/test_auto_update_price_and_context_window_file.py -v (21 passed)
• make lint passes locally
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Setup for both runs: litellm installed editable from the checkout, LITELLM_LOCAL_MODEL_COST_MAP=True so the local cost map is read, not the remote one

import litellm
from litellm.proxy.health_check import resolve_health_check_mode

m = "vercel_ai_gateway/cohere/embed-v4.0"
info = litellm.get_model_info(m)
print("get_model_info mode:", info.get("mode"))
print("get_model_info max_input_tokens:", info.get("max_input_tokens"))
print("resolve_health_check_mode:", resolve_health_check_mode({}, {"model": m}))

Before (7dd9aff)

vercel_ai_gateway/cohere/embed-v4.0

1. Run the script above on main
2. Output:

get_model_info mode: chat
get_model_info max_input_tokens: 0
resolve_health_check_mode: chat

After (0c13188)

vercel_ai_gateway/cohere/embed-v4.0

1. Run the same script on this branch
2. Output:

get_model_info mode: embedding
get_model_info max_input_tokens: None
resolve_health_check_mode: embedding

Type

Bug Fix

Caveats (if any)

Medium

• Proof calls resolve_health_check_mode directly instead of curl against a live proxy, since no Vercel AI Gateway key is available and a live call would fail both before and after
• Only cohere/embed-v4.0 was run through the proof script; the other 3 entries are covered by the cost-map guard test
• The 4 entries now have no max_*_tokens, since real limits were not guessed

Low

• Any other gateway row with embed in its id is also labeled embedding by the fallback
• A gateway row that reports a limit of 0 now ends up with no limit field instead of a 0

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[greptile-apps]

Confidence Score: 5/5

[Medium risk] Updates model classification and pricing data for embedding models.

The PR appears safe to merge, with no outstanding actionable findings

What we checked:

• Missing limits remain readable: get_model_info reads these fields with None defaults, so absent fields do not cause a missing-key error
• Health checks select embeddings: resolve_health_check_mode falls back to catalog metadata. The health-check handler maps embedding to aembedding

Summary

Corrects four Vercel AI Gateway models to embedding in both shipped catalogs. The sync reads the declared model type, keeps the embed name fallback, and omits nonpositive token limits

• Adds tests for declared types, fallback names, zero limits, positive limits, and both catalogs
• The three supplied previous threads are unnumbered. Their issues are addressed: zero limits are filtered, redundant comments are removed, and new declarations have types and Final annotations
• thefrojex acknowledges that the fallback labels any ID containing embed as an embedding model and that zero limits become absent fields. These behaviors are intentional

_{Reviews (2) · Last reviewed commit: "fix(vercel_ai_gateway): label embedding ..." · Reviewed by Greptile}

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing thefrojex:fix/vercel-embedding-mode (0c13188) with main (f824d11)}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[thefrojex]

@greptileai review