scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Add Vercel AI Gateway provider support

Browse source
This commit is contained in:
joshualipman123 2025-07-30 10:01:05 -07:00
parent 2da6d457ef
commit 1ea28ccf67
16 changed files with 3689 additions and 272 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
.github/workflows/auto_update_price_and_context_window_file.py vendored
View file

@ -43,8 +43,8 @@ def write_to_file(file_path, data):
# Print an error message if writing to file fails
print("Error updating JSON file:", e)
# Update the existing models and add the missing models
def transform_remote_data(data):
# Update the existing models and add the missing models for OpenRouter
def transform_openrouter_data(data):
transformed = {}
for row in data:
# Add the fields 'max_tokens' and 'input_cost_per_token'
@ -81,6 +81,27 @@ def transform_remote_data(data):
return transformed
# Update the existing models and add the missing models for Vercel AI Gateway
def transform_vercel_ai_gateway_data(data):
transformed = {}
for row in data:
obj = {
"max_tokens": row["context_window"],
"input_cost_per_token": float(row["pricing"]["input"]),
"output_cost_per_token": float(row["pricing"]["output"]),
}
obj['max_input_tokens'] = row["context_window"]
obj.update({
"litellm_provider": "vercel_ai_gateway",
"mode": "chat"
})
transformed[f'vercel_ai_gateway/{row["id"]}'] = obj
return transformed
# Load local data from a specified file
def load_local_data(file_path):
@ -100,22 +121,32 @@ def load_local_data(file_path):
def main():
local_file_path = "model_prices_and_context_window.json" # Path to the local data file
url = "https://openrouter.ai/api/v1/models" # URL to fetch remote data
openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data
vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data
# Load local data from file
local_data = load_local_data(local_file_path)
# Fetch remote data asynchronously
remote_data = asyncio.run(fetch_data(url))
# Transform the fetched remote data
remote_data = transform_remote_data(remote_data)
# Fetch OpenRouter data
openrouter_data = asyncio.run(fetch_data(openrouter_url))
# Transform the fetched OpenRouter data
openrouter_data = transform_openrouter_data(openrouter_data)
# Fetch Vercel AI Gateway data
vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
# Transform the fetched Vercel AI Gateway data
vercel_data = transform_vercel_ai_gateway_data(vercel_data)
# Combine both datasets
all_remote_data = {**openrouter_data, **vercel_data}
# If both local and remote data are available, synchronize and save
if local_data and remote_data:
sync_local_data_with_remote(local_data, remote_data)
# If both local and openrouter data are available, synchronize and save
if local_data and all_remote_data:
sync_local_data_with_remote(local_data, all_remote_data)
write_to_file(local_file_path, local_data)
else:
print("Failed to fetch model data from either local file or URL.")
# Entry point of the script
if __name__ == "__main__":
main()
main()

144
docs/my-website/docs/index.md
View file

