""" Mock tests for vercel_ai_gateway provider """ import json from unittest.mock import MagicMock, patch import pytest import respx import litellm from litellm import completion from litellm.llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig from litellm.cost_calculator import cost_per_token import math @pytest.fixture def vercel_ai_gateway_response(): """Mock response from Vercel AI Gateway API""" return { "id": "chatcmpl-vercel-123", "object": "chat.completion", "created": 1677652288, "model": "openai/gpt-3.5-turbo", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "Hello! This is a test response from Vercel AI Gateway."}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 15, "total_tokens": 25}, } 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(respx_mock, vercel_ai_gateway_response, monkeypatch): """Test completion call with vercel_ai_gateway provider using mocked response""" monkeypatch.setenv("VERCEL_AI_GATEWAY_API_KEY", "test-api-key") litellm.disable_aiohttp_transport = True respx_mock.post("https://ai-gateway.vercel.sh/v1/chat/completions").respond(json=vercel_ai_gateway_response) 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 == "Hello! This is a test response from Vercel AI Gateway." assert response.model == "vercel_ai_gateway/openai/gpt-3.5-turbo" assert response.usage.total_tokens == 25 assert len(respx_mock.calls) == 1 request = respx_mock.calls[0].request assert request.method == "POST" assert "ai-gateway.vercel.sh" in str(request.url) assert "Authorization" in request.headers assert request.headers["Authorization"] == "Bearer test-api-key" @pytest.mark.asyncio async def test_vercel_ai_gateway_with_oidc_token(respx_mock, vercel_ai_gateway_response, monkeypatch): """Test completion call with vercel_ai_gateway provider using VERCEL_OIDC_TOKEN""" monkeypatch.setenv("VERCEL_OIDC_TOKEN", "test-oidc-token") litellm.disable_aiohttp_transport = True respx_mock.post("https://ai-gateway.vercel.sh/v1/chat/completions").respond(json=vercel_ai_gateway_response) 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 == "Hello! This is a test response from Vercel AI Gateway." assert response.model == "vercel_ai_gateway/openai/gpt-3.5-turbo" assert response.usage.total_tokens == 25 assert len(respx_mock.calls) == 1 request = respx_mock.calls[0].request assert "Authorization" in request.headers assert request.headers["Authorization"] == "Bearer test-oidc-token" 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(respx_mock, vercel_ai_gateway_response, monkeypatch): """Test synchronous completion call""" monkeypatch.setenv("VERCEL_AI_GATEWAY_API_KEY", "test-api-key") litellm.disable_aiohttp_transport = True respx_mock.post("https://ai-gateway.vercel.sh/v1/chat/completions").respond(json=vercel_ai_gateway_response) 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 == "Hello! This is a test response from Vercel AI Gateway." assert response.model == "vercel_ai_gateway/openai/gpt-3.5-turbo" assert response.usage.total_tokens == 25 def test_vercel_ai_gateway_with_provider_options(respx_mock, vercel_ai_gateway_response, monkeypatch): """Test vercel_ai_gateway with providerOptions parameter""" monkeypatch.setenv("VERCEL_AI_GATEWAY_API_KEY", "test-api-key") litellm.disable_aiohttp_transport = True respx_mock.post("https://ai-gateway.vercel.sh/v1/chat/completions").respond(json=vercel_ai_gateway_response) 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 == "Hello! This is a test response from Vercel AI Gateway." assert response.model == "vercel_ai_gateway/openai/gpt-3.5-turbo" assert response.usage.total_tokens == 25 assert len(respx_mock.calls) == 1 request = respx_mock.calls[0].request request_data = json.loads(request.content.decode("utf-8")) assert "providerOptions" in request_data assert request_data["providerOptions"]["gateway"]["order"] == ["azure", "openai"] def test_vercel_ai_gateway_models_endpoint(): """Test the get_models functionality""" config = VercelAIGatewayConfig() with patch("litellm.module_level_client.get") as mock_get: mock_response = 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-3-sonnet"}] } mock_get.return_value = mock_response models = config.get_models() assert models == ["openai/gpt-4o", "openai/gpt-3.5-turbo", "anthropic/claude-3-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 patch("litellm.module_level_client.get") as mock_get: mock_response = 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() def test_vercel_ai_gateway_glm46_cost_math(): """Test the cost math for glm-4.6""" with open("model_prices_and_context_window.json", "r") as f: litellm.model_cost = json.load(f) key = "vercel_ai_gateway/zai/glm-4.6" info = litellm.model_cost[key] prompt_cost, completion_cost = cost_per_token( model="vercel_ai_gateway/zai/glm-4.6", prompt_tokens=1000, completion_tokens=500, ) assert math.isclose(prompt_cost, 1000 * info["input_cost_per_token"], rel_tol=1e-12) assert math.isclose(completion_cost, 500 * info["output_cost_per_token"], rel_tol=1e-12)