Add MoonshotAI provider with Kimi-K2 model support (pydantic#2211)

Co-authored-by: Douwe Maan <[email protected]>

Author: 2 people

docs/models/openai.md

@@ -401,6 +401,24 @@ agent = Agent(model)
 ...
 ```
 
+### MoonshotAI
+
+Create an API key in the [Moonshot Console](https://platform.moonshot.ai/console).
+With that key you can instantiate the [`MoonshotAIProvider`][pydantic_ai.providers.moonshotai.MoonshotAIProvider]:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.moonshotai import MoonshotAIProvider
+
+model = OpenAIModel(
+    'kimi-k2-0711-preview',
+    provider=MoonshotAIProvider(api_key='your-moonshot-api-key'),
+)
+agent = Agent(model)
+...
+```
+
 ### GitHub Models
 
 To use [GitHub Models](https://docs.github.com/en/github-models), you'll need a GitHub personal access token with the `models: read` permission.

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -233,6 +233,15 @@
         'mistral:mistral-large-latest',
         'mistral:mistral-moderation-latest',
         'mistral:mistral-small-latest',
+        'moonshotai:moonshot-v1-8k',
+        'moonshotai:moonshot-v1-32k',
+        'moonshotai:moonshot-v1-128k',
+        'moonshotai:moonshot-v1-8k-vision-preview',
+        'moonshotai:moonshot-v1-32k-vision-preview',
+        'moonshotai:moonshot-v1-128k-vision-preview',
+        'moonshotai:kimi-latest',
+        'moonshotai:kimi-thinking-preview',
+        'moonshotai:kimi-k2-0711-preview',
         'o1',
         'o1-2024-12-17',
         'o1-mini',
@@ -617,6 +626,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'openrouter',
         'vercel',
         'grok',
+        'moonshotai',
         'fireworks',
         'together',
         'heroku',

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -195,6 +195,7 @@ def __init__(
             'deepseek',
             'azure',
             'openrouter',
+            'moonshotai',
             'vercel',
             'grok',
             'fireworks',
@@ -299,7 +300,10 @@ async def _completions_create(
         tools = self._get_tools(model_request_parameters)
         if not tools:
             tool_choice: Literal['none', 'required', 'auto'] | None = None
-        elif not model_request_parameters.allow_text_output:
+        elif (
+            not model_request_parameters.allow_text_output
+            and OpenAIModelProfile.from_profile(self.profile).openai_supports_tool_choice_required
+        ):
             tool_choice = 'required'
         else:
             tool_choice = 'auto'

pydantic_ai_slim/pydantic_ai/profiles/openai.py

@@ -21,6 +21,14 @@ class OpenAIModelProfile(ModelProfile):
     openai_supports_sampling_settings: bool = True
     """Turn off to don't send sampling settings like `temperature` and `top_p` to models that don't support them, like OpenAI's o-series reasoning models."""
 
+    # Some OpenAI-compatible providers (e.g. MoonshotAI) currently do **not** accept
+    # `tool_choice="required"`.  This flag lets the calling model know whether it's
+    # safe to pass that value along.  Default is `True` to preserve existing
+    # behaviour for OpenAI itself and most providers.
+    openai_supports_tool_choice_required: bool = True
+    """Whether the provider accepts the value ``tool_choice='required'`` in the
+    request payload."""
+
 
 def openai_model_profile(model_name: str) -> ModelProfile:
     """Get the model profile for an OpenAI model."""

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -103,6 +103,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .grok import GrokProvider
 
         return GrokProvider
+    elif provider == 'moonshotai':
+        from .moonshotai import MoonshotAIProvider
+
+        return MoonshotAIProvider
     elif provider == 'fireworks':
         from .fireworks import FireworksProvider
 

pydantic_ai_slim/pydantic_ai/providers/moonshotai.py

@@ -0,0 +1,97 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import Literal, overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
+from pydantic_ai.profiles.openai import (
+    OpenAIJsonSchemaTransformer,
+    OpenAIModelProfile,
+)
+from pydantic_ai.providers import Provider
+
+MoonshotAIModelName = Literal[
+    'moonshot-v1-8k',
+    'moonshot-v1-32k',
+    'moonshot-v1-128k',
+    'moonshot-v1-8k-vision-preview',
+    'moonshot-v1-32k-vision-preview',
+    'moonshot-v1-128k-vision-preview',
+    'kimi-latest',
+    'kimi-thinking-preview',
+    'kimi-k2-0711-preview',
+]
+
+
+class MoonshotAIProvider(Provider[AsyncOpenAI]):
+    """Provider for MoonshotAI platform (Kimi models)."""
+
+    @property
+    def name(self) -> str:
+        return 'moonshotai'
+
+    @property
+    def base_url(self) -> str:
+        # OpenAI-compatible endpoint, see MoonshotAI docs
+        return 'https://api.moonshot.ai/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        profile = moonshotai_model_profile(model_name)
+
+        # As the MoonshotAI API is OpenAI-compatible, let's assume we also need OpenAIJsonSchemaTransformer,
+        # unless json_schema_transformer is set explicitly.
+        # Also, MoonshotAI does not support strict tool definitions
+        # https://platform.moonshot.ai/docs/guide/migrating-from-openai-to-kimi#about-tool_choice
+        # "Please note that the current version of Kimi API does not support the tool_choice=required parameter."
+        return OpenAIModelProfile(
+            json_schema_transformer=OpenAIJsonSchemaTransformer,
+            openai_supports_tool_choice_required=False,
+            supports_json_object_output=True,
+        ).update(profile)
+
+    # ---------------------------------------------------------------------
+    # Construction helpers
+    # ---------------------------------------------------------------------
+    @overload
+    def __init__(self) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        api_key = api_key or os.getenv('MOONSHOTAI_API_KEY')
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `MOONSHOTAI_API_KEY` environment variable or pass it via '
+                '`MoonshotAIProvider(api_key=...)` to use the MoonshotAI provider.'
+            )
+
+        if openai_client is not None:
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            http_client = cached_async_http_client(provider='moonshotai')
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)

tests/models/test_model.py

@@ -72,6 +72,14 @@
         'github',
         'OpenAIModel',
     ),
+    (
+        'MOONSHOTAI_API_KEY',
+        'moonshotai:kimi-k2-0711-preview',
+        'kimi-k2-0711-preview',
+        'moonshotai',
+        'moonshotai',
+        'OpenAIModel',
+    ),
     (
         'GROK_API_KEY',
         'grok:grok-3',

tests/models/test_model_names.py

@@ -20,6 +20,7 @@
     from pydantic_ai.models.mistral import MistralModelName
     from pydantic_ai.models.openai import OpenAIModelName
     from pydantic_ai.providers.grok import GrokModelName
+    from pydantic_ai.providers.moonshotai import MoonshotAIModelName
 
 pytestmark = [
     pytest.mark.skipif(not imports_successful(), reason='some model package was not installed'),
@@ -51,6 +52,7 @@ def get_model_names(model_name_type: Any) -> Iterator[str]:
     ]
     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)] + [
         n for n in get_model_names(OpenAIModelName) if n.startswith('o1') or n.startswith('gpt') or n.startswith('o3')
@@ -68,6 +70,7 @@ def get_model_names(model_name_type: Any) -> Iterator[str]:
         + grok_names
         + groq_names
         + mistral_names
+        + moonshotai_names
         + openai_names
         + bedrock_names
         + deepseek_names

tests/models/test_openai.py

@@ -36,6 +36,7 @@
     ToolReturnPart,
     UserPromptPart,
 )