@ -73,154 +73,20 @@ response = completion(
```
</TabItem>
<TabItem value="anthropic" label="Anthropic">
<TabItem value="vercel" label="Vercel AI Gateway">
```python
from litellm import completion
import os
## set ENV variables
os.environ["ANTHROPIC_API_KEY"] = "your-api-key"
os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key"
response = completion(
model="anthropic/claude-3-sonnet-20240229",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
</TabItem>
<TabItem value="xai" label="xAI">
```python
from litellm import completion
import os
## set ENV variables
os.environ["XAI_API_KEY"] = "your-api-key"
response = completion(
model="xai/grok-2-latest",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
</TabItem>
<TabItem value="vertex" label="VertexAI">
```python
from litellm import completion
import os
# auth: run 'gcloud auth application-default'
os.environ["VERTEXAI_PROJECT"] = "hardy-device-386718"
os.environ["VERTEXAI_LOCATION"] = "us-central1"
response = completion(
model="vertex_ai/gemini-1.5-pro",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
</TabItem>
<TabItem value="nvidia" label="NVIDIA">
```python
from litellm import completion
import os
## set ENV variables
os.environ["NVIDIA_NIM_API_KEY"] = "nvidia_api_key"
os.environ["NVIDIA_NIM_API_BASE"] = "nvidia_nim_endpoint_url"
response = completion(
model="nvidia_nim/<model_name>",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
</TabItem>
<TabItem value="hugging" label="HuggingFace">
```python
from litellm import completion
import os
os.environ["HUGGINGFACE_API_KEY"] = "huggingface_api_key"
# e.g. Call 'WizardLM/WizardCoder-Python-34B-V1.0' hosted on HF Inference endpoints
response = completion(
model="huggingface/WizardLM/WizardCoder-Python-34B-V1.0",
model="vercel_ai_gateway/openai/gpt-4o",
messages=[{ "content": "Hello, how are you?","role": "user"}],
api_base="https://my-endpoint.huggingface.cloud"
)
print(response)
```
</TabItem>
<TabItem value="azure" label="Azure OpenAI">
```python
from litellm import completion
import os
## set ENV variables
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
# azure call
response = completion(
"azure/<your_deployment_name>",
messages = [{ "content": "Hello, how are you?","role": "user"}]
)
```
</TabItem>
<TabItem value="ollama" label="Ollama">
```python
from litellm import completion
response = completion(
model="ollama/llama2",
messages = [{ "content": "Hello, how are you?","role": "user"}],
api_base="http://localhost:11434"
)
```
</TabItem>
<TabItem value="or" label="Openrouter">
```python
from litellm import completion
import os
## set ENV variables
os.environ["OPENROUTER_API_KEY"] = "openrouter_api_key"
response = completion(
model="openrouter/google/palm-2-chat-bison",
messages = [{ "content": "Hello, how are you?","role": "user"}],
)
```
</TabItem>
<TabItem value="novita" label="Novita AI">
```python
from litellm import completion
import os
## set ENV variables. Visit https://novita.ai/settings/key-management to get your API key
os.environ["NOVITA_API_KEY"] = "novita-api-key"
response = completion(
model="novita/deepseek/deepseek-r1",
messages=[{ "content": "Hello, how are you?","role": "user"}]
stream=True,
)
```

209
docs/my-website/docs/providers/vercel_ai_gateway.md Normal file
View file

@ -0,0 +1,209 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Vercel AI Gateway
## Overview
| Property | Details |
|-------|-------|
| Description | Vercel AI Gateway provides a unified interface to access multiple AI providers through a single endpoint, with built-in caching, rate limiting, and analytics. |
| Provider Route on LiteLLM | `vercel_ai_gateway/` |
| Link to Provider Doc | [Vercel AI Gateway Documentation ↗](https://vercel.com/docs/ai-gateway) |
| Base URL | `https://ai-gateway.vercel.sh/v1` |
| Supported Operations | `/chat/completions`, `/models` |
<br />
<br />
https://vercel.com/docs/ai-gateway
**We support ALL models available through Vercel AI Gateway, just set `vercel_ai_gateway/` as a prefix when sending completion requests**
## Required Variables
```python showLineNumbers title="Environment Variables"
os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "" # your Vercel AI Gateway API key
# OR
os.environ["VERCEL_OIDC_TOKEN"] = "" # your Vercel OIDC token for authentication
```
## Usage - LiteLLM Python SDK
### Non-streaming
```python showLineNumbers title="Vercel AI Gateway Non-streaming Completion"
import os
import litellm
from litellm import completion
os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key"
messages = [{"content": "Hello, how are you?", "role": "user"}]
# Vercel AI Gateway call
response = completion(
model="vercel_ai_gateway/openai/gpt-4o",
messages=messages
)
print(response)
```
### Streaming
```python showLineNumbers title="Vercel AI Gateway Streaming Completion"
import os
import litellm
from litellm import completion
os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key"
messages = [{"content": "Hello, how are you?", "role": "user"}]
# Vercel AI Gateway call with streaming
response = completion(
model="vercel_ai_gateway/openai/gpt-4o",
messages=messages,
stream=True
)
for chunk in response:
print(chunk)
```
## Usage - LiteLLM Proxy
Add the following to your LiteLLM Proxy configuration file:
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: gpt-4o-gateway
litellm_params:
model: vercel_ai_gateway/openai/gpt-4o
api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY
- model_name: claude-4-sonnet-gateway
litellm_params:
model: vercel_ai_gateway/anthropic/claude-4-sonnet
api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY
```
Start your LiteLLM Proxy server:
```bash showLineNumbers title="Start LiteLLM Proxy"
litellm --config config.yaml
# RUNNING on http://0.0.0.0:4000
```
<Tabs>
<TabItem value="openai-sdk" label="OpenAI SDK">
```python showLineNumbers title="Vercel AI Gateway via Proxy - Non-streaming"
from openai import OpenAI
# Initialize client with your proxy URL
client = OpenAI(
base_url="http://localhost:4000", # Your proxy URL
api_key="your-proxy-api-key" # Your proxy API key
)
# Non-streaming response
response = client.chat.completions.create(
model="gpt-4o-gateway",
messages=[{"role": "user", "content": "Hello, how are you?"}]
)
print(response.choices[0].message.content)
```
```python showLineNumbers title="Vercel AI Gateway via Proxy - Streaming"
from openai import OpenAI
# Initialize client with your proxy URL
client = OpenAI(
base_url="http://localhost:4000", # Your proxy URL
api_key="your-proxy-api-key" # Your proxy API key
)
# Streaming response
response = client.chat.completions.create(
model="gpt-4o-gateway",
messages=[{"role": "user", "content": "Hello, how are you?"}],
stream=True
)
for chunk in response:
if chunk.choices[0].delta.content is not None:
print(chunk.choices[0].delta.content, end="")
```
</TabItem>
<TabItem value="litellm-sdk" label="LiteLLM SDK">
```python showLineNumbers title="Vercel AI Gateway via Proxy - LiteLLM SDK"
import litellm
# Configure LiteLLM to use your proxy
response = litellm.completion(
model="litellm_proxy/gpt-4o-gateway",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_base="http://localhost:4000",
api_key="your-proxy-api-key"
)
print(response.choices[0].message.content)
```
```python showLineNumbers title="Vercel AI Gateway via Proxy - LiteLLM SDK Streaming"
import litellm
# Configure LiteLLM to use your proxy with streaming
response = litellm.completion(
model="litellm_proxy/gpt-4o-gateway",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_base="http://localhost:4000",
api_key="your-proxy-api-key",
stream=True
)
for chunk in response:
if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta.content is not None:
print(chunk.choices[0].delta.content, end="")
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="Vercel AI Gateway via Proxy - cURL"
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-proxy-api-key" \
-d '{
"model": "gpt-4o-gateway",
"messages": [{"role": "user", "content": "Hello, how are you?"}]
}'
```
```bash showLineNumbers title="Vercel AI Gateway via Proxy - cURL Streaming"
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-proxy-api-key" \
-d '{
"model": "gpt-4o-gateway",
"messages": [{"role": "user", "content": "Hello, how are you?"}],
"stream": true
}'
```
</TabItem>
</Tabs>
For more detailed information on using the LiteLLM Proxy, see the [LiteLLM Proxy documentation](../providers/litellm_proxy).
## Additional Resources
- [Vercel AI Gateway Documentation](https://vercel.com/docs/ai-gateway)

1
docs/my-website/sidebars.js
View file

@ -451,6 +451,7 @@ const sidebars = {
"providers/replicate",
"providers/togetherai",
"providers/v0",
"providers/vercel_ai_gateway",
"providers/morph",
"providers/lambda_ai",
"providers/novita",

7
litellm/__init__.py
View file

@ -209,6 +209,7 @@ vertex_location: Optional[str] = None
predibase_tenant_id: Optional[str] = None
togetherai_api_key: Optional[str] = None
cloudflare_api_key: Optional[str] = None
vercel_ai_gateway_key: Optional[str] = None
baseten_key: Optional[str] = None
llama_api_key: Optional[str] = None
aleph_alpha_key: Optional[str] = None
@ -433,6 +434,7 @@ mistral_chat_models: List = []
text_completion_codestral_models: List = []
anthropic_models: List = []
openrouter_models: List = []
vercel_ai_gateway_models: List = []
datarobot_models: List = []
vertex_language_models: List = []
vertex_vision_models: List = []
@ -549,6 +551,8 @@ def add_known_models():
empower_models.append(key)
elif value.get("litellm_provider") == "openrouter":
openrouter_models.append(key)
elif value.get("litellm_provider") == "vercel_ai_gateway":
vercel_ai_gateway_models.append(key)
elif value.get("litellm_provider") == "datarobot":
datarobot_models.append(key)
elif value.get("litellm_provider") == "vertex_ai-text-models":
@ -718,6 +722,7 @@ model_list = (
+ anthropic_models
+ replicate_models
+ openrouter_models
+ vercel_ai_gateway_models
+ datarobot_models
+ huggingface_models
+ vertex_chat_models
@ -787,6 +792,7 @@ models_by_provider: dict = {
"together_ai": together_ai_models,
"baseten": baseten_models,
"openrouter": openrouter_models,
"vercel_ai_gateway": vercel_ai_gateway_models,
"datarobot": datarobot_models,
"vertex_ai": vertex_chat_models
+ vertex_text_models
@ -1165,6 +1171,7 @@ from .llms.v0.chat.transformation import V0ChatConfig
from .llms.morph.chat.transformation import MorphChatConfig
from .llms.lambda_ai.chat.transformation import LambdaAIChatConfig
from .llms.hyperbolic.chat.transformation import HyperbolicChatConfig
from .llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig
from .main import * # type: ignore
from .integrations import *
from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients

3
litellm/constants.py
View file

@ -281,6 +281,7 @@ LITELLM_CHAT_PROVIDERS = [
"v0",
"morph",
"lambda_ai",
"vercel_ai_gateway",
]
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [
@ -413,6 +414,7 @@ openai_compatible_endpoints: List = [
"https://api.morphllm.com/v1",
"https://api.lambda.ai/v1",
"https://api.hyperbolic.xyz/v1",
"https://ai-gateway.vercel.sh/v1",
]
@ -454,6 +456,7 @@ openai_compatible_providers: List = [
"morph",
"lambda_ai",
"hyperbolic",
"vercel_ai_gateway",
]
openai_text_completion_compatible_providers: List = (
[ # providers that support `/v1/completions`

10
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -246,6 +246,9 @@ def get_llm_provider( # noqa: PLR0915
elif endpoint == "https://api.hyperbolic.xyz/v1":
custom_llm_provider = "hyperbolic"
dynamic_api_key = get_secret_str("HYPERBOLIC_API_KEY")
elif endpoint == "https://ai-gateway.vercel.sh/v1":
custom_llm_provider = "vercel_ai_gateway"
dynamic_api_key = get_secret_str("VERCEL_AI_GATEWAY_API_KEY")
if api_base is not None and not isinstance(api_base, str):
raise Exception(
@ -718,6 +721,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
) = litellm.HyperbolicChatConfig()._get_openai_compatible_provider_info(
api_base, api_key
)
elif custom_llm_provider == "vercel_ai_gateway":
(
api_base,
dynamic_api_key,
) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info(
api_base, api_key
)
if api_base is not None and not isinstance(api_base, str):
raise Exception("api base needs to be a string. api_base={}".format(api_base))

2
litellm/litellm_core_utils/get_supported_openai_params.py
View file

@ -125,6 +125,8 @@ def get_supported_openai_params( # noqa: PLR0915
return litellm.AzureOpenAIConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "openrouter":
return litellm.OpenrouterConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "vercel_ai_gateway":
return litellm.VercelAIGatewayConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "mistral" or custom_llm_provider == "codestral":
# mistal and codestral api have the exact same params
if request_type == "chat_completion":

112
litellm/llms/vercel_ai_gateway/chat/transformation.py Normal file
View file

@ -0,0 +1,112 @@
"""
Support for OpenAI's `/v1/chat/completions` endpoint.
Calls done in OpenAI/openai.py as Vercel AI Gateway is openai-compatible.
Docs: https://vercel.com/docs/ai-gateway
"""
from typing import Any, List, Optional, Tuple, Union
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import AllMessageValues
from litellm.secret_managers.main import get_secret_str
import litellm
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
from ..common_utils import VercelAIGatewayException
class VercelAIGatewayConfig(OpenAIGPTConfig):
@property
def custom_llm_provider(self) -> Optional[str]:
return "vercel_ai_gateway"
def get_supported_openai_params(self, model: str) -> list:
base_params = super().get_supported_openai_params(model)
if "extra_body" not in base_params:
base_params.append("extra_body")
return base_params
def _get_openai_compatible_provider_info(
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:
api_base = (
api_base
or get_secret_str("VERCEL_AI_GATEWAY_API_BASE")
or "https://ai-gateway.vercel.sh/v1"
)
user_api_key = (
api_key
or get_secret_str("VERCEL_AI_GATEWAY_API_KEY")
or get_secret_str("VERCEL_OIDC_TOKEN")
)
return api_base, user_api_key
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
mapped_openai_params = super().map_openai_params(
non_default_params, optional_params, model, drop_params
)
# Vercel AI Gateway-only parameters
extra_body = {}
provider_options = non_default_params.pop("providerOptions", None)
if provider_options is not None:
extra_body["providerOptions"] = provider_options
mapped_openai_params["extra_body"] = extra_body # openai client supports `extra_body` param
return mapped_openai_params
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the overall request to be sent to the API.
Returns:
dict: The transformed request. Sent as the body of the API call.
"""
return super().transform_request(
model, messages, optional_params, litellm_params, headers
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return VercelAIGatewayException(
message=error_message,
status_code=status_code,
headers=headers,
)
def get_models(
self, api_key: Optional[str] = None, api_base: Optional[str] = None
) -> List[str]:
api_base, _ = self._get_openai_compatible_provider_info(api_base, api_key)
if api_base is None:
api_base = "https://ai-gateway.vercel.sh/v1"
models_url = f"{api_base}/models"
response = litellm.module_level_client.get(url=models_url)
if response.status_code != 200:
raise Exception(f"Failed to get models: {response.text}")
models = response.json()["data"]
return [model["id"] for model in models]

5
litellm/llms/vercel_ai_gateway/common_utils.py Normal file
View file

@ -0,0 +1,5 @@
from litellm.llms.base_llm.chat.transformation import BaseLLMException
class VercelAIGatewayException(BaseLLMException):
pass

64
litellm/main.py
View file

@ -2561,6 +2561,70 @@ def completion( # type: ignore # noqa: PLR0915
logging.post_call(
input=messages, api_key=openai.api_key, original_response=response
)
elif custom_llm_provider == "vercel_ai_gateway":
api_base = (
api_base
or litellm.api_base
or get_secret_str("VERCEL_AI_GATEWAY_API_BASE")
or "https://ai-gateway.vercel.sh/v1"
)
api_key = (
api_key
or litellm.api_key
or get_secret("VERCEL_AI_GATEWAY_API_KEY")
)
vercel_site_url = get_secret("VERCEL_SITE_URL") or "https://litellm.ai"
vercel_app_name = get_secret("VERCEL_APP_NAME") or "liteLLM"
vercel_headers = {
"http-referer": vercel_site_url,
"x-title": vercel_app_name,
}
_headers = headers or litellm.headers
if _headers:
vercel_headers.update(_headers)
headers = vercel_headers
## Load Config
config = litellm.VercelAIGatewayConfig.get_config()
for k, v in config.items():
if k == "extra_body":
# we use openai 'extra_body' to pass vercel specific params - providerOptions
if "extra_body" in optional_params:
optional_params[k].update(v)
else:
optional_params[k] = v
elif k not in optional_params:
optional_params[k] = v
data = {"model": model, "messages": messages, **optional_params}
## COMPLETION CALL
response = base_llm_http_handler.completion(
model=model,
stream=stream,
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
custom_llm_provider="vercel_ai_gateway",
timeout=timeout,
headers=headers,
encoding=encoding,
api_key=api_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
client=client,
)
## LOGGING
logging.post_call(
input=messages, api_key=openai.api_key, original_response=response
)
elif (
custom_llm_provider == "together_ai"
or ("togethercomputer" in model)

2
litellm/types/utils.py
View file

@ -2319,7 +2319,7 @@ class LlmProviders(str, Enum):
HYPERBOLIC = "hyperbolic"
RECRAFT = "recraft"
AUTO_ROUTER = "auto_router"
VERCEL_AI_GATEWAY = "vercel_ai_gateway"
# Create a set of all provider values for quick lookup
LlmProvidersSet = {provider.value for provider in LlmProviders}

28
litellm/utils.py
View file

@ -2328,6 +2328,9 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915
split_string = key.split("/", 1)
if key not in litellm.openrouter_models:
litellm.openrouter_models.append(split_string[1])
elif value.get("litellm_provider") == "vercel_ai_gateway":
if key not in litellm.vercel_ai_gateway_models:
litellm.vercel_ai_gateway_models.append(key)
elif value.get("litellm_provider") == "vertex_ai-text-models":
if key not in litellm.vertex_text_models:
litellm.vertex_text_models.append(key)
@ -3150,6 +3153,7 @@ def pre_process_optional_params(
and custom_llm_provider != "bedrock"
and custom_llm_provider != "ollama_chat"
and custom_llm_provider != "openrouter"
and custom_llm_provider != "vercel_ai_gateway"
and custom_llm_provider != "nebius"
and custom_llm_provider not in litellm.openai_compatible_providers
):
@ -3826,6 +3830,17 @@ def get_optional_params( # noqa: PLR0915
else False
),
)
elif custom_llm_provider == "vercel_ai_gateway":
optional_params = litellm.VercelAIGatewayConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(
drop_params
if drop_params is not None and isinstance(drop_params, bool)
else False
),
)
elif custom_llm_provider == "watsonx":
optional_params = litellm.IBMWatsonXChatConfig().map_openai_params(
@ -5209,6 +5224,11 @@ def validate_environment( # noqa: PLR0915
keys_in_environment = True
else:
missing_keys.append("OPENROUTER_API_KEY")
elif custom_llm_provider == "vercel_ai_gateway":
if "VERCEL_AI_GATEWAY_API_KEY" in os.environ:
keys_in_environment = True
else:
missing_keys.append("VERCEL_AI_GATEWAY_API_KEY")
elif custom_llm_provider == "datarobot":
if "DATAROBOT_API_TOKEN" in os.environ:
keys_in_environment = True
@ -5426,6 +5446,12 @@ def validate_environment( # noqa: PLR0915
keys_in_environment = True
else:
missing_keys.append("OPENROUTER_API_KEY")
## vercel_ai_gateway
elif model in litellm.vercel_ai_gateway_models:
if "VERCEL_AI_GATEWAY_API_KEY" in os.environ:
keys_in_environment = True
else:
missing_keys.append("VERCEL_AI_GATEWAY_API_KEY")
## datarobot
elif model in litellm.datarobot_models:
if "DATAROBOT_API_TOKEN" in os.environ:
@ -6790,6 +6816,8 @@ class ProviderConfigManager:
return litellm.TogetherAIConfig()
elif litellm.LlmProviders.OPENROUTER == provider:
return litellm.OpenrouterConfig()
elif litellm.LlmProviders.VERCEL_AI_GATEWAY == provider:
return litellm.VercelAIGatewayConfig()
elif litellm.LlmProviders.DATAROBOT == provider:
return litellm.DataRobotConfig()
elif litellm.LlmProviders.GEMINI == provider:

3004
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

205
tests/llm_translation/test_vercel_ai_gateway.py Normal file
View file

@ -0,0 +1,205 @@
"""
Tests for vercel_ai_gateway provider integration
"""
import os
from unittest import mock
import pytest
import litellm
from litellm import completion
from litellm.llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig
def test_vercel_ai_gateway_config_initialization():
"""Test VercelAIGatewayConfig initializes correctly"""
config = VercelAIGatewayConfig()
assert config.custom_llm_provider == "vercel_ai_gateway"
def test_get_llm_provider_vercel_ai_gateway():
"""Test that get_llm_provider correctly identifies vercel_ai_gateway"""
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
# Test with vercel_ai_gateway/provider/model-name format
model, provider, api_key, api_base = get_llm_provider("vercel_ai_gateway/openai/gpt-4o")
assert model == "openai/gpt-4o"
assert provider == "vercel_ai_gateway"
# Test with api_base containing vercel ai gateway endpoint
model, provider, api_key, api_base = get_llm_provider(
"gpt-4o", api_base="https://ai-gateway.vercel.sh/v1"
)
assert model == "gpt-4o"
assert provider == "vercel_ai_gateway"
assert api_base == "https://ai-gateway.vercel.sh/v1"
def test_vercel_ai_gateway_in_provider_lists():
"""Test that vercel_ai_gateway is registered in all necessary provider lists"""
assert "vercel_ai_gateway" in litellm.openai_compatible_providers
assert "vercel_ai_gateway" in litellm.provider_list
assert "https://ai-gateway.vercel.sh/v1" in litellm.openai_compatible_endpoints
@pytest.mark.asyncio
async def test_vercel_ai_gateway_completion_call():
"""Test completion call with vercel_ai_gateway provider (requires VERCEL_AI_GATEWAY_API_KEY)"""
# Skip if no API key is available
if not os.getenv("VERCEL_AI_GATEWAY_API_KEY"):
pytest.skip("VERCEL_AI_GATEWAY_API_KEY not set")
try:
response = await litellm.acompletion(
model="vercel_ai_gateway/openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello, this is a test"}],
max_tokens=20,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
# If the API key is invalid or there's a network issue, that's okay
# The important thing is that the provider was recognized
if "vercel_ai_gateway" not in str(e) and "provider" not in str(e).lower():
# Re-raise if it's not a provider-related error
raise
@pytest.mark.asyncio
async def test_vercel_ai_gateway_with_oidc_token():
"""Test completion call with vercel_ai_gateway provider using VERCEL_OIDC_TOKEN"""
# Skip if no OIDC token is available
if not os.getenv("VERCEL_OIDC_TOKEN"):
pytest.skip("VERCEL_OIDC_TOKEN not set")
try:
response = await litellm.acompletion(
model="vercel_ai_gateway/openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello, this is a test"}],
max_tokens=20,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
# If the OIDC token is invalid or there's a network issue, that's okay
# The important thing is that the provider was recognized
if "vercel_ai_gateway" not in str(e) and "provider" not in str(e).lower():
# Re-raise if it's not a provider-related error
raise
def test_vercel_ai_gateway_supported_params():
"""Test that vercel_ai_gateway returns the supported parameters"""
config = VercelAIGatewayConfig()
supported_params = config.get_supported_openai_params("vercel_ai_gateway/openai/gpt-3.5-turbo")
# vercel_ai_gateway should include all base OpenAI params plus extra_body
expected_base_params = [
"frequency_penalty",
"logit_bias",
"logprobs",
"top_logprobs",
"max_tokens",
"max_completion_tokens",
"modalities",
"prediction",
"n",
"presence_penalty",
"seed",
"stop",
"stream",
"stream_options",
"temperature",
"top_p",
"tools",
"tool_choice",
"function_call",
"functions",
"max_retries",
"extra_headers",
"parallel_tool_calls",
"audio",
"web_search_options",
"extra_body"
]
for param in expected_base_params:
assert param in supported_params, f"Expected parameter '{param}' not found in supported params"
assert "extra_body" in supported_params
def test_vercel_ai_gateway_sync_completion():
"""Test synchronous completion call"""
if not os.getenv("VERCEL_AI_GATEWAY_API_KEY"):
pytest.skip("VERCEL_AI_GATEWAY_API_KEY not set")
try:
response = completion(
model="vercel_ai_gateway/openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=20,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
if "vercel_ai_gateway" not in str(e) and "provider" not in str(e).lower():
raise
def test_vercel_ai_gateway_with_provider_options():
"""Test vercel_ai_gateway with providerOptions parameter"""
if not os.getenv("VERCEL_AI_GATEWAY_API_KEY"):
pytest.skip("VERCEL_AI_GATEWAY_API_KEY not set")
try:
response = completion(
model="vercel_ai_gateway/openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello"}],
providerOptions={"gateway": {"order": ["azure", "openai"]}},
max_tokens=20,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
if "vercel_ai_gateway" not in str(e) and "provider" not in str(e).lower():
raise
def test_vercel_ai_gateway_models_endpoint():
"""Test the get_models functionality"""
config = VercelAIGatewayConfig()
with mock.patch("litellm.module_level_client.get") as mock_get:
mock_response = mock.MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": [
{"id": "openai/gpt-4o"},
{"id": "openai/gpt-3.5-turbo"},
{"id": "anthropic/claude-4-sonnet"}
]
}
mock_get.return_value = mock_response
models = config.get_models()
assert models == ["openai/gpt-4o", "openai/gpt-3.5-turbo", "anthropic/claude-4-sonnet"]
mock_get.assert_called_once_with(url="https://ai-gateway.vercel.sh/v1/models")
def test_vercel_ai_gateway_models_endpoint_failure():
"""Test the get_models functionality with failure"""
config = VercelAIGatewayConfig()
with mock.patch("litellm.module_level_client.get") as mock_get:
mock_response = mock.MagicMock()
mock_response.status_code = 404
mock_response.text = "Not found"
mock_get.return_value = mock_response
with pytest.raises(Exception, match="Failed to get models: Not found"):
config.get_models()

112
tests/test_litellm/llms/vercel_ai_gateway/chat/test_vercel_ai_gateway_transformation.py Normal file
View file

@ -0,0 +1,112 @@
import os
import sys
from unittest.mock import patch
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.vercel_ai_gateway.chat.transformation import (
VercelAIGatewayConfig,
)
from litellm.llms.vercel_ai_gateway.common_utils import VercelAIGatewayException
def test_vercel_ai_gateway_extra_body_transformation():
"""Test that providerOptions is correctly moved to extra_body"""
transformed_request = VercelAIGatewayConfig().transform_request(
model="vercel_ai_gateway/openai/gpt-4o",
messages=[{"role": "user", "content": "Hello, world!"}],
optional_params={
"extra_body": {
"providerOptions": {
"gateway": {"order": ["azure", "openai"]}
}
}
},
litellm_params={},
headers={},
)
assert transformed_request["extra_body"]["providerOptions"]["gateway"]["order"] == ["azure", "openai"]
assert transformed_request["messages"] == [
{"role": "user", "content": "Hello, world!"}
]
def test_vercel_ai_gateway_provider_options_mapping():
"""Test that providerOptions from non_default_params is moved to extra_body"""
config = VercelAIGatewayConfig()
non_default_params = {
"providerOptions": {
"gateway": {"order": ["azure", "openai"]}
}
}
optional_params = {}
model = "vercel_ai_gateway/openai/gpt-4o"
result = config.map_openai_params(
non_default_params, optional_params, model, drop_params=False
)
assert result["extra_body"]["providerOptions"]["gateway"]["order"] == ["azure", "openai"]
assert "providerOptions" not in result
def test_vercel_ai_gateway_get_supported_openai_params():
"""Test that extra_body is included in supported params"""
config = VercelAIGatewayConfig()
supported_params = config.get_supported_openai_params("vercel_ai_gateway/openai/gpt-4o")
assert "extra_body" in supported_params
assert "temperature" in supported_params
assert "max_tokens" in supported_params
assert "stream" in supported_params
def test_vercel_ai_gateway_get_openai_compatible_provider_info():
"""Test provider info retrieval with environment variables"""
config = VercelAIGatewayConfig()
with patch.dict(
"os.environ",
{
"VERCEL_AI_GATEWAY_API_BASE": "https://env.vercel.sh/v1",
"VERCEL_AI_GATEWAY_API_KEY": "env_api_key",
},
):
api_base, api_key = config._get_openai_compatible_provider_info(None, None)
assert api_base == "https://env.vercel.sh/v1"
assert api_key == "env_api_key"
def test_vercel_ai_gateway_error_class():
"""Test error class creation"""
config = VercelAIGatewayConfig()
error_message = "Test error"
status_code = 400
headers = {"Content-Type": "application/json"}
error_class = config.get_error_class(error_message, status_code, headers)
assert isinstance(error_class, VercelAIGatewayException)
assert error_class.message == error_message
assert error_class.status_code == status_code
assert error_class.headers == headers
def test_vercel_ai_gateway_exception_inheritance():
"""Test that VercelAIGatewayException inherits from BaseLLMException"""
from litellm.llms.base_llm.chat.transformation import BaseLLMException
exception = VercelAIGatewayException(
message="test",
status_code=500,
headers={}
)
assert isinstance(exception, BaseLLMException)