/tests/providers/test_gateway.py

import os
import re
from typing import Any, Literal
from unittest.mock import patch
from urllib.parse import urlparse

import httpx
import pytest

from pydantic_ai import Agent, UserError

from .._inline_snapshot import raises, snapshot
from ..conftest import TestEnv, try_import

with try_import() as imports_successful:
    from pydantic_ai.models.anthropic import AnthropicModel
    from pydantic_ai.models.bedrock import BedrockConverseModel
    from pydantic_ai.models.google import GoogleModel
    from pydantic_ai.models.groq import GroqModel
    from pydantic_ai.models.openai import OpenAIChatModel, OpenAIResponsesModel
    from pydantic_ai.providers import Provider
    from pydantic_ai.providers.anthropic import AnthropicProvider
    from pydantic_ai.providers.bedrock import BedrockProvider
    from pydantic_ai.providers.gateway import gateway_provider
    from pydantic_ai.providers.google_cloud import GoogleCloudProvider
    from pydantic_ai.providers.groq import GroqProvider
    from pydantic_ai.providers.openai import OpenAIProvider


if not imports_successful():
    pytest.skip('Providers not installed', allow_module_level=True)  # pragma: lax no cover

pytestmark = [pytest.mark.anyio, pytest.mark.vcr]

# Any URL works here — these tests exercise the explicit `PYDANTIC_AI_GATEWAY_BASE_URL` override path.
GATEWAY_BASE_URL = 'https://gateway.pydantic.dev/proxy'


@pytest.mark.parametrize(
    'provider_name, provider_cls, route',
    [
        # PAIG exposes a single canonical `openai` route; Chat vs Responses is selected by the
        # OpenAI SDK's sub-path (/chat/completions vs /responses), not by the URL prefix.
        ('openai', OpenAIProvider, 'openai'),
        ('openai-chat', OpenAIProvider, 'openai'),
        ('openai-responses', OpenAIProvider, 'openai'),
    ],
)
def test_init_with_base_url(
    provider_name: Literal['openai', 'openai-chat', 'openai-responses'], provider_cls: type[Provider[Any]], route: str
):
    provider = gateway_provider(provider_name, base_url='https://example.com/', api_key='foobar')
    assert isinstance(provider, provider_cls)
    assert provider.base_url == f'https://example.com/{route}/'
    assert provider.client.api_key == 'foobar'


def test_init_gateway_without_api_key_raises_error(env: TestEnv):
    env.remove('PYDANTIC_AI_GATEWAY_API_KEY')
    with pytest.raises(
        UserError,
        match=re.escape(
            'Set the `PYDANTIC_AI_GATEWAY_API_KEY` environment variable or pass it via `gateway_provider(..., api_key=...)` to use the Pydantic AI Gateway provider.'
        ),
    ):
        gateway_provider('openai')


async def test_init_with_http_client():
    async with httpx.AsyncClient() as http_client:
        provider = gateway_provider('openai', http_client=http_client, api_key='foobar', base_url=GATEWAY_BASE_URL)
        assert provider.client._client == http_client  # pyright: ignore[reportPrivateUsage]


async def test_init_with_http_client_preserves_existing_event_hooks():
    # Unit (not VCR): this checks local HTTPX hook merging by inspecting and invoking event hooks directly;
    # cassette playback would not exercise hook ordering or preservation.
    async def existing_request_hook(request: httpx.Request) -> None:
        request.headers['X-Existing-Request-Hook'] = 'kept'

    async def existing_response_hook(response: httpx.Response) -> None:
        response.headers['X-Existing-Response-Hook'] = 'kept'

    async with httpx.AsyncClient(
        event_hooks={'request': [existing_request_hook], 'response': [existing_response_hook]}
    ) as http_client:
        provider = gateway_provider('openai', http_client=http_client, api_key='foobar', base_url=GATEWAY_BASE_URL)
        assert provider.client._client == http_client  # pyright: ignore[reportPrivateUsage]
        assert existing_request_hook in http_client.event_hooks['request']
        assert existing_response_hook in http_client.event_hooks['response']

        request = httpx.Request('GET', provider.base_url)
        for hook in http_client.event_hooks['request']:
            await hook(request)

        assert request.headers['X-Existing-Request-Hook'] == 'kept'
        assert request.headers['Authorization'] == 'Bearer foobar'

        response = httpx.Response(200, request=request)
        for hook in http_client.event_hooks['response']:
            await hook(response)

        assert response.headers['X-Existing-Response-Hook'] == 'kept'


