feat(vercel_ai_gateway): add embeddings support

feat(vercel_ai_gateway): add embeddings support

docs/my-website/docs/providers/vercel_ai_gateway.md

@@ -11,7 +11,7 @@ import TabItem from '@theme/TabItem';
 | 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` |
+| Supported Operations | `/chat/completions`, `/embeddings`, `/models` |
 
 <br />
 <br />
@@ -73,7 +73,7 @@ messages = [{"content": "Hello, how are you?", "role": "user"}]
 
 # Vercel AI Gateway call with streaming
 response = completion(
-    model="vercel_ai_gateway/openai/gpt-4o", 
+    model="vercel_ai_gateway/openai/gpt-4o",
     messages=messages,
     stream=True
 )
@@ -82,6 +82,33 @@ for chunk in response:
     print(chunk)
 ```
 
+### Embeddings
+
+```python showLineNumbers title="Vercel AI Gateway Embeddings"
+import os
+from litellm import embedding
+
+os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key"
+
+# Vercel AI Gateway embedding call
+response = embedding(
+    model="vercel_ai_gateway/openai/text-embedding-3-small",
+    input="Hello world"
+)
+
+print(response.data[0]["embedding"][:5])  # Print first 5 dimensions
+```
+
+You can also specify the `dimensions` parameter:
+
+```python showLineNumbers title="Vercel AI Gateway Embeddings with Dimensions"
+response = embedding(
+    model="vercel_ai_gateway/openai/text-embedding-3-small",
+    input=["Hello world", "Goodbye world"],
+    dimensions=768
+)
+```
+
 ## Usage - LiteLLM Proxy
 
 Add the following to your LiteLLM Proxy configuration file:
@@ -97,6 +124,11 @@ model_list:
     litellm_params:
       model: vercel_ai_gateway/anthropic/claude-4-sonnet
       api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY
+
+  - model_name: text-embedding-3-small-gateway
+    litellm_params:
+      model: vercel_ai_gateway/openai/text-embedding-3-small
+      api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY
 ```
 
 Start your LiteLLM Proxy server:

litellm/llms/vercel_ai_gateway/embedding/transformation.py

@@ -0,0 +1,176 @@
+"""
+Vercel AI Gateway Embedding API Configuration.
+
+This module provides the configuration for Vercel AI Gateway's Embedding API.
+Vercel AI Gateway is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint.
+
+Docs: https://vercel.com/docs/ai-gateway/openai-compat/embeddings
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+import httpx
+
+from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.openai import AllEmbeddingInputValues
+from litellm.types.utils import EmbeddingResponse
+from litellm.utils import convert_to_model_response_object
+
+from ..common_utils import VercelAIGatewayException
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+
+    LiteLLMLoggingObj = _LiteLLMLoggingObj
+else:
+    LiteLLMLoggingObj = Any
+
+
+class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig):
+    """
+    Configuration for Vercel AI Gateway's Embedding API.
+
+    Reference: https://vercel.com/docs/ai-gateway/openai-compat/embeddings
+    """
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: list,
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        """
+        Validate environment and set up headers for Vercel AI Gateway API.
+
+        Vercel AI Gateway requires:
+        - Authorization header with Bearer token (API key or OIDC token)
+        """
+        vercel_headers = {
+            "Content-Type": "application/json",
+        }
+
+        # Add Authorization header if api_key is provided
+        if api_key:
+            vercel_headers["Authorization"] = f"Bearer {api_key}"
+
+        # Merge with existing headers (user's extra_headers take priority)
+        merged_headers = {**vercel_headers, **headers}
+
+        return merged_headers
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        """
+        Get the complete URL for Vercel AI Gateway Embedding API endpoint.
+        """
+        if api_base:
+            api_base = api_base.rstrip("/")
+        else:
+            api_base = (
+                get_secret_str("VERCEL_AI_GATEWAY_API_BASE")
+                or "https://ai-gateway.vercel.sh/v1"
+            )
+
+        return f"{api_base}/embeddings"
+
+    def transform_embedding_request(
+        self,
+        model: str,
+        input: AllEmbeddingInputValues,
+        optional_params: dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform embedding request to Vercel AI Gateway format (OpenAI-compatible).
+        """
+        # Ensure input is a list
+        if isinstance(input, str):
+            input = [input]
+
+        # Strip 'vercel_ai_gateway/' prefix if present
+        if model.startswith("vercel_ai_gateway/"):
+            model = model.replace("vercel_ai_gateway/", "", 1)
+
+        return {
+            "model": model,
+            "input": input,
+            **optional_params,
+        }
+
+    def transform_embedding_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: EmbeddingResponse,
+        logging_obj: LiteLLMLoggingObj,
+        api_key: Optional[str],
+        request_data: dict,
+        optional_params: dict,
+        litellm_params: dict,
+    ) -> EmbeddingResponse:
+        """
+        Transform embedding response from Vercel AI Gateway format (OpenAI-compatible).
+        """
+        logging_obj.post_call(original_response=raw_response.text)
+
+        # Vercel AI Gateway returns standard OpenAI-compatible embedding response
+        response_json = raw_response.json()
+
+        return convert_to_model_response_object(
+            response_object=response_json,
+            model_response_object=model_response,
+            response_type="embedding",
+        )
+
+    def get_supported_openai_params(self, model: str) -> list:
+        """
+        Get list of supported OpenAI parameters for Vercel AI Gateway embeddings.
+
+        Vercel AI Gateway supports the standard OpenAI embeddings parameters
+        and auto-maps 'dimensions' to each provider's expected field.
+        """
+        return [
+            "timeout",
+            "dimensions",
+            "encoding_format",
+            "user",
+        ]
+
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        """
+        Map OpenAI parameters to Vercel AI Gateway format.
+        """
+        for param, value in non_default_params.items():
+            if param in self.get_supported_openai_params(model):
+                optional_params[param] = value
+        return optional_params
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Any
+    ) -> Any:
+        """
+        Get the error class for Vercel AI Gateway errors.
+        """
+        return VercelAIGatewayException(
+            message=error_message,
+            status_code=status_code,
+            headers=headers,
+        )

litellm/main.py

@@ -4866,6 +4866,36 @@ def embedding(  # noqa: PLR0915
 
             headers = openrouter_headers
 
+            response = base_llm_http_handler.embedding(
+                model=model,
+                input=input,
+                custom_llm_provider=custom_llm_provider,
+                api_base=api_base,
+                api_key=api_key,
+                logging_obj=logging,
+                timeout=timeout,
+                model_response=EmbeddingResponse(),
+                optional_params=optional_params,
+                client=client,
+                aembedding=aembedding,
+                litellm_params=litellm_params_dict,
+                headers=headers,
+            )
+        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_str("VERCEL_AI_GATEWAY_API_KEY")
+                or get_secret_str("VERCEL_OIDC_TOKEN")
+            )
+
             response = base_llm_http_handler.embedding(
                 model=model,
                 input=input,

litellm/utils.py

@@ -8046,6 +8046,12 @@ class ProviderConfigManager:
             )
 
             return OpenrouterEmbeddingConfig()
+        elif litellm.LlmProviders.VERCEL_AI_GATEWAY == provider:
+            from litellm.llms.vercel_ai_gateway.embedding.transformation import (
+                VercelAIGatewayEmbeddingConfig,
+            )
+
+            return VercelAIGatewayEmbeddingConfig()
         elif litellm.LlmProviders.GIGACHAT == provider:
             return litellm.GigaChatEmbeddingConfig()
         elif litellm.LlmProviders.SAGEMAKER == provider:

tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py

@@ -0,0 +1,218 @@
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import httpx
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.llms.vercel_ai_gateway.embedding.transformation import (
+    VercelAIGatewayEmbeddingConfig,
+)
+from litellm.llms.vercel_ai_gateway.common_utils import VercelAIGatewayException
+from litellm.types.utils import EmbeddingResponse
+
+
+def test_vercel_ai_gateway_embedding_get_complete_url():
+    """Test URL generation for embeddings endpoint"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    # Test with default API base
+    url = config.get_complete_url(
+        api_base=None,
+        api_key=None,
+        model="openai/text-embedding-3-small",
+        optional_params={},
+        litellm_params={},
+    )
+    assert url == "https://ai-gateway.vercel.sh/v1/embeddings"
+
+    # Test with custom API base
+    url = config.get_complete_url(
+        api_base="https://custom.vercel.sh/v1",
+        api_key=None,
+        model="openai/text-embedding-3-small",
+        optional_params={},
+        litellm_params={},
+    )
+    assert url == "https://custom.vercel.sh/v1/embeddings"
+
+    # Test with trailing slash
+    url = config.get_complete_url(
+        api_base="https://custom.vercel.sh/v1/",
+        api_key=None,
+        model="openai/text-embedding-3-small",
+        optional_params={},
+        litellm_params={},
+    )
+    assert url == "https://custom.vercel.sh/v1/embeddings"
+
+
+def test_vercel_ai_gateway_embedding_transform_request():
+    """Test request transformation for embeddings"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    # Test with string input
+    request = config.transform_embedding_request(
+        model="openai/text-embedding-3-small",
+        input="Hello world",
+        optional_params={},
+        headers={},
+    )
+    assert request["model"] == "openai/text-embedding-3-small"
+    assert request["input"] == ["Hello world"]
+
+    # Test with list input
+    request = config.transform_embedding_request(
+        model="openai/text-embedding-3-small",
+        input=["Hello", "World"],
+        optional_params={},
+        headers={},
+    )
+    assert request["model"] == "openai/text-embedding-3-small"
+    assert request["input"] == ["Hello", "World"]
+
+    # Test stripping vercel_ai_gateway/ prefix
+    request = config.transform_embedding_request(
+        model="vercel_ai_gateway/openai/text-embedding-3-small",
+        input="Hello",
+        optional_params={},
+        headers={},
+    )
+    assert request["model"] == "openai/text-embedding-3-small"
+
+
+def test_vercel_ai_gateway_embedding_transform_request_with_dimensions():
+    """Test request transformation with dimensions parameter"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    request = config.transform_embedding_request(
+        model="openai/text-embedding-3-small",
+        input="Hello world",
+        optional_params={"dimensions": 768},
+        headers={},
+    )
+    assert request["model"] == "openai/text-embedding-3-small"
+    assert request["input"] == ["Hello world"]
+    assert request["dimensions"] == 768
+
+
+def test_vercel_ai_gateway_embedding_validate_environment():
+    """Test header validation and setup"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    headers = config.validate_environment(
+        headers={},
+        model="openai/text-embedding-3-small",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key="test_key",
+    )
+    assert headers["Content-Type"] == "application/json"
+    assert headers["Authorization"] == "Bearer test_key"
+
+    # Test with existing headers (should merge)
+    headers = config.validate_environment(
+        headers={"X-Custom": "value"},
+        model="openai/text-embedding-3-small",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key="test_key",
+    )
+    assert headers["X-Custom"] == "value"
+    assert headers["Authorization"] == "Bearer test_key"
+
+
+def test_vercel_ai_gateway_embedding_get_supported_params():
+    """Test supported OpenAI parameters"""
+    config = VercelAIGatewayEmbeddingConfig()
+    supported = config.get_supported_openai_params("openai/text-embedding-3-small")
+
+    assert "dimensions" in supported
+    assert "encoding_format" in supported
+    assert "timeout" in supported
+    assert "user" in supported
+
+
+def test_vercel_ai_gateway_embedding_map_openai_params():
+    """Test OpenAI parameter mapping"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    optional_params = config.map_openai_params(
+        non_default_params={"dimensions": 768, "encoding_format": "float"},
+        optional_params={},
+        model="openai/text-embedding-3-small",
+        drop_params=False,
+    )
+    assert optional_params["dimensions"] == 768
+    assert optional_params["encoding_format"] == "float"
+
+
+def test_vercel_ai_gateway_embedding_error_class():
+    """Test error class creation"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    error = config.get_error_class(
+        error_message="Test error",
+        status_code=400,
+        headers={"Content-Type": "application/json"},
+    )
+
+    assert isinstance(error, VercelAIGatewayException)
+    assert error.message == "Test error"
+    assert error.status_code == 400
+
+
+def test_vercel_ai_gateway_embedding_transform_response():
+    """Test response transformation"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    mock_response = MagicMock(spec=httpx.Response)
+    mock_response.text = '{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}],"model":"openai/text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}'
+    mock_response.json.return_value = {
+        "object": "list",
+        "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}],
+        "model": "openai/text-embedding-3-small",
+        "usage": {"prompt_tokens": 2, "total_tokens": 2},
+    }
+
+    mock_logging = MagicMock()
+
+    response = config.transform_embedding_response(
+        model="openai/text-embedding-3-small",
+        raw_response=mock_response,
+        model_response=EmbeddingResponse(),
+        logging_obj=mock_logging,
+        api_key="test_key",
+        request_data={},
+        optional_params={},
+        litellm_params={},
+    )
+
+    assert response is not None
+    mock_logging.post_call.assert_called_once()
+
+
+def test_vercel_ai_gateway_embedding_env_vars():
+    """Test environment variable handling"""
+    config = VercelAIGatewayEmbeddingConfig()
+
+    with patch.dict(
+        os.environ,
+        {
+            "VERCEL_AI_GATEWAY_API_BASE": "https://env.vercel.sh/v1",
+        },
+    ):
+        url = config.get_complete_url(
+            api_base=None,
+            api_key=None,
+            model="openai/text-embedding-3-small",
+            optional_params={},
+            litellm_params={},
+        )
+        assert url == "https://env.vercel.sh/v1/embeddings"