pydantic-ai-slim 0.4.6__py3-none-any.whl → 0.4.8__py3-none-any.whl

pydantic_ai/messages.py

@@ -85,7 +85,7 @@ class SystemPromptPart:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-@dataclass(repr=False)
+@dataclass(init=False, repr=False)
 class FileUrl(ABC):
     """Abstract base class for any URL-based file."""
 
@@ -106,11 +106,29 @@ class FileUrl(ABC):
     - `GoogleModel`: `VideoUrl.vendor_metadata` is used as `video_metadata`: https://ai.google.dev/gemini-api/docs/video-understanding#customize-video-processing
     """
 
-    @property
+    _media_type: str | None = field(init=False, repr=False)
+
+    def __init__(
+        self,
+        url: str,
+        force_download: bool = False,
+        vendor_metadata: dict[str, Any] | None = None,
+        media_type: str | None = None,
+    ) -> None:
+        self.url = url
+        self.vendor_metadata = vendor_metadata
+        self.force_download = force_download
+        self._media_type = media_type
+
     @abstractmethod
-    def media_type(self) -> str:
+    def _infer_media_type(self) -> str:
         """Return the media type of the file, based on the url."""
 
+    @property
+    def media_type(self) -> str:
+        """Return the media type of the file, based on the url or the provided `_media_type`."""
+        return self._media_type or self._infer_media_type()
+
     @property
     @abstractmethod
     def format(self) -> str:
@@ -119,7 +137,7 @@ class FileUrl(ABC):
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-@dataclass(repr=False)
+@dataclass(init=False, repr=False)
 class VideoUrl(FileUrl):
     """A URL to a video."""
 
@@ -129,8 +147,18 @@ class VideoUrl(FileUrl):
     kind: Literal['video-url'] = 'video-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def media_type(self) -> VideoMediaType:
+    def __init__(
+        self,
+        url: str,
+        force_download: bool = False,
+        vendor_metadata: dict[str, Any] | None = None,
+        media_type: str | None = None,
+        kind: Literal['video-url'] = 'video-url',
+    ) -> None:
+        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        self.kind = kind
+
+    def _infer_media_type(self) -> VideoMediaType:
         """Return the media type of the video, based on the url."""
         if self.url.endswith('.mkv'):
             return 'video/x-matroska'
@@ -170,7 +198,7 @@ class VideoUrl(FileUrl):
         return _video_format_lookup[self.media_type]
 
 
-@dataclass(repr=False)
+@dataclass(init=False, repr=False)
 class AudioUrl(FileUrl):
     """A URL to an audio file."""
 
@@ -180,8 +208,18 @@ class AudioUrl(FileUrl):
     kind: Literal['audio-url'] = 'audio-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def media_type(self) -> AudioMediaType:
+    def __init__(
+        self,
+        url: str,
+        force_download: bool = False,
+        vendor_metadata: dict[str, Any] | None = None,
+        media_type: str | None = None,
+        kind: Literal['audio-url'] = 'audio-url',
+    ) -> None:
+        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        self.kind = kind
+
+    def _infer_media_type(self) -> AudioMediaType:
         """Return the media type of the audio file, based on the url.
 
         References:
@@ -208,7 +246,7 @@ class AudioUrl(FileUrl):
         return _audio_format_lookup[self.media_type]
 
 
-@dataclass(repr=False)
+@dataclass(init=False, repr=False)
 class ImageUrl(FileUrl):
     """A URL to an image."""
 
@@ -218,8 +256,18 @@ class ImageUrl(FileUrl):
     kind: Literal['image-url'] = 'image-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def media_type(self) -> ImageMediaType:
+    def __init__(
+        self,
+        url: str,
+        force_download: bool = False,
+        vendor_metadata: dict[str, Any] | None = None,
+        media_type: str | None = None,
+        kind: Literal['image-url'] = 'image-url',
+    ) -> None:
+        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        self.kind = kind
+
+    def _infer_media_type(self) -> ImageMediaType:
         """Return the media type of the image, based on the url."""
         if self.url.endswith(('.jpg', '.jpeg')):
             return 'image/jpeg'
@@ -241,7 +289,7 @@ class ImageUrl(FileUrl):
         return _image_format_lookup[self.media_type]
 
 
-@dataclass(repr=False)
+@dataclass(init=False, repr=False)
 class DocumentUrl(FileUrl):
     """The URL of the document."""
 
@@ -251,8 +299,18 @@ class DocumentUrl(FileUrl):
     kind: Literal['document-url'] = 'document-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def media_type(self) -> str:
+    def __init__(
+        self,
+        url: str,
+        force_download: bool = False,
+        vendor_metadata: dict[str, Any] | None = None,
+        media_type: str | None = None,
+        kind: Literal['document-url'] = 'document-url',
+    ) -> None:
+        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        self.kind = kind
+
+    def _infer_media_type(self) -> str:
         """Return the media type of the document, based on the url."""
         type_, _ = guess_type(self.url)
         if type_ is None:
