Skip to content

Commit 985f186

Browse files
dmontagudmontagu
authored
Clean up known model names (#781)
1 parent d8c8959 commit 985f186

File tree

2 files changed

+168
-44
lines changed

2 files changed

+168
-44
lines changed

‎pydantic_ai_slim/pydantic_ai/models/__init__.py‎

Lines changed: 109 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -27,39 +27,86 @@
2727

2828

2929
KnownModelName = Literal[
30-
'openai:gpt-4o',
31-
'openai:gpt-4o-mini',
32-
'openai:gpt-4-turbo',
33-
'openai:gpt-4',
34-
'openai:o1-preview',
35-
'openai:o1-mini',
36-
'openai:o1',
37-
'openai:gpt-3.5-turbo',
38-
'groq:llama-3.3-70b-versatile',
39-
'groq:llama-3.3-70b-versatile',
40-
'groq:llama3-groq-70b-8192-tool-use-preview',
41-
'groq:llama3-groq-8b-8192-tool-use-preview',
42-
'groq:llama-3.1-70b-specdec',
30+
'anthropic:claude-3-5-haiku-latest',
31+
'anthropic:claude-3-5-sonnet-latest',
32+
'anthropic:claude-3-opus-latest',
33+
'claude-3-5-haiku-latest',
34+
'claude-3-5-sonnet-latest',
35+
'claude-3-opus-latest',
36+
'cohere:c4ai-aya-expanse-32b',
37+
'cohere:c4ai-aya-expanse-8b',
38+
'cohere:command',
39+
'cohere:command-light',
40+
'cohere:command-light-nightly',
41+
'cohere:command-nightly',
42+
'cohere:command-r',
43+
'cohere:command-r-03-2024',
44+
'cohere:command-r-08-2024',
45+
'cohere:command-r-plus',
46+
'cohere:command-r-plus-04-2024',
47+
'cohere:command-r-plus-08-2024',
48+
'cohere:command-r7b-12-2024',
49+
'google-gla:gemini-1.0-pro',
50+
'google-gla:gemini-1.5-flash',
51+
'google-gla:gemini-1.5-flash-8b',
52+
'google-gla:gemini-1.5-pro',
53+
'google-gla:gemini-2.0-flash-exp',
54+
'google-vertex:gemini-1.0-pro',
55+
'google-vertex:gemini-1.5-flash',
56+
'google-vertex:gemini-1.5-flash-8b',
57+
'google-vertex:gemini-1.5-pro',
58+
'google-vertex:gemini-2.0-flash-exp',
59+
'gpt-3.5-turbo',
60+
'gpt-3.5-turbo-0125',
61+
'gpt-3.5-turbo-0301',
62+
'gpt-3.5-turbo-0613',
63+
'gpt-3.5-turbo-1106',
64+
'gpt-3.5-turbo-16k',
65+
'gpt-3.5-turbo-16k-0613',
66+
'gpt-4',
67+
'gpt-4-0125-preview',
68+
'gpt-4-0314',
69+
'gpt-4-0613',
70+
'gpt-4-1106-preview',
71+
'gpt-4-32k',
72+
'gpt-4-32k-0314',
73+
'gpt-4-32k-0613',
74+
'gpt-4-turbo',
75+
'gpt-4-turbo-2024-04-09',
76+
'gpt-4-turbo-preview',
77+
'gpt-4-vision-preview',
78+
'gpt-4o',
79+
'gpt-4o-2024-05-13',
80+
'gpt-4o-2024-08-06',
81+
'gpt-4o-2024-11-20',
82+
'gpt-4o-audio-preview',
83+
'gpt-4o-audio-preview-2024-10-01',
84+
'gpt-4o-audio-preview-2024-12-17',
85+
'gpt-4o-mini',
86+
'gpt-4o-mini-2024-07-18',
87+
'gpt-4o-mini-audio-preview',
88+
'gpt-4o-mini-audio-preview-2024-12-17',
89+
'groq:gemma2-9b-it',
4390
'groq:llama-3.1-8b-instant',
91+
'groq:llama-3.2-11b-vision-preview',
4492
'groq:llama-3.2-1b-preview',
4593
'groq:llama-3.2-3b-preview',
46-
'groq:llama-3.2-11b-vision-preview',
4794
'groq:llama-3.2-90b-vision-preview',
95+
'groq:llama-3.3-70b-specdec',
96+
'groq:llama-3.3-70b-versatile',
4897
'groq:llama3-70b-8192',
4998
'groq:llama3-8b-8192',
5099
'groq:mixtral-8x7b-32768',
51-
'groq:gemma2-9b-it',
52-
'groq:gemma-7b-it',
53-
'google-gla:gemini-1.5-flash',
54-
'google-gla:gemini-1.5-pro',
55-
'google-gla:gemini-2.0-flash-exp',
56-
'google-vertex:gemini-1.5-flash',
57-
'google-vertex:gemini-1.5-pro',
58-
'google-vertex:gemini-2.0-flash-exp',
59-
'mistral:mistral-small-latest',
60-
'mistral:mistral-large-latest',
61100
'mistral:codestral-latest',
101+
'mistral:mistral-large-latest',
62102
'mistral:mistral-moderation-latest',
103+
'mistral:mistral-small-latest',
104+
'o1',
105+
'o1-2024-12-17',
106+
'o1-mini',
107+
'o1-mini-2024-09-12',
108+
'o1-preview',
109+
'o1-preview-2024-09-12',
63110
'ollama:codellama',
64111
'ollama:deepseek-r1',
65112
'ollama:gemma',
@@ -74,30 +121,48 @@
74121
'ollama:mixtral',
75122
'ollama:phi3',
76123
'ollama:phi4',
77-
'ollama:qwq',
78124
'ollama:qwen',
79125
'ollama:qwen2',
80126
'ollama:qwen2.5',
127+
'ollama:qwq',
81128
'ollama:starcoder2',
82-
'anthropic:claude-3-5-haiku-latest',
83-
'anthropic:claude-3-5-sonnet-latest',
84-
'anthropic:claude-3-opus-latest',
85-
'claude-3-5-haiku-latest',
86-
'claude-3-5-sonnet-latest',
87-
'claude-3-opus-latest',
88-
'cohere:c4ai-aya-expanse-32b',
89-
'cohere:c4ai-aya-expanse-8b',
90-
'cohere:command',
91-
'cohere:command-light',
92-
'cohere:command-light-nightly',
93-
'cohere:command-nightly',
94-
'cohere:command-r',
95-
'cohere:command-r-03-2024',
96-
'cohere:command-r-08-2024',
97-
'cohere:command-r-plus',
98-
'cohere:command-r-plus-04-2024',
99-
'cohere:command-r-plus-08-2024',
100-
'cohere:command-r7b-12-2024',
129+
'openai:chatgpt-4o-latest',
130+
'openai:gpt-3.5-turbo',
131+
'openai:gpt-3.5-turbo-0125',
132+
'openai:gpt-3.5-turbo-0301',
133+
'openai:gpt-3.5-turbo-0613',
134+
'openai:gpt-3.5-turbo-1106',
135+
'openai:gpt-3.5-turbo-16k',
136+
'openai:gpt-3.5-turbo-16k-0613',
137+
'openai:gpt-4',
138+
'openai:gpt-4-0125-preview',
139+
'openai:gpt-4-0314',
140+
'openai:gpt-4-0613',
141+
'openai:gpt-4-1106-preview',
142+
'openai:gpt-4-32k',
143+
'openai:gpt-4-32k-0314',
144+
'openai:gpt-4-32k-0613',
145+
'openai:gpt-4-turbo',
146+
'openai:gpt-4-turbo-2024-04-09',
147+
'openai:gpt-4-turbo-preview',
148+
'openai:gpt-4-vision-preview',
149+
'openai:gpt-4o',
150+
'openai:gpt-4o-2024-05-13',
151+
'openai:gpt-4o-2024-08-06',
152+
'openai:gpt-4o-2024-11-20',
153+
'openai:gpt-4o-audio-preview',
154+
'openai:gpt-4o-audio-preview-2024-10-01',
155+
'openai:gpt-4o-audio-preview-2024-12-17',
156+
'openai:gpt-4o-mini',
157+
'openai:gpt-4o-mini-2024-07-18',
158+
'openai:gpt-4o-mini-audio-preview',
159+
'openai:gpt-4o-mini-audio-preview-2024-12-17',
160+
'openai:o1',
161+
'openai:o1-2024-12-17',
162+
'openai:o1-mini',
163+
'openai:o1-mini-2024-09-12',
164+
'openai:o1-preview',
165+
'openai:o1-preview-2024-09-12',
101166
'test',
102167
]
103168
"""Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].

‎tests/models/test_model_names.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
from collections.abc import Iterator
2+
from typing import Any, get_args
3+
4+
import pytest
5+
6+
from pydantic_ai.models import KnownModelName
7+
8+
from ..conftest import try_import
9+
10+
with try_import() as imports_successful:
11+
from pydantic_ai.models.anthropic import AnthropicModelName
12+
from pydantic_ai.models.cohere import CohereModelName
13+
from pydantic_ai.models.gemini import GeminiModelName
14+
from pydantic_ai.models.groq import GroqModelName
15+
from pydantic_ai.models.mistral import MistralModelName
16+
from pydantic_ai.models.ollama import OllamaModelName
17+
from pydantic_ai.models.openai import OpenAIModelName
18+
19+
pytestmark = [
20+
pytest.mark.skipif(not imports_successful(), reason='some model package was not installed'),
21+
]
22+
23+
24+
def test_known_model_names():
25+
def get_model_names(model_name_type: Any) -> Iterator[str]:
26+
for arg in get_args(model_name_type):
27+
if isinstance(arg, str):
28+
yield arg
29+
else:
30+
yield from get_model_names(arg)
31+
32+
anthropic_names = [f'anthropic:{n}' for n in get_model_names(AnthropicModelName)] + [
33+
n for n in get_model_names(AnthropicModelName) if n.startswith('claude')
34+
]
35+
cohere_names = [f'cohere:{n}' for n in get_model_names(CohereModelName)]
36+
google_names = [f'google-gla:{n}' for n in get_model_names(GeminiModelName)] + [
37+
f'google-vertex:{n}' for n in get_model_names(GeminiModelName)
38+
]
39+
groq_names = [f'groq:{n}' for n in get_model_names(GroqModelName)]
40+
mistral_names = [f'mistral:{n}' for n in get_model_names(MistralModelName)]
41+
ollama_names = [f'ollama:{n}' for n in get_model_names(OllamaModelName)]
42+
openai_names = [f'openai:{n}' for n in get_model_names(OpenAIModelName)] + [
43+
n for n in get_model_names(OpenAIModelName) if n.startswith('o1') or n.startswith('gpt')
44+
]
45+
extra_names = ['test']
46+
47+
generated_names = sorted(
48+
anthropic_names
49+
+ cohere_names
50+
+ google_names
51+
+ groq_names
52+
+ mistral_names
53+
+ ollama_names
54+
+ openai_names
55+
+ extra_names
56+
)
57+
known_model_names = sorted(get_args(KnownModelName))
58+
59+
assert generated_names == known_model_names

0 commit comments

Comments
 (0)