Add grok-4 and groq kimi-k2 models (#2235)

Author: zachmayer

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -185,13 +185,22 @@
         'gpt-4o-mini-search-preview-2025-03-11',
         'gpt-4o-search-preview',
         'gpt-4o-search-preview-2025-03-11',
+        'grok:grok-4',
+        'grok:grok-4-0709',
+        'grok:grok-3',
+        'grok:grok-3-mini',
+        'grok:grok-3-fast',
+        'grok:grok-3-mini-fast',
+        'grok:grok-2-vision-1212',
+        'grok:grok-2-image-1212',
         'groq:distil-whisper-large-v3-en',
         'groq:gemma2-9b-it',
         'groq:llama-3.3-70b-versatile',
         'groq:llama-3.1-8b-instant',
         'groq:llama-guard-3-8b',
         'groq:llama3-70b-8192',
         'groq:llama3-8b-8192',
+        'groq:moonshotai/kimi-k2-instruct',
         'groq:whisper-large-v3',
         'groq:whisper-large-v3-turbo',
         'groq:playai-tts',

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -79,6 +79,7 @@
     'llama-3.2-3b-preview',
     'llama-3.2-11b-vision-preview',
     'llama-3.2-90b-vision-preview',
+    'moonshotai/kimi-k2-instruct',
 ]
 """Preview Groq models from <https://console.groq.com/docs/models#preview-models>."""
 

pydantic_ai_slim/pydantic_ai/profiles/moonshotai.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def moonshotai_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a MoonshotAI model."""
+    return None

pydantic_ai_slim/pydantic_ai/providers/grok.py

@@ -1,7 +1,7 @@
 from __future__ import annotations as _annotations
 
 import os
-from typing import overload
+from typing import Literal, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from openai import AsyncOpenAI
@@ -21,6 +21,18 @@
         'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
     ) from _import_error
 
+# https://docs.x.ai/docs/models
+GrokModelName = Literal[
+    'grok-4',
+    'grok-4-0709',
+    'grok-3',
+    'grok-3-mini',
+    'grok-3-fast',
+    'grok-3-mini-fast',
+    'grok-2-vision-1212',
+    'grok-2-image-1212',
+]
+
 
 class GrokProvider(Provider[AsyncOpenAI]):
     """Provider for Grok API."""

pydantic_ai_slim/pydantic_ai/providers/groq.py

@@ -12,6 +12,7 @@
 from pydantic_ai.profiles.google import google_model_profile
 from pydantic_ai.profiles.meta import meta_model_profile
 from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
 from pydantic_ai.profiles.qwen import qwen_model_profile
 from pydantic_ai.providers import Provider
 
@@ -47,6 +48,7 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
             'qwen': qwen_model_profile,
             'deepseek': deepseek_model_profile,
             'mistral': mistral_model_profile,
+            'moonshotai/': moonshotai_model_profile,
         }
 
         for prefix, profile_func in prefix_to_profile.items():

tests/models/cassettes/test_model_names/test_known_model_names.yaml

@@ -15,7 +15,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '545'
+      - '550'
       content-security-policy:
       - default-src 'none'; frame-ancestors 'none'
       content-type:
@@ -46,6 +46,7 @@ interactions:
       - text-to-text
     - model_id: claude-4-sonnet
       regions:
+      - eu
       - us
       type:
       - text-to-text

tests/models/test_model.py

@@ -72,6 +72,22 @@
         'github',
         'OpenAIModel',
     ),
+    (
+        'GROK_API_KEY',
+        'grok:grok-3',
+        'grok-3',
+        'grok',
+        'grok',
+        'OpenAIModel',
+    ),
+    (
+        'GROK_API_KEY',
+        'grok-4',  # Note that the provider and model name are both "grok", so the plain string grok with no prefix works because its also the provider name
+        'grok-4',
+        'grok',
+        'grok',
+        'OpenAIModel',
+    ),
 ]
 
 

tests/models/test_model_names.py

@@ -19,6 +19,7 @@
     from pydantic_ai.models.huggingface import HuggingFaceModelName
     from pydantic_ai.models.mistral import MistralModelName
     from pydantic_ai.models.openai import OpenAIModelName
+    from pydantic_ai.providers.grok import GrokModelName
 
 pytestmark = [
     pytest.mark.skipif(not imports_successful(), reason='some model package was not installed'),
@@ -48,6 +49,7 @@ def get_model_names(model_name_type: Any) -> Iterator[str]:
     google_names = [f'google-gla:{n}' for n in get_model_names(GeminiModelName)] + [
         f'google-vertex:{n}' for n in get_model_names(GeminiModelName)
     ]
+    grok_names = [f'grok:{n}' for n in get_model_names(GrokModelName)]
     groq_names = [f'groq:{n}' for n in get_model_names(GroqModelName)]
     mistral_names = [f'mistral:{n}' for n in get_model_names(MistralModelName)]
     openai_names = [f'openai:{n}' for n in get_model_names(OpenAIModelName)] + [
@@ -63,6 +65,7 @@ def get_model_names(model_name_type: Any) -> Iterator[str]:
         anthropic_names
         + cohere_names
         + google_names
+        + grok_names
         + groq_names
         + mistral_names
         + openai_names

tests/providers/test_groq.py

@@ -12,6 +12,7 @@
 from pydantic_ai.profiles.google import GoogleJsonSchemaTransformer, google_model_profile
 from pydantic_ai.profiles.meta import meta_model_profile
 from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
 from pydantic_ai.profiles.qwen import qwen_model_profile
 
 from ..conftest import TestEnv, try_import
@@ -74,6 +75,7 @@ def test_groq_provider_model_profile(mocker: MockerFixture):
     google_model_profile_mock = mocker.patch(f'{ns}.google_model_profile', wraps=google_model_profile)
     mistral_model_profile_mock = mocker.patch(f'{ns}.mistral_model_profile', wraps=mistral_model_profile)
     qwen_model_profile_mock = mocker.patch(f'{ns}.qwen_model_profile', wraps=qwen_model_profile)
+    moonshotai_model_profile_mock = mocker.patch(f'{ns}.moonshotai_model_profile', wraps=moonshotai_model_profile)
 
     meta_profile = provider.model_profile('meta-llama/Llama-Guard-4-12B')
     meta_model_profile_mock.assert_called_with('llama-guard-4-12b')
@@ -103,5 +105,10 @@ def test_groq_provider_model_profile(mocker: MockerFixture):
     assert qwen_profile is not None
     assert qwen_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
 
+    # MoonshotAI model should remove the "moonshotai/" prefix before passing to profile
+    moonshotai_profile = provider.model_profile('moonshotai/kimi-k2-instruct')
+    moonshotai_model_profile_mock.assert_called_with('kimi-k2-instruct')
+    assert moonshotai_profile is None
+
     unknown_profile = provider.model_profile('unknown-model')
     assert unknown_profile is None

tests/test_cli.py

@@ -144,6 +144,7 @@ def test_list_models(capfd: CaptureFixture[str]):
         'cohere',
         'deepseek',
         'heroku',
+        'grok',
         'huggingface',
     )
     models = {line.strip().split(' ')[0] for line in output[3:]}