@@ -632,7 +690,7 @@ class ThinkingPart:
 
     def has_content(self) -> bool:
         """Return `True` if the thinking content is non-empty."""
-        return bool(self.content)  # pragma: no cover
+        return bool(self.content)
 
     __repr__ = _utils.dataclasses_no_defaults_repr
 
@@ -757,11 +815,16 @@ class ModelResponse:
                         },
                     }
                 )
-            elif isinstance(part, TextPart):
-                if body.get('content'):
-                    body = new_event_body()
-                if settings.include_content:
-                    body['content'] = part.content
+            elif isinstance(part, (TextPart, ThinkingPart)):
+                kind = part.part_kind
+                body.setdefault('content', []).append(
+                    {'kind': kind, **({'text': part.content} if settings.include_content else {})}
+                )
+
+        if content := body.get('content'):
+            text_content = content[0].get('text')
+            if content == [{'kind': 'text', 'text': text_content}]:
+                body['content'] = text_content
 
         return result
 

pydantic_ai/models/__init__.py

@@ -233,6 +233,15 @@ KnownModelName = TypeAliasType(
         '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',
@@ -615,7 +624,9 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'deepseek',
         'azure',
         'openrouter',
+        'vercel',
         'grok',
+        'moonshotai',
         'fireworks',
         'together',
         'heroku',

pydantic_ai/models/anthropic.py

@@ -470,7 +470,7 @@ class AnthropicStreamedResponse(StreamedResponse):
     _response: AsyncIterable[BetaRawMessageStreamEvent]
     _timestamp: datetime
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         current_block: BetaContentBlock | None = None
 
         async for event in self._response:
@@ -479,7 +479,11 @@ class AnthropicStreamedResponse(StreamedResponse):
             if isinstance(event, BetaRawContentBlockStartEvent):
                 current_block = event.content_block
                 if isinstance(current_block, BetaTextBlock) and current_block.text:
-                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=current_block.text)
+                    maybe_event = self._parts_manager.handle_text_delta(
+                        vendor_part_id='content', content=current_block.text
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
                 elif isinstance(current_block, BetaThinkingBlock):
                     yield self._parts_manager.handle_thinking_delta(
                         vendor_part_id='thinking',
@@ -498,7 +502,11 @@ class AnthropicStreamedResponse(StreamedResponse):
 
             elif isinstance(event, BetaRawContentBlockDeltaEvent):
                 if isinstance(event.delta, BetaTextDelta):
-                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=event.delta.text)
+                    maybe_event = self._parts_manager.handle_text_delta(
+                        vendor_part_id='content', content=event.delta.text
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
                 elif isinstance(event.delta, BetaThinkingDelta):
                     yield self._parts_manager.handle_thinking_delta(
                         vendor_part_id='thinking', content=event.delta.thinking

pydantic_ai/models/bedrock.py

@@ -572,7 +572,7 @@ class BedrockStreamedResponse(StreamedResponse):
     _event_stream: EventStream[ConverseStreamOutputTypeDef]
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         """Return an async iterator of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
 
         This method should be implemented by subclasses to translate the vendor-specific stream of events into
@@ -618,7 +618,9 @@ class BedrockStreamedResponse(StreamedResponse):
                             UserWarning,
                         )
                 if 'text' in delta:
-                    yield self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
+                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
+                    if maybe_event is not None:
+                        yield maybe_event
                 if 'toolUse' in delta:
                     tool_use = delta['toolUse']
                     maybe_event = self._parts_manager.handle_tool_call_delta(

pydantic_ai/models/cohere.py

@@ -38,15 +38,15 @@ try:
         AssistantChatMessageV2,
         AsyncClientV2,
         ChatMessageV2,
-        ChatResponse,
         SystemChatMessageV2,
-        TextAssistantMessageContentItem,
+        TextAssistantMessageV2ContentItem,
         ToolCallV2,
         ToolCallV2Function,
         ToolChatMessageV2,
         ToolV2,
         ToolV2Function,
         UserChatMessageV2,
+        V2ChatResponse,
     )
     from cohere.core.api_error import ApiError
     from cohere.v2.client import OMIT
@@ -164,7 +164,7 @@ class CohereModel(Model):
         messages: list[ModelMessage],
         model_settings: CohereModelSettings,
         model_request_parameters: ModelRequestParameters,
-    ) -> ChatResponse:
+    ) -> V2ChatResponse:
         tools = self._get_tools(model_request_parameters)
         cohere_messages = self._map_messages(messages)
         try:
@@ -185,7 +185,7 @@ class CohereModel(Model):
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
             raise  # pragma: no cover
 
-    def _process_response(self, response: ChatResponse) -> ModelResponse:
+    def _process_response(self, response: V2ChatResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         parts: list[ModelResponsePart] = []
         if response.message.content is not None and len(response.message.content) > 0:
@@ -227,7 +227,7 @@ class CohereModel(Model):
                         assert_never(item)
                 message_param = AssistantChatMessageV2(role='assistant')
                 if texts:
-                    message_param.content = [TextAssistantMessageContentItem(text='\n\n'.join(texts))]
+                    message_param.content = [TextAssistantMessageV2ContentItem(text='\n\n'.join(texts))]
                 if tool_calls:
                     message_param.tool_calls = tool_calls
                 cohere_messages.append(message_param)
@@ -294,7 +294,7 @@ class CohereModel(Model):
                 assert_never(part)
 
 
-def _map_usage(response: ChatResponse) -> usage.Usage:
+def _map_usage(response: V2ChatResponse) -> usage.Usage:
     u = response.usage
     if u is None:
         return usage.Usage()

pydantic_ai/models/function.py

@@ -264,7 +264,9 @@ class FunctionStreamedResponse(StreamedResponse):
             if isinstance(item, str):
                 response_tokens = _estimate_string_tokens(item)
                 self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
-                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=item)
+                maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=item)
+                if maybe_event is not None:  # pragma: no branch
+                    yield maybe_event
             elif isinstance(item, dict) and item:
                 for dtc_index, delta in item.items():
                     if isinstance(delta, DeltaThinkingPart):
@@ -286,7 +288,7 @@ class FunctionStreamedResponse(StreamedResponse):
                             args=delta.json_args,
                             tool_call_id=delta.tool_call_id,
                         )
-                        if maybe_event is not None:
+                        if maybe_event is not None:  # pragma: no branch
                             yield maybe_event
                     else:
                         assert_never(delta)

pydantic_ai/models/gemini.py

@@ -438,7 +438,11 @@ class GeminiStreamedResponse(StreamedResponse):
                 if 'text' in gemini_part:
                     # Using vendor_part_id=None means we can produce multiple text parts if their deltas are sprinkled
                     # amongst the tool call deltas
-                    yield self._parts_manager.handle_text_delta(vendor_part_id=None, content=gemini_part['text'])
+                    maybe_event = self._parts_manager.handle_text_delta(
+                        vendor_part_id=None, content=gemini_part['text']
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
 
                 elif 'function_call' in gemini_part:
                     # Here, we assume all function_call parts are complete and don't have deltas.

pydantic_ai/models/google.py

@@ -411,7 +411,12 @@ class GoogleModel(Model):
                         file_data_dict['video_metadata'] = item.vendor_metadata
                     content.append(file_data_dict)  # type: ignore
                 elif isinstance(item, FileUrl):
-                    if self.system == 'google-gla' or item.force_download:
+                    if item.force_download or (
+                        # google-gla does not support passing file urls directly, except for youtube videos
+                        # (see above) and files uploaded to the file API (which cannot be downloaded anyway)
+                        self.system == 'google-gla'
+                        and not item.url.startswith(r'https://generativelanguage.googleapis.com/v1beta/files')
+                    ):
                         downloaded_item = await download_item(item, data_format='base64')
                         inline_data = {'data': downloaded_item['data'], 'mime_type': downloaded_item['data_type']}
                         content.append({'inline_data': inline_data})  # type: ignore
@@ -453,7 +458,9 @@ class GeminiStreamedResponse(StreamedResponse):
                     if part.thought:
                         yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
                     else:
-                        yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=part.text)
+                        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=part.text)
+                        if maybe_event is not None:  # pragma: no branch
+                            yield maybe_event
                 elif part.function_call:
                     maybe_event = self._parts_manager.handle_tool_call_delta(
                         vendor_part_id=uuid4(),

pydantic_ai/models/groq.py

@@ -415,7 +415,11 @@ class GroqStreamedResponse(StreamedResponse):
             # Handle the text part of the response
             content = choice.delta.content
             if content is not None:
-                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content)
+                maybe_event = self._parts_manager.handle_text_delta(
+                    vendor_part_id='content', content=content, extract_think_tags=True
+                )
+                if maybe_event is not None:  # pragma: no branch
+                    yield maybe_event
 
             # Handle the tool calls
             for dtc in choice.delta.tool_calls or []:
@@ -444,7 +448,7 @@ def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> us
     if isinstance(completion, chat.ChatCompletion):
         response_usage = completion.usage
     elif completion.x_groq is not None:
-        response_usage = completion.x_groq.usage  # pragma: no cover
+        response_usage = completion.x_groq.usage
 
     if response_usage is None:
         return usage.Usage()

pydantic_ai/models/huggingface.py

@@ -426,8 +426,12 @@ class HuggingFaceStreamedResponse(StreamedResponse):
 
             # Handle the text part of the response
             content = choice.delta.content
-            if content is not None:
-                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content)
+            if content:
+                maybe_event = self._parts_manager.handle_text_delta(
+                    vendor_part_id='content', content=content, extract_think_tags=True
+                )
+                if maybe_event is not None:  # pragma: no branch
+                    yield maybe_event
 
             for dtc in choice.delta.tool_calls or []:
                 maybe_event = self._parts_manager.handle_tool_call_delta(

pydantic_ai/models/mistral.py

@@ -601,7 +601,9 @@ class MistralStreamedResponse(StreamedResponse):
                             tool_call_id=maybe_tool_call_part.tool_call_id,
                         )
                 else:
-                    yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=text)
+                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=text)
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
 
             # Handle the explicit tool calls
             for index, dtc in enumerate(choice.delta.tool_calls or []):

pydantic_ai/models/openai.py

@@ -17,7 +17,7 @@ from pydantic_ai.providers import Provider, infer_provider
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
-from .._utils import guard_tool_call_id as _guard_tool_call_id, number_to_datetime
+from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc, number_to_datetime
 from ..messages import (
     AudioUrl,
     BinaryContent,
@@ -191,7 +191,17 @@ class OpenAIModel(Model):
         model_name: OpenAIModelName,
         *,
         provider: Literal[
-            'openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'heroku', 'github'
+            'openai',
+            'deepseek',
+            'azure',
+            'openrouter',
+            'moonshotai',
+            'vercel',
+            'grok',
+            'fireworks',
+            'together',
+            'heroku',
+            'github',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -290,7 +300,10 @@ class OpenAIModel(Model):
         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'
@@ -357,11 +370,17 @@ class OpenAIModel(Model):
         if not isinstance(response, chat.ChatCompletion):
             raise UnexpectedModelBehavior('Invalid response from OpenAI chat completions endpoint, expected JSON data')
 
+        if response.created:
+            timestamp = number_to_datetime(response.created)
+        else:
+            timestamp = _now_utc()
+            response.created = int(timestamp.timestamp())
+
         try:
             response = chat.ChatCompletion.model_validate(response.model_dump())
         except ValidationError as e:
             raise UnexpectedModelBehavior(f'Invalid response from OpenAI chat completions endpoint: {e}') from e
-        timestamp = number_to_datetime(response.created)
+
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
         # The `reasoning_content` is only present in DeepSeek models.
@@ -1003,8 +1022,12 @@ class OpenAIStreamedResponse(StreamedResponse):
 
             # Handle the text part of the response
             content = choice.delta.content
-            if content is not None:
-                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content)
+            if content:
+                maybe_event = self._parts_manager.handle_text_delta(
+                    vendor_part_id='content', content=content, extract_think_tags=True
+                )
+                if maybe_event is not None:  # pragma: no branch
+                    yield maybe_event
 
             # Handle reasoning part of the response, present in DeepSeek models
             if reasoning_content := getattr(choice.delta, 'reasoning_content', None):
@@ -1121,7 +1144,11 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                 )
 
             elif isinstance(chunk, responses.ResponseTextDeltaEvent):
-                yield self._parts_manager.handle_text_delta(vendor_part_id=chunk.content_index, content=chunk.delta)
+                maybe_event = self._parts_manager.handle_text_delta(
+                    vendor_part_id=chunk.content_index, content=chunk.delta
+                )
+                if maybe_event is not None:  # pragma: no branch
+                    yield maybe_event
 
             elif isinstance(chunk, responses.ResponseTextDoneEvent):
                 pass  # there's nothing we need to do here

pydantic_ai/models/test.py

@@ -269,10 +269,14 @@ class TestStreamedResponse(StreamedResponse):
                     mid = len(text) // 2
                     words = [text[:mid], text[mid:]]
                 self._usage += _get_string_usage('')
-                yield self._parts_manager.handle_text_delta(vendor_part_id=i, content='')
+                maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=i, content='')
+                if maybe_event is not None:  # pragma: no branch
+                    yield maybe_event
                 for word in words:
                     self._usage += _get_string_usage(word)
-                    yield self._parts_manager.handle_text_delta(vendor_part_id=i, content=word)
+                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=i, content=word)
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
             elif isinstance(part, ToolCallPart):
                 yield self._parts_manager.handle_tool_call_part(
                     vendor_part_id=i, tool_name=part.tool_name, args=part.args, tool_call_id=part.tool_call_id

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

@@ -62,6 +62,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .openrouter import OpenRouterProvider
 
         return OpenRouterProvider
+    elif provider == 'vercel':
+        from .vercel import VercelProvider
+
+        return VercelProvider
     elif provider == 'azure':
         from .azure import AzureProvider
 
@@ -99,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