async def test_init_with_http_client_replaces_existing_gateway_hook():
    # Unit (not VCR): this checks local HTTPX hook replacement by inspecting and invoking event hooks directly;
    # cassette playback would not exercise Gateway hook deduplication.
    async def existing_request_hook(request: httpx.Request) -> None:
        request.headers['X-Existing-Request-Hook'] = 'kept'

    async with httpx.AsyncClient(event_hooks={'request': [existing_request_hook]}) as http_client:
        first_provider = gateway_provider('openai', http_client=http_client, api_key='first', base_url=GATEWAY_BASE_URL)
        second_provider = gateway_provider(
            'openai', http_client=http_client, api_key='second', base_url=GATEWAY_BASE_URL
        )

        assert first_provider.client._client == http_client  # pyright: ignore[reportPrivateUsage]
        assert second_provider.client._client == http_client  # pyright: ignore[reportPrivateUsage]
        assert http_client.event_hooks['request'][0] == existing_request_hook
        assert len(http_client.event_hooks['request']) == 2

        request = httpx.Request('GET', second_provider.base_url)
        for hook in http_client.event_hooks['request']:
            await hook(request)

        assert request.headers['X-Existing-Request-Hook'] == 'kept'
        assert request.headers['Authorization'] == 'Bearer second'


@pytest.fixture
def gateway_api_key():
    return os.getenv('PYDANTIC_AI_GATEWAY_API_KEY', 'test-api-key')


@pytest.fixture(scope='module')
def vcr_config():
    return {
        'ignore_localhost': False,
        # Note: additional header filtering is done inside the serializer
        'filter_headers': ['authorization', 'x-api-key'],
        'decode_compressed_response': True,
    }


@patch.dict(
    os.environ, {'PYDANTIC_AI_GATEWAY_API_KEY': 'test-api-key', 'PYDANTIC_AI_GATEWAY_BASE_URL': GATEWAY_BASE_URL}
)
@pytest.mark.parametrize(
    'provider_name, provider_cls, route',
    [
        ('openai', OpenAIProvider, 'openai'),
        ('openai-chat', OpenAIProvider, 'openai'),
        ('openai-responses', OpenAIProvider, 'openai'),
        ('groq', GroqProvider, 'groq'),
        ('google', GoogleCloudProvider, 'google-vertex'),
        ('google-cloud', GoogleCloudProvider, 'google-vertex'),
        ('anthropic', AnthropicProvider, 'anthropic'),
        ('bedrock', BedrockProvider, 'bedrock'),
    ],
)
def test_gateway_provider(provider_name: str, provider_cls: type[Provider[Any]], route: str):
    provider = gateway_provider(provider_name)
    assert isinstance(provider, provider_cls)

    # Some providers add a trailing slash, others don't
    assert provider.base_url in (f'{GATEWAY_BASE_URL}/{route}/', f'{GATEWAY_BASE_URL}/{route}')


@patch.dict(
    os.environ, {'PYDANTIC_AI_GATEWAY_API_KEY': 'test-api-key', 'PYDANTIC_AI_GATEWAY_BASE_URL': GATEWAY_BASE_URL}
)
@pytest.mark.parametrize('removed_alias', ['foo', 'google-vertex', 'gemini'])
def test_gateway_provider_unknown(removed_alias: str):
    # `google-vertex` and `gemini` were removed in v2 alongside their bare-prefix counterparts —
    # `gateway/google-vertex:` and `gateway/gemini:` raise the same `UserError` as any other unknown alias.
    with pytest.raises(UserError, match=f'Unknown upstream provider: {removed_alias}'):
        gateway_provider(removed_alias)


async def test_gateway_provider_with_openai(allow_model_requests: None, gateway_api_key: str):
    provider = gateway_provider('openai-chat', api_key=gateway_api_key, base_url='http://localhost:8787')
    model = OpenAIChatModel('gpt-5', provider=provider)
    agent = Agent(model)

    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('Paris.')


async def test_gateway_provider_with_openai_responses(allow_model_requests: None, gateway_api_key: str):
    provider = gateway_provider('openai-responses', api_key=gateway_api_key, base_url='http://localhost:8787')
    model = OpenAIResponsesModel('gpt-5', provider=provider)
    agent = Agent(model)

    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('Paris.')


async def test_gateway_provider_with_groq(allow_model_requests: None, gateway_api_key: str):
    provider = gateway_provider('groq', api_key=gateway_api_key, base_url='http://localhost:8787')
    model = GroqModel('llama-3.3-70b-versatile', provider=provider)
    agent = Agent(model)

    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('The capital of France is Paris.')


async def test_gateway_provider_with_google_cloud(allow_model_requests: None, gateway_api_key: str):
    provider = gateway_provider('google-cloud', api_key=gateway_api_key, base_url='http://localhost:8787')
    model = GoogleModel('gemini-2.5-flash', provider=provider)
    agent = Agent(model)

    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('The capital of France is **Paris**.')


async def test_gateway_provider_with_anthropic(allow_model_requests: None, gateway_api_key: str):
    provider = gateway_provider('anthropic', api_key=gateway_api_key, base_url='http://localhost:8787')
    model = AnthropicModel('claude-sonnet-4-5', provider=provider)
    agent = Agent(model)

    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('The capital of France is Paris.')


async def test_gateway_provider_with_bedrock(allow_model_requests: None, gateway_api_key: str):
    provider = gateway_provider('bedrock', api_key=gateway_api_key, base_url='http://localhost:8787')
    model = BedrockConverseModel('amazon.nova-micro-v1:0', provider=provider)
    agent = Agent(model)

    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot(
        'The capital of France is Paris. Paris is not only the capital city but also the most populous city in France, and it is a major center for culture, commerce, fashion, and international diplomacy. The city is known for its historical landmarks, such as the Eiffel Tower, the Louvre Museum, Notre-Dame Cathedral, and the Champs-Élysées, among many other attractions.'
    )


@patch.dict(
    os.environ, {'PYDANTIC_AI_GATEWAY_API_KEY': 'test-api-key', 'PYDANTIC_AI_GATEWAY_BASE_URL': GATEWAY_BASE_URL}
)
async def test_model_provider_argument():
    model = OpenAIChatModel('gpt-5', provider='gateway')
    assert urlparse(model._provider.base_url).hostname == urlparse(GATEWAY_BASE_URL).hostname  # pyright: ignore[reportPrivateUsage]

    model = OpenAIResponsesModel('gpt-5', provider='gateway')
    assert urlparse(model._provider.base_url).hostname == urlparse(GATEWAY_BASE_URL).hostname  # pyright: ignore[reportPrivateUsage]

    model = GroqModel('llama-3.3-70b-versatile', provider='gateway')
    assert urlparse(model._provider.base_url).hostname == urlparse(GATEWAY_BASE_URL).hostname  # pyright: ignore[reportPrivateUsage]

    model = GoogleModel('gemini-1.5-flash', provider='gateway')
    assert urlparse(model._provider.base_url).hostname == urlparse(GATEWAY_BASE_URL).hostname  # pyright: ignore[reportPrivateUsage]

    model = AnthropicModel('claude-sonnet-4-5', provider='gateway')
    assert urlparse(model._provider.base_url).hostname == urlparse(GATEWAY_BASE_URL).hostname  # pyright: ignore[reportPrivateUsage]

    model = BedrockConverseModel('amazon.nova-micro-v1:0', provider='gateway')
    assert urlparse(model._provider.base_url).hostname == urlparse(GATEWAY_BASE_URL).hostname  # pyright: ignore[reportPrivateUsage]


async def test_gateway_provider_routing_group(gateway_api_key: str):
    provider = gateway_provider('openai', route='potato', api_key=gateway_api_key, base_url=GATEWAY_BASE_URL)
    assert provider.client.base_url.path.endswith('/potato/')


@pytest.mark.parametrize(
    'api_key, expected_base_url',
    [
        pytest.param('pylf_v1_us_abc123', 'gateway-us.pydantic.dev', id='us-region'),
        pytest.param('pylf_v1_eu_abc123', 'gateway-eu.pydantic.dev', id='eu-region'),
        pytest.param('pylf_v1_stagingus_abc123', 'gateway.pydantic.info', id='staging'),
        pytest.param('pylf_v1_ap_abc123', 'gateway-ap.pydantic.dev', id='any-region'),
    ],
)
def test_infer_base_url(api_key: str, expected_base_url: str):
    provider = gateway_provider('openai', api_key=api_key)
    assert urlparse(provider.base_url).netloc == expected_base_url


def test_infer_base_url_no_region():
    """An API key that doesn't encode a region used to fall back to a shared Gateway URL; that URL
    is dead, so it now raises instead of silently routing to a dead host."""
    with raises(
        snapshot(
            'UserError: Could not infer the Pydantic AI Gateway base URL: the API key does not encode a region. '
            'Generate a new key from the Pydantic AI Gateway, or set the `PYDANTIC_AI_GATEWAY_BASE_URL` '
            'environment variable explicitly.'
        )
    ):
        gateway_provider('openai', api_key='not-a-pylf-token')