+from pydantic_ai.models import ModelRequestParameters
 from pydantic_ai.models.gemini import GeminiModel
 from pydantic_ai.output import NativeOutput, PromptedOutput, TextOutput, ToolOutput
 from pydantic_ai.profiles import ModelProfile
@@ -2631,3 +2632,22 @@ async def test_process_response_no_created_timestamp(allow_model_requests: None)
     response_message = messages[1]
     assert isinstance(response_message, ModelResponse)
     assert response_message.timestamp == IsNow(tz=timezone.utc)
+
+
+@pytest.mark.anyio()
+async def test_tool_choice_fallback(allow_model_requests: None) -> None:
+    profile = OpenAIModelProfile(openai_supports_tool_choice_required=False).update(openai_model_profile('stub'))
+
+    mock_client = MockOpenAI.create_mock(completion_message(ChatCompletionMessage(content='ok', role='assistant')))
+    model = OpenAIModel('stub', provider=OpenAIProvider(openai_client=mock_client), profile=profile)
+
+    params = ModelRequestParameters(function_tools=[ToolDefinition(name='x')], allow_text_output=False)
+
+    await model._completions_create(  # pyright: ignore[reportPrivateUsage]
+        messages=[],
+        stream=False,
+        model_settings={},
+        model_request_parameters=params,
+    )
+
+    assert get_mock_chat_completion_kwargs(mock_client)[0]['tool_choice'] == 'auto'

