# -*- coding: utf-8 -*- """The Moonshot AI chat model implementation.""" from collections import OrderedDict from datetime import datetime from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type from pydantic import BaseModel, Field from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES from .._model_response import ChatResponse, StructuredResponse from .._model_usage import ChatUsage from ..._utils._common import _generate_id from ...credential import MoonshotCredential from ...formatter import FormatterBase, MoonshotChatFormatter from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock from ...tool import ToolChoice if TYPE_CHECKING: from openai.types.chat import ChatCompletion from openai import AsyncStream else: ChatCompletion = Any AsyncStream = Any class MoonshotChatModel(ChatModelBase): """The Moonshot AI chat model.""" class Parameters(BaseModel): """The parameters for the Moonshot AI chat model.""" max_tokens: int | None = Field( default=None, title="Max Tokens", description="The maximum number of tokens for the LLM output.", gt=0, ) thinking_enable: bool = Field( default=False, title="Thinking", description=( "Whether to enable thinking mode. For kimi-k2-thinking, " "thinking is always enabled. For kimi-k2.6, thinking is " "enabled by default but can be disabled." ), ) reasoning_effort: Literal["low", "high", "max"] | None = Field( default=None, title="Reasoning Effort", description=( "The reasoning effort level for kimi-k3. " "Supports 'low', 'high', and 'max' " "(default 'max')." ), ) temperature: float | None = Field( default=None, title="Temperature", description="The temperature for the LLM output.", ge=0, le=1, ) top_p: float | None = Field( default=None, title="Top P", description="The top P value for the LLM output.", gt=0, le=1, ) type: Literal["moonshot_chat"] = "moonshot_chat" """The type of the chat model.""" def __init__( self, credential: MoonshotCredential, model: str, parameters: "MoonshotChatModel.Parameters | None" = None, stream: bool = True, max_retries: int = 3, retry_delay: float = 1.0, context_size: int = 131072, formatter: FormatterBase | None = None, client_kwargs: dict[str, Any] | None = None, ) -> None: """Initialize the Moonshot AI chat model. Args: credential (`MoonshotCredential`): The Moonshot AI credential used to authenticate API calls. model (`str`): The model name, e.g. ``moonshot-v1-8k``, ``kimi-k2.6``. parameters (`MoonshotChatModel.Parameters | None`, defaults to \ `None`): The API parameters. When ``None``, the default parameters will be used. stream (`bool`, defaults to `True`): Whether to enable streaming output. max_retries (`int`, defaults to `3`): The maximum number of retries for the API. retry_delay (`float`, defaults to `1.0`): Seconds to sleep between retry attempts. context_size (`int`, defaults to `131072`): The model context size used for context compression. formatter (`FormatterBase | None`, defaults to `None`): The formatter that converts ``Msg`` objects to the format required by the API. When ``None``, a ``MoonshotChatFormatter`` instance will be used. client_kwargs (`dict[str, Any] | None`, defaults to `None`): Extra keyword arguments forwarded to ``openai.AsyncClient`` (e.g. ``timeout``, ``default_headers``, ``http_client``). """ super().__init__( credential=credential, model=model, parameters=parameters or self.Parameters(), stream=stream, max_retries=max_retries, retry_delay=retry_delay, context_size=context_size, ) self.formatter = formatter or MoonshotChatFormatter() self.client_kwargs = client_kwargs or {} import openai self.client: openai.AsyncClient = openai.AsyncClient( api_key=self.credential.api_key.get_secret_value(), base_url=self.credential.base_url, **self.client_kwargs, ) @classmethod def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: import openai return ( openai.APIConnectionError, openai.APITimeoutError, openai.RateLimitError, openai.InternalServerError, ) async def _call_api( self, model_name: str, messages: list[Msg], tools: list[dict] | None = None, tool_choice: ToolChoice | None = None, **generate_kwargs: Any, ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: """Call the Moonshot AI chat API (OpenAI-compatible). Args: model_name (`str`): The model name to use for this call. messages (`list`): A list of message dicts with ``role`` and ``content`` keys. tools (`list[dict]`, default `None`): The tools JSON schemas. tool_choice (`ToolChoice | None`, optional): Controls which (if any) tool is called by the model. **generate_kwargs (`Any`): Extra keyword arguments forwarded to the API. Returns: `ChatResponse | AsyncGenerator[ChatResponse, None]`: A ``ChatResponse`` when streaming is disabled, or an async generator of ``ChatResponse`` objects when streaming is enabled. """ formatted_messages = await self.formatter.format(messages) kwargs: dict[str, Any] = { "model": model_name, "messages": formatted_messages, "stream": self.stream, } is_k3 = model_name == "kimi-k3" if self.parameters.max_tokens is not None: if is_k3: kwargs["max_completion_tokens"] = self.parameters.max_tokens else: kwargs["max_tokens"] = self.parameters.max_tokens if self.parameters.temperature is not None: kwargs["temperature"] = self.parameters.temperature if self.parameters.top_p is not None: kwargs["top_p"] = self.parameters.top_p kwargs.update(generate_kwargs) if is_k3: if self.parameters.reasoning_effort is not None: kwargs["reasoning_effort"] = self.parameters.reasoning_effort else: thinking_type = ( "enabled" if self.parameters.thinking_enable else "disabled" ) kwargs.setdefault("extra_body", {}) kwargs["extra_body"].setdefault("thinking", {}) kwargs["extra_body"]["thinking"].setdefault( "type", thinking_type, ) fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) if fmt_tools: kwargs["tools"] = fmt_tools if fmt_tool_choice is not None: kwargs["tool_choice"] = fmt_tool_choice if self.stream: kwargs["stream_options"] = {"include_usage": True} start_datetime = datetime.now() response = await self.client.chat.completions.create(**kwargs) if self.stream: return self._parse_stream_response(start_datetime, response) return self._parse_completion_response(start_datetime, response) async def _parse_stream_response( self, start_datetime: datetime, response: AsyncStream, ) -> AsyncGenerator[ChatResponse, None]: """Parse the Moonshot AI streaming response. Args: start_datetime (`datetime`): The start datetime of the response generation. response (`AsyncStream`): The OpenAI-compatible async stream object. Yields: `ChatResponse`: Incremental ``ChatResponse`` objects with ``is_last=False`` followed by a final one with ``is_last=True``. """ usage = None response_id: str = _generate_id() text_id: str = _generate_id() thinking_id: str = _generate_id() # The mapping from index to tool call id tool_call_mapping: dict = OrderedDict() async with response as stream: async for chunk in stream: delta_res = ChatResponse( content=[], is_last=False, id=response_id, ) # Update the response ID if exists response_id = getattr(chunk, "id", None) or response_id delta_res.id = response_id if chunk.usage: u = chunk.usage usage = ChatUsage( input_tokens=u.prompt_tokens, output_tokens=u.completion_tokens, time=(datetime.now() - start_datetime).total_seconds(), cache_input_tokens=getattr( u, "cached_tokens", 0, ), ) if not chunk.choices: continue choice = chunk.choices[0] delta = choice.delta # Thinking if getattr(delta, "reasoning_content", None): delta_res.append_thinking( block_id=thinking_id, thinking=delta.reasoning_content, ) # Text if getattr(delta, "content", None): delta_res.append_text( block_id=text_id, text=delta.content, ) # Tool call for tool_call in getattr(delta, "tool_calls", None) or []: index = tool_call.index fn = getattr(tool_call, "function", None) delta_name = getattr(fn, "name", None) if fn else None delta_args = getattr(fn, "arguments", None) if fn else None # Record the id and name in case following deltas # don't provide them if index not in tool_call_mapping: tool_call_mapping[index] = ( tool_call.id, delta_name or "unknown", ) stored_id, stored_name = tool_call_mapping[index] delta_res.append_tool_call( block_id=tool_call.id or stored_id, name=delta_name or stored_name, input=delta_args or "", ) if delta_res.content or usage: delta_res.usage = usage yield delta_res def _parse_completion_response( self, start_datetime: datetime, response: ChatCompletion, ) -> ChatResponse: """Parse the Moonshot AI non-streaming response. Args: start_datetime (`datetime`): The start datetime of the response generation. response (`ChatCompletion`): The OpenAI-compatible chat completion object. Returns: `ChatResponse`: A single ``ChatResponse`` with ``is_last=True``. """ content_blocks: List[ThinkingBlock | TextBlock | ToolCallBlock] = [] if response.choices: choice = response.choices[0] reasoning = getattr(choice.message, "reasoning_content", None) if reasoning: content_blocks.append(ThinkingBlock(thinking=reasoning)) if choice.message.content: content_blocks.append(TextBlock(text=choice.message.content)) for tool_call in choice.message.tool_calls or []: content_blocks.append( ToolCallBlock( id=tool_call.id, name=tool_call.function.name, input=tool_call.function.arguments, ), ) usage = None if response.usage: u = response.usage usage = ChatUsage( input_tokens=u.prompt_tokens, output_tokens=u.completion_tokens, time=(datetime.now() - start_datetime).total_seconds(), cache_input_tokens=getattr(u, "cached_tokens", 0), ) resp_kwargs: dict[str, Any] = { "content": content_blocks, "is_last": True, "usage": usage, } response_id = getattr(response, "id", None) if response_id: resp_kwargs["id"] = response_id return ChatResponse(**resp_kwargs) def _format_tools( self, tools: list[dict] | None, tool_choice: ToolChoice | None, ) -> tuple[list[dict] | None, str | dict | None]: """Validate, filter, and format tools and tool_choice for the Moonshot AI API. When ``tool_choice.tools`` is specified the schemas list is filtered to only those tools. When ``tool_choice.mode`` is a specific tool name (str) the model is forced to call exactly that tool without needing to filter the list, preserving prompt-cache efficiency. Args: tools (`list[dict] | None`, optional): The raw tool schemas. tool_choice (`ToolChoice | None`, optional): The tool choice configuration. Returns: `tuple[list[dict] | None, str | dict | None]`: A tuple of (formatted_tools, formatted_tool_choice). """ if tool_choice and tools: self._validate_tool_choice(tool_choice, tools) if tool_choice.tools: allowed = set(tool_choice.tools) tools = [t for t in tools if t["function"]["name"] in allowed] if not tool_choice: return tools, None mode = tool_choice.mode if mode not in _TOOL_CHOICE_LITERAL_MODES: return tools, {"type": "function", "function": {"name": mode}} return tools, mode async def _call_api_with_structured_output( self, model_name: str, messages: list[Msg], structured_model: Type[BaseModel] | dict, tool_choice: ToolChoice | None = None, **kwargs: Any, ) -> StructuredResponse: """Moonshot-specific override for structured output. Moonshot rejects ``tool_choice="required"`` or an object-form ``tool_choice`` when thinking mode is enabled. In that case we default ``tool_choice`` to ``"auto"`` and rely on the base class's injected system-reminder prompt to guide the model. When thinking is disabled, this falls through to the base implementation. Args: model_name (`str`): The model name to use for this call. messages (`list[Msg]`): The context for the LLM to generate the structured output. structured_model (`Type[BaseModel] | dict`): A Pydantic model class or a JSON schema dict describing the required output structure. tool_choice (`ToolChoice | None`, defaults to `None`): The tool_choice forwarded to ``_call_api``. When ``None`` and thinking mode is enabled, it is downgraded to ``ToolChoice(mode="auto")``; otherwise the base default (force the structured-output tool) is used. **kwargs (`Any`): Additional keyword arguments forwarded to ``_call_api``. Returns: `StructuredResponse`: The structured response whose ``content`` is the validated output dict matching ``structured_model``. """ if tool_choice is None and self.parameters.thinking_enable: tool_choice = ToolChoice(mode="auto") return await super()._call_api_with_structured_output( model_name=model_name, messages=messages, structured_model=structured_model, tool_choice=tool_choice, **kwargs, )