feat: add MoonshotAI provider with Kimi-K2 model support

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -286,6 +286,8 @@
         'openai:o3-mini-2025-01-31',
         'openai:o4-mini',
         'openai:o4-mini-2025-04-16',
+        'openai:computer-use-preview-2025-03-11',
+        'moonshotai:kimi-k2-0711-preview',
         'test',
     ],
 )
@@ -588,6 +590,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'azure',
         'openrouter',
         'grok',
+        'moonshotai',
         'fireworks',
         'together',
         'heroku',

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -190,7 +190,16 @@ def __init__(
         model_name: OpenAIModelName,
         *,
         provider: Literal[
-            'openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'heroku', 'github'
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'grok',
+            'moonshotai',
+            'fireworks',
+            'together',
+            'heroku',
+            'github',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -598,7 +607,18 @@ def __init__(
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together']
+        provider: Literal[
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'grok',
+            'moonshotai',
+            'fireworks',
+            'together',
+            'heroku',
+            'github',
+        ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,

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/__init__.py

@@ -99,6 +99,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,84 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import 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
+
+
+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_strict_tool_definition=False,
+        ).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('MOONSHOT_API_KEY')
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `MOONSHOT_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',
     ),
+    (
+        'MOONSHOT_API_KEY',
+        'moonshotai:kimi-k2-0711-preview',
+        'kimi-k2-0711-preview',
+        'moonshotai',
+        'openai',
+        'OpenAIModel',
+    ),
 ]
 
 

tests/models/test_model_names.py

@@ -49,6 +49,7 @@ def get_model_names(model_name_type: Any) -> Iterator[str]:
         f'google-vertex:{n}' for n in get_model_names(GeminiModelName)
     ]
     groq_names = [f'groq:{n}' for n in get_model_names(GroqModelName)]
+    moonshotai_names = ['moonshotai:kimi-k2-0711-preview']
     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')
@@ -64,6 +65,7 @@ def get_model_names(model_name_type: Any) -> Iterator[str]:
         + cohere_names
         + google_names
         + groq_names
+        + moonshotai_names
         + mistral_names
         + openai_names
         + bedrock_names

tests/providers/test_moonshotai.py

@@ -0,0 +1,69 @@
+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('MOONSHOT_API_KEY')
+    with pytest.raises(
+        UserError,
+        match=re.escape(
+            'Set the `MOONSHOT_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_strict_tool_definition is False

tests/providers/test_provider_names.py

@@ -26,6 +26,7 @@
     from pydantic_ai.providers.groq import GroqProvider
     from pydantic_ai.providers.heroku import HerokuProvider
     from pydantic_ai.providers.mistral import MistralProvider
+    from pydantic_ai.providers.moonshotai import MoonshotAIProvider
     from pydantic_ai.providers.openai import OpenAIProvider
     from pydantic_ai.providers.openrouter import OpenRouterProvider
     from pydantic_ai.providers.together import TogetherProvider
@@ -42,6 +43,7 @@
         ('groq', GroqProvider, 'GROQ_API_KEY'),
         ('mistral', MistralProvider, 'MISTRAL_API_KEY'),
         ('grok', GrokProvider, 'GROK_API_KEY'),
+        ('moonshotai', MoonshotAIProvider, 'MOONSHOT_API_KEY'),
         ('fireworks', FireworksProvider, 'FIREWORKS_API_KEY'),
         ('together', TogetherProvider, 'TOGETHER_API_KEY'),
         ('heroku', HerokuProvider, 'HEROKU_INFERENCE_KEY'),

tests/test_cli.py

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