tests/providers/test_moonshotai.py

@@ -0,0 +1,70 @@
+import re
+
+import httpx
+import pytest
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
+
+from ..conftest import TestEnv, try_import
+
+with try_import() as imports_successful:
+    import openai
+
+    from pydantic_ai.models.openai import OpenAIModel
+    from pydantic_ai.providers.moonshotai import MoonshotAIProvider
+
+pytestmark = pytest.mark.skipif(not imports_successful(), reason='openai not installed')
+
+
+def test_moonshotai_provider():
+    """Test basic MoonshotAI provider initialization."""
+    provider = MoonshotAIProvider(api_key='api-key')
+    assert provider.name == 'moonshotai'
+    assert provider.base_url == 'https://api.moonshot.ai/v1'
+    assert isinstance(provider.client, openai.AsyncOpenAI)
+    assert provider.client.api_key == 'api-key'
+
+
+def test_moonshotai_provider_need_api_key(env: TestEnv) -> None:
+    """Test that MoonshotAI provider requires an API key."""
+    env.remove('MOONSHOTAI_API_KEY')
+    with pytest.raises(
+        UserError,
+        match=re.escape(
+            'Set the `MOONSHOTAI_API_KEY` environment variable or pass it via `MoonshotAIProvider(api_key=...)`'
+            ' to use the MoonshotAI provider.'
+        ),
+    ):
+        MoonshotAIProvider()
+
+
+def test_moonshotai_provider_pass_http_client() -> None:
+    """Test passing a custom HTTP client to MoonshotAI provider."""
+    http_client = httpx.AsyncClient()
+    provider = MoonshotAIProvider(http_client=http_client, api_key='api-key')
+    assert provider.client._client == http_client  # type: ignore[reportPrivateUsage]
+
+
+def test_moonshotai_pass_openai_client() -> None:
+    """Test passing a custom OpenAI client to MoonshotAI provider."""
+    openai_client = openai.AsyncOpenAI(api_key='api-key')
+    provider = MoonshotAIProvider(openai_client=openai_client)
+    assert provider.client == openai_client
+
+
+def test_moonshotai_provider_with_cached_http_client() -> None:
+    """Test MoonshotAI provider using cached HTTP client (covers line 76)."""
+    # This should use the else branch with cached_async_http_client
+    provider = MoonshotAIProvider(api_key='api-key')
+    assert isinstance(provider.client, openai.AsyncOpenAI)
+    assert provider.client.api_key == 'api-key'
+
+
+def test_moonshotai_model_profile():
+    provider = MoonshotAIProvider(api_key='api-key')
+    model = OpenAIModel('kimi-k2-0711-preview', provider=provider)
+    assert isinstance(model.profile, OpenAIModelProfile)
+    assert model.profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+    assert model.profile.openai_supports_tool_choice_required is False
+    assert model.profile.supports_json_object_output is True