Repository: pydantic/pydantic-ai
File: tests/models/test_model_names.py (Python)

Code before:
        'before_record_response': partial(modify_response, filter_headers=['cache-control', 'connection']),
    }


def test_known_model_names():  # pragma: lax no cover
    # Coverage seems to be misbehaving..?
    def get_model_names(model_name_type: Any) -> Iterator[str]:
        for arg in get_args(model_name_type):
            if isinstance(arg, str):
                yield arg
            else:
                yield from get_model_names(arg)

    anthropic_names = [f'anthropic:{n}' for n in get_model_names(AnthropicModelName)]
    cohere_names = [f'cohere:{n}' for n in get_model_names(CohereModelName)]
    google_names = [f'google-gla:{n}' for n in get_model_names(GoogleModelName)] + [
        f'google-vertex:{n}' for n in get_model_names(GoogleModelName)
    ]
    grok_names = [f'grok:{n}' for n in get_model_names(GrokModelName)]
    groq_names = [f'groq:{n}' for n in get_model_names(GroqModelName)]
    moonshotai_names = [f'moonshotai:{n}' for n in get_model_names(MoonshotAIModelName)]
    mistral_names = [f'mistral:{n}' for n in get_model_names(MistralModelName)]
    openai_names = [f'openai:{n}' for n in get_model_names(OpenAIModelName)]
    bedrock_names = [f'bedrock:{n}' for n in get_model_names(BedrockModelName)]
    deepseek_names = ['deepseek:deepseek-chat', 'deepseek:deepseek-reasoner']
    gateway_names = (
        [f'gateway/anthropic:{n}' for n in get_model_names(AnthropicModelName)]
        + [f'gateway/bedrock:{n}' for n in get_model_names(BedrockModelName)]
        + [f'gateway/google-vertex:{n}' for n in get_model_names(GoogleModelName)]
        + [f'gateway/groq:{n}' for n in get_model_names(GroqModelName)]
        + [f'gateway/openai:{n}' for n in get_model_names(OpenAIModelName)]
    )
    huggingface_names = [f'huggingface:{n}' for n in get_model_names(HuggingFaceModelName)]
    heroku_names = get_heroku_model_names()
    cerebras_names = get_cerebras_model_names()
    extra_names = ['test']

    generated_names = sorted(
        anthropic_names
        + cohere_names
        + google_names
        + grok_names
        + groq_names
        + mistral_names
        + moonshotai_names
        + openai_names
        + bedrock_names
        + deepseek_names
        + gateway_names
        + huggingface_names
        + heroku_names

Review comment (suggestion):
Could we use the `pydantic_ai.providers.gateway.UpstreamProvider` type (possibly needs splitting up into 2 so that we don't include the APi types here), and then use those names as keys into a dict of things like `{'openai': OpenAIModelName}`? That way when we update it in the gateway file, this list will update too

Code after:
        'record_mode': 'rewrite',
        'filter_headers': ['accept-encoding'],
        'before_record_response': partial(modify_response, filter_headers=['cache-control', 'connection']),
    }


_PROVIDER_TO_MODEL_NAMES = {
    'anthropic': AnthropicModelName,
    'bedrock': BedrockModelName,
    'cohere': CohereModelName,
    'deepseek': DeepSeekModelName,
    'google-gla': GoogleModelName,
    'google-vertex': GoogleModelName,
    'grok': GrokModelName,
    'groq': GroqModelName,
    'huggingface': HuggingFaceModelName,
    'mistral': MistralModelName,
    'moonshotai': MoonshotAIModelName,
    'openai': OpenAIModelName,
}


def test_known_model_names():  # pragma: lax no cover
    # Coverage seems to be misbehaving..?
    def get_model_names(model_name_type: Any) -> Iterator[str]:
        for arg in get_args(model_name_type):
            if isinstance(arg, str):
                yield arg
            else:
                yield from get_model_names(arg)

    all_generated_names = [
        f'{provider}:{n}'
        for provider, model_names in _PROVIDER_TO_MODEL_NAMES.items()
        for n in get_model_names(model_names)
    ]

    cerebras_names = get_cerebras_model_names()
    heroku_names = get_heroku_model_names()
    gateway_names = [
        f'gateway/{provider}:{model_name}'
        for provider in GatewayModelProvider.__args__
        for model_name in get_model_names(_PROVIDER_TO_MODEL_NAMES[provider])
    ]

    extra_names = ['test']

    generated_names = sorted(all_generated_names + gateway_names + heroku_names + cerebras_names + extra_names)

    known_model_names = sorted(get_args(KnownModelName.__value__))
    assert generated_names == known_model_names