diff --git a/rlm/clients/__init__.py b/rlm/clients/__init__.py index 6a2f7f22..09d3008b 100644 --- a/rlm/clients/__init__.py +++ b/rlm/clients/__init__.py @@ -36,6 +36,11 @@ def get_client( backend_kwargs.setdefault("base_url", "https://openrouter.ai/api/v1") return OpenAIClient(**backend_kwargs) + elif backend == "vercel": + from rlm.clients.openai import OpenAIClient + + backend_kwargs.setdefault("base_url", "https://ai-gateway.vercel.sh/v1") + return OpenAIClient(**backend_kwargs) elif backend == "litellm": from rlm.clients.litellm import LiteLLMClient @@ -54,5 +59,5 @@ def get_client( return AzureOpenAIClient(**backend_kwargs) else: raise ValueError( - f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic', 'azure_openai', 'gemini']" + f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic', 'azure_openai', 'gemini', 'vercel']" ) diff --git a/rlm/clients/openai.py b/rlm/clients/openai.py index 48e8ed37..9b46a42e 100644 --- a/rlm/clients/openai.py +++ b/rlm/clients/openai.py @@ -13,6 +13,7 @@ # Load API keys from environment variables DEFAULT_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") DEFAULT_OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY") +DEFAULT_VERCEL_API_KEY = os.getenv("AI_GATEWAY_API_KEY") DEFAULT_PRIME_INTELLECT_BASE_URL = "https://api.pinference.ai/api/v1/" @@ -35,6 +36,8 @@ def __init__( api_key = DEFAULT_OPENAI_API_KEY elif base_url == "https://openrouter.ai/api/v1": api_key = DEFAULT_OPENROUTER_API_KEY + elif base_url == "https://ai-gateway.vercel.sh/v1": + api_key = DEFAULT_VERCEL_API_KEY # For vLLM, set base_url to local vLLM server address. self.client = openai.OpenAI(api_key=api_key, base_url=base_url) diff --git a/rlm/core/types.py b/rlm/core/types.py index c4ce16bb..7e913835 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -3,7 +3,15 @@ from typing import Any, Literal ClientBackend = Literal[ - "openai", "portkey", "openrouter", "vllm", "litellm", "anthropic", "azure_openai", "gemini" + "openai", + "portkey", + "openrouter", + "vercel", + "vllm", + "litellm", + "anthropic", + "azure_openai", + "gemini", ] EnvironmentType = Literal["local", "prime", "modal"]