kosong.chat_provider.kimi

  1import copy
  2import inspect
  3import mimetypes
  4import os
  5import uuid
  6from collections.abc import AsyncIterator, Mapping, Sequence
  7from typing import TYPE_CHECKING, Any, Literal, Self, Unpack, cast
  8
  9import httpx
 10from openai import AsyncOpenAI, AsyncStream, BaseModel, OpenAIError, omit
 11from openai._types import RequestFiles, RequestOptions
 12from openai.types.chat import (
 13    ChatCompletion,
 14    ChatCompletionChunk,
 15    ChatCompletionMessageFunctionToolCall,
 16    ChatCompletionMessageParam,
 17    ChatCompletionToolParam,
 18)
 19from openai.types.completion_usage import CompletionUsage
 20from typing_extensions import TypedDict
 21
 22from kosong.chat_provider import (
 23    ChatProvider,
 24    ChatProviderError,
 25    RetryableChatProvider,
 26    StreamedMessagePart,
 27    ThinkingEffort,
 28    TokenUsage,
 29)
 30from kosong.chat_provider.openai_common import (
 31    close_replaced_openai_client,
 32    convert_error,
 33    create_openai_client,
 34    tool_to_openai,
 35)
 36from kosong.message import (
 37    ContentPart,
 38    Message,
 39    TextPart,
 40    ThinkPart,
 41    ToolCall,
 42    ToolCallPart,
 43    VideoURLPart,
 44)
 45from kosong.tooling import Tool
 46from kosong.utils.jsonschema import JsonDict, ensure_property_types
 47
 48if TYPE_CHECKING:
 49
 50    def type_check(kimi: "Kimi"):
 51        _: ChatProvider = kimi
 52        _: RetryableChatProvider = kimi
 53
 54
 55class ThinkingConfig(TypedDict, total=False):
 56    type: Literal["enabled", "disabled"]
 57    keep: Any
 58    """Moonshot-specific ``thinking.keep`` switch for preserved thinking.
 59    Forwarded verbatim to the API; callers are responsible for choosing a value
 60    the server accepts (e.g. ``"all"``)."""
 61
 62
 63class ExtraBody(TypedDict, total=False, extra_items=Any):
 64    thinking: ThinkingConfig
 65
 66
 67class Kimi:
 68    """
 69    A chat provider that uses the Kimi API.
 70
 71    >>> chat_provider = Kimi(model="kimi-k2-turbo-preview", api_key="sk-1234567890")
 72    >>> chat_provider.name
 73    'kimi'
 74    >>> chat_provider.model_name
 75    'kimi-k2-turbo-preview'
 76    >>> chat_provider.with_generation_kwargs(temperature=0)._generation_kwargs
 77    {'temperature': 0}
 78    >>> chat_provider._generation_kwargs
 79    {}
 80    """
 81
 82    name = "kimi"
 83
 84    class GenerationKwargs(TypedDict, total=False):
 85        """
 86        See https://platform.moonshot.ai/docs/api/chat#request-body.
 87        """
 88
 89        max_completion_tokens: int | None
 90        max_tokens: int | None
 91        """Deprecated alias. Normalized to ``max_completion_tokens`` before requests."""
 92        temperature: float | None
 93        top_p: float | None
 94        n: int | None
 95        presence_penalty: float | None
 96        frequency_penalty: float | None
 97        stop: str | list[str] | None
 98        prompt_cache_key: str | None
 99        reasoning_effort: str | None
100        """Legacy explicit passthrough. `with_thinking` uses `extra_body.thinking` instead."""
101        extra_body: ExtraBody | None
102
103    def __init__(
104        self,
105        *,
106        model: str,
107        api_key: str | None = None,
108        base_url: str | None = None,
109        stream: bool = True,
110        **client_kwargs: Any,
111    ):
112        if api_key is None:
113            api_key = os.getenv("KIMI_API_KEY")
114        if api_key is None:
115            raise ChatProviderError(
116                "The api_key client option or the KIMI_API_KEY environment variable is not set"
117            )
118        if base_url is None:
119            base_url = os.getenv("KIMI_BASE_URL", "https://api.moonshot.ai/v1")
120
121        self.model: str = model
122        """The name of the model to use."""
123        self.stream: bool = stream
124        """Whether to generate responses as a stream."""
125        self._api_key: str | None = api_key
126        self._base_url: str | None = base_url
127        self._client_kwargs: dict[str, Any] = dict(client_kwargs)
128        self.client: AsyncOpenAI = create_openai_client(
129            api_key=self._api_key,
130            base_url=self._base_url,
131            client_kwargs=self._client_kwargs,
132        )
133        """The underlying `AsyncOpenAI` client."""
134        self._generation_kwargs: Kimi.GenerationKwargs = {}
135        self._thinking_effort: ThinkingEffort | None = None
136        """Thinking state kept separately from parameters serialized onto the wire."""
137
138    @property
139    def model_name(self) -> str:
140        return self.model
141
142    @property
143    def thinking_effort(self) -> ThinkingEffort | None:
144        return self._thinking_effort
145
146    async def generate(
147        self,
148        system_prompt: str,
149        tools: Sequence[Tool],
150        history: Sequence[Message],
151        *,
152        generation_overrides: Mapping[str, Any] | None = None,
153    ) -> "KimiStreamedMessage":
154        messages: list[ChatCompletionMessageParam] = []
155        if system_prompt:
156            messages.append({"role": "system", "content": system_prompt})
157        messages.extend(_convert_message(message) for message in history)
158
159        generation_kwargs: dict[str, Any] = dict(self._generation_kwargs)
160        if generation_overrides:
161            generation_kwargs.update(
162                _normalize_generation_kwargs(
163                    cast(Kimi.GenerationKwargs, dict(generation_overrides))
164                )
165            )
166        if generation_kwargs.get("max_completion_tokens") is None:
167            generation_kwargs.pop("max_completion_tokens", None)
168
169        try:
170            if self.stream:
171                # ``with_raw_response`` eagerly reads the response body in the
172                # OpenAI SDK. Use the normal streaming path so callers receive
173                # the AsyncStream as soon as response headers arrive.
174                parsed_response = await cast(Any, self.client.chat.completions.create)(
175                    model=self.model,
176                    messages=messages,
177                    tools=(_convert_tool(tool) for tool in tools),
178                    stream=True,
179                    stream_options={"include_usage": True},
180                    **generation_kwargs,
181                )
182                trace_id = parsed_response.response.headers.get("x-trace-id")
183            else:
184                raw_response = await self.client.chat.completions.with_raw_response.create(
185                    model=self.model,
186                    messages=messages,
187                    tools=(_convert_tool(tool) for tool in tools),
188                    stream=False,
189                    stream_options=omit,
190                    **generation_kwargs,
191                )
192                trace_id = raw_response.headers.get("x-trace-id")
193                parsed_response = raw_response.parse()
194                if inspect.isawaitable(parsed_response):
195                    parsed_response = await parsed_response
196            return KimiStreamedMessage(parsed_response, trace_id=trace_id)
197        except (OpenAIError, httpx.HTTPError) as e:
198            raise convert_error(e) from e
199
200    def on_retryable_error(self, error: BaseException) -> bool:
201        old_client = self.client
202        # Read api_key from the live client (not self._api_key) so that
203        # OAuth token refreshes applied via client.api_key are preserved.
204        current_api_key = old_client.api_key
205        self.client = create_openai_client(
206            api_key=current_api_key,
207            base_url=self._base_url,
208            client_kwargs=self._client_kwargs,
209        )
210        self._api_key = current_api_key
211        close_replaced_openai_client(old_client, client_kwargs=self._client_kwargs)
212        return True
213
214    def with_thinking(self, effort: ThinkingEffort) -> Self:
215        new_self = self.with_extra_body(
216            {
217                "thinking": {
218                    "type": "enabled" if effort != "off" else "disabled",
219                }
220            }
221        )
222        new_self._thinking_effort = effort
223        return new_self
224
225    def with_generation_kwargs(self, **kwargs: Unpack[GenerationKwargs]) -> Self:
226        """
227        Copy the chat provider, updating the generation kwargs with the given values.
228
229        Returns:
230            Self: A new instance of the chat provider with updated generation kwargs.
231        """
232        new_self = copy.copy(self)
233        new_self._generation_kwargs = copy.deepcopy(self._generation_kwargs)
234        new_self._generation_kwargs.update(_normalize_generation_kwargs(kwargs))
235        return new_self
236
237    def with_extra_body(self, extra_body: ExtraBody) -> Self:
238        """
239        Copy the chat provider, updating the extra_body in generation kwargs.
240
241        Top-level keys follow last-writer-wins semantics, except for the
242        ``thinking`` key: its sub-dict is merged field-by-field so that a
243        later call adding ``thinking.keep`` does not erase a ``thinking.type``
244        installed by an earlier ``with_thinking`` call.
245
246        Returns:
247            Self: A new instance of the chat provider with updated extra_body.
248        """
249        new_self = copy.copy(self)
250        new_self._generation_kwargs = copy.deepcopy(self._generation_kwargs)
251        old_extra_body = new_self._generation_kwargs.get("extra_body") or {}
252        new_extra_body: ExtraBody = {**old_extra_body, **extra_body}
253        old_thinking = old_extra_body.get("thinking")
254        new_thinking = extra_body.get("thinking")
255        if old_thinking is not None and new_thinking is not None:
256            new_extra_body["thinking"] = {**old_thinking, **new_thinking}
257        new_self._generation_kwargs["extra_body"] = new_extra_body
258        return new_self
259
260    @property
261    def model_parameters(self) -> dict[str, Any]:
262        """
263        The parameters of the model to use.
264
265        For tracing/logging purposes.
266        """
267
268        model_parameters: dict[str, Any] = {"base_url": str(self.client.base_url)}
269        model_parameters.update(self._generation_kwargs)
270        return model_parameters
271
272    @property
273    def files(self) -> "KimiFiles":
274        return KimiFiles(self.client)
275
276
277class KimiFiles:
278    def __init__(self, client: AsyncOpenAI) -> None:
279        self._client = client
280
281    async def upload_video(self, *, data: bytes, mime_type: str) -> VideoURLPart:
282        """Upload a video to Kimi files API and return a video URL content part."""
283        if not mime_type.startswith("video/"):
284            raise ChatProviderError(f"Expected a video mime type, got {mime_type}")
285        url = await self._upload_file(data=data, mime_type=mime_type, purpose="video")
286        return VideoURLPart(video_url=VideoURLPart.VideoURL(url=url))
287
288    async def _upload_file(self, *, data: bytes, mime_type: str, purpose: "KimiFilePurpose") -> str:
289        filename = _guess_filename(mime_type)
290        files: RequestFiles = {"file": (filename, data, mime_type)}
291        options: RequestOptions = {"headers": {"Content-Type": "multipart/form-data"}}
292        try:
293            response: KimiFileObject = await self._client.post(
294                "/files",
295                cast_to=KimiFileObject,
296                body={"purpose": purpose},
297                files=files,
298                options=options,
299            )
300        except (OpenAIError, httpx.HTTPError) as e:
301            raise convert_error(e) from e
302        return f"ms://{response.id}"
303
304
305class KimiFileObject(BaseModel):
306    id: str
307
308
309type KimiFilePurpose = Literal["video", "image"]
310
311
312def _guess_filename(mime_type: str) -> str:
313    extension = mimetypes.guess_extension(mime_type) or ".bin"
314    return f"upload{extension}"
315
316
317def _normalize_generation_kwargs(kwargs: Kimi.GenerationKwargs) -> Kimi.GenerationKwargs:
318    normalized: dict[str, Any] = dict(kwargs)
319    if "max_tokens" in normalized:
320        max_tokens = normalized.pop("max_tokens")
321        if "max_completion_tokens" not in normalized:
322            normalized["max_completion_tokens"] = max_tokens
323    return cast(Kimi.GenerationKwargs, normalized)
324
325
326def _convert_message(message: Message) -> ChatCompletionMessageParam:
327    message = message.model_copy(deep=True)
328    reasoning_content: str = ""
329    content: list[ContentPart] = []
330    has_reasoning = False
331    for part in message.content:
332        if isinstance(part, ThinkPart):
333            has_reasoning = True
334            reasoning_content += part.think
335        else:
336            content.append(part)
337    message.content = content
338    dumped_message = message.model_dump(exclude_none=True)
339    if (
340        message.role == "assistant"
341        and message.tool_calls
342        and _is_effectively_empty_content_parts(content)
343    ):
344        # OpenAI-compatible APIs allow assistant tool-call messages to omit
345        # `content`, but the Kimi-for-Coding compat layer rejects a content
346        # list that contains an empty text part (observed: `content:
347        # [{"type": "text", "text": ""}]` -> 400 "text content is empty").
348        # Dropping `content` entirely is always accepted, so do that whenever
349        # the visible content is effectively empty alongside a tool call.
350        dumped_message.pop("content", None)
351    if has_reasoning:
352        dumped_message["reasoning_content"] = reasoning_content
353    return cast(ChatCompletionMessageParam, dumped_message)
354
355
356def _is_effectively_empty_content_parts(content: Sequence[ContentPart]) -> bool:
357    for part in content:
358        if not isinstance(part, TextPart):
359            return False
360        if part.text.strip():
361            return False
362    return True
363
364
365def _convert_tool(tool: Tool) -> ChatCompletionToolParam:
366    if tool.name.startswith("$"):
367        # Kimi builtin functions start with `$`
368        return cast(
369            ChatCompletionToolParam,
370            {
371                "type": "builtin_function",
372                "function": {
373                    "name": tool.name,
374                    # no need to set description and parameters
375                },
376            },
377        )
378    converted = tool_to_openai(tool)
379    # Moonshot's API rejects parameter schemas whose nested properties omit
380    # `type` (e.g. enum-only properties exposed by some MCP servers). Patch
381    # the schema locally so such tools keep working against Kimi without
382    # requiring every MCP server author to tighten their schemas.
383    function = converted["function"]
384    parameters = function.get("parameters")
385    if isinstance(parameters, dict):
386        normalized = ensure_property_types(cast(JsonDict, parameters))
387        function["parameters"] = cast(dict[str, object], normalized)
388    return converted
389
390
391class KimiStreamedMessage:
392    """The streamed message of the Kimi chat provider."""
393
394    def __init__(
395        self,
396        response: ChatCompletion | AsyncStream[ChatCompletionChunk],
397        *,
398        trace_id: str | None = None,
399    ):
400        if isinstance(response, ChatCompletion):
401            self._iter = self._convert_non_stream_response(response)
402        else:
403            self._iter = self._convert_stream_response(response)
404        self._id: str | None = None
405        self._usage: CompletionUsage | None = None
406        self._trace_id = trace_id
407
408    def __aiter__(self) -> AsyncIterator[StreamedMessagePart]:
409        return self
410
411    async def __anext__(self) -> StreamedMessagePart:
412        return await self._iter.__anext__()
413
414    @property
415    def id(self) -> str | None:
416        return self._id
417
418    @property
419    def trace_id(self) -> str | None:
420        return self._trace_id
421
422    @property
423    def usage(self) -> TokenUsage | None:
424        if self._usage:
425            cached = 0
426            other_input = self._usage.prompt_tokens
427            if hasattr(self._usage, "cached_tokens"):
428                # https://platform.moonshot.cn/docs/api/chat#%E8%BF%94%E5%9B%9E%E5%86%85%E5%AE%B9
429                # TODO: delete this when Moonshot API becomes compatible with OpenAI API
430                cached = getattr(self._usage, "cached_tokens") or 0  # noqa: B009
431                other_input -= cached
432            elif (
433                self._usage.prompt_tokens_details
434                and self._usage.prompt_tokens_details.cached_tokens
435            ):
436                cached = self._usage.prompt_tokens_details.cached_tokens
437                other_input -= cached
438            return TokenUsage(
439                input_other=other_input,
440                output=self._usage.completion_tokens,
441                input_cache_read=cached,
442            )
443        return None
444
445    async def _convert_non_stream_response(
446        self,
447        response: ChatCompletion,
448    ) -> AsyncIterator[StreamedMessagePart]:
449        self._id = response.id
450        self._usage = response.usage
451        message = response.choices[0].message
452        reasoning_content = getattr(message, "reasoning_content", None)
453        if reasoning_content is not None:
454            assert isinstance(reasoning_content, str)
455            yield ThinkPart(think=reasoning_content)
456        if message.content:
457            yield TextPart(text=message.content)
458        if message.tool_calls:
459            for tool_call in message.tool_calls:
460                if isinstance(tool_call, ChatCompletionMessageFunctionToolCall):
461                    yield ToolCall(
462                        id=tool_call.id or str(uuid.uuid4()),
463                        function=ToolCall.FunctionBody(
464                            name=tool_call.function.name,
465                            arguments=tool_call.function.arguments,
466                        ),
467                    )
468
469    async def _convert_stream_response(
470        self,
471        response: AsyncIterator[ChatCompletionChunk],
472    ) -> AsyncIterator[StreamedMessagePart]:
473        try:
474            async for chunk in response:
475                if chunk.id:
476                    self._id = chunk.id
477                if usage := extract_usage_from_chunk(chunk):
478                    self._usage = usage
479
480                if not chunk.choices:
481                    continue
482
483                delta = chunk.choices[0].delta
484
485                # convert thinking content — an empty string means "reasoned
486                # but empty", not "no reasoning": keep it as a ThinkPart so
487                # the distinction round-trips to the server (preserved-thinking
488                # backends require reasoning_content on every assistant turn)
489                reasoning_content = getattr(delta, "reasoning_content", None)
490                if reasoning_content is not None:
491                    assert isinstance(reasoning_content, str)
492                    yield ThinkPart(think=reasoning_content)
493
494                # convert text content
495                if delta.content:
496                    yield TextPart(text=delta.content)
497
498                # convert tool calls
499                for tool_call in delta.tool_calls or []:
500                    if not tool_call.function:
501                        continue
502
503                    if tool_call.function.name:
504                        yield ToolCall(
505                            id=tool_call.id or str(uuid.uuid4()),
506                            function=ToolCall.FunctionBody(
507                                name=tool_call.function.name,
508                                arguments=tool_call.function.arguments,
509                            ),
510                        )
511                    elif tool_call.function.arguments:
512                        yield ToolCallPart(
513                            arguments_part=tool_call.function.arguments,
514                        )
515                    else:
516                        # skip empty tool calls
517                        pass
518        except (OpenAIError, httpx.HTTPError) as e:
519            raise convert_error(e) from e
520
521
522def extract_usage_from_chunk(chunk: ChatCompletionChunk) -> CompletionUsage | None:
523    if chunk.usage:
524        return chunk.usage
525    if not chunk.choices:
526        return None
527    choice_dump: dict[str, object] = chunk.choices[0].model_dump()
528    raw_usage = choice_dump.get("usage")
529    if isinstance(raw_usage, CompletionUsage):
530        return raw_usage
531    if isinstance(raw_usage, dict):
532        return CompletionUsage.model_validate(raw_usage)
533    return None
534
535
536if __name__ == "__main__":
537
538    async def _dev_main():
539        chat = Kimi(model="kimi-k2-turbo-preview", stream=False)
540        system_prompt = ""
541        history = [
542            Message(role="user", content="Hello, who is Confucius?"),
543        ]
544        stream = await chat.with_generation_kwargs(
545            temperature=0,
546            max_completion_tokens=1000,
547        ).generate(system_prompt, [], history)
548        async for part in stream:
549            print(part.model_dump(exclude_none=True))
550        print("id:", stream.id)
551        print("usage:", stream.usage)
552
553    import asyncio
554
555    from dotenv import load_dotenv
556
557    load_dotenv()
558    asyncio.run(_dev_main())