rig_core/providers/moonshot/
extension.rs

1//! Moonshot's typed request options and reply extras. One entry serves its
2//! OpenAI-format API (<https://platform.kimi.ai/docs/api/chat>) and its
3//! Anthropic-format Messages API, which takes no Moonshot field of its own.
4//!
5//! ```
6//! use rig_core::completion::CompletionRequest;
7//! use rig_core::providers::moonshot::extension::MoonshotOptions;
8//!
9//! let options = MoonshotOptions::new().prompt_cache_key("session-7");
10//! let request = CompletionRequest::new("hi").provider_option(options);
11//! # let _ = request;
12//! ```
13
14use serde::Serialize;
15use serde_json::Value;
16
17use crate::completion::provider_options::reply_field;
18use crate::completion::{ExtensionOptions, ProviderExtension, ReplyExtras};
19use crate::message::Api;
20use crate::providers::anthropic::extension::MessagesStop;
21use crate::providers::anthropic::wire::MESSAGES_API;
22
23/// Moonshot's extension marker.
24#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
25pub struct MoonshotExt;
26
27impl ProviderExtension for MoonshotExt {
28    const PROVIDER: &'static str = super::PROVIDER_NAME;
29    type Options = MoonshotOptions;
30    type Extras = MoonshotExtras;
31}
32
33/// Moonshot's request options, by route. The field setters here write the
34/// same field as [`chat`](Self::chat), which replaces every field set
35/// before it; the Messages API takes none of them.
36#[non_exhaustive]
37#[derive(Clone, Debug, Default, PartialEq, Serialize)]
38pub struct MoonshotOptions {
39    /// Sent on the OpenAI-format API only.
40    #[serde(rename = "openai.chat")]
41    pub chat: MoonshotChat,
42}
43
44impl MoonshotOptions {
45    /// No option set.
46    pub fn new() -> Self {
47        Self::default()
48    }
49
50    /// Send `chat` on the OpenAI-format API.
51    pub fn chat(mut self, chat: MoonshotChat) -> Self {
52        self.chat = chat;
53        self
54    }
55
56    /// Keep `keep` earlier turns' thinking in context, on the OpenAI-format
57    /// API, as [`MoonshotChat::thinking_keep`].
58    pub fn thinking_keep(mut self, keep: ThinkingKeep) -> Self {
59        self.chat = std::mem::take(&mut self.chat).thinking_keep(keep);
60        self
61    }
62
63    /// Route the prompt cache by `key`, on the OpenAI-format API, as
64    /// [`MoonshotChat::prompt_cache_key`].
65    pub fn prompt_cache_key(mut self, key: impl Into<String>) -> Self {
66        self.chat = std::mem::take(&mut self.chat).prompt_cache_key(key);
67        self
68    }
69}
70
71impl ExtensionOptions for MoonshotOptions {
72    type Ext = MoonshotExt;
73}
74
75/// The fields Moonshot's OpenAI-format API takes.
76#[non_exhaustive]
77#[derive(Clone, Debug, Default, PartialEq, Serialize)]
78pub struct MoonshotChat {
79    /// Thinking fields sent beside the mapped thinking type.
80    #[serde(skip_serializing_if = "MoonshotThinking::is_empty")]
81    pub thinking: MoonshotThinking,
82    /// Routes requests that share a prefix to the same cache.
83    #[serde(skip_serializing_if = "Option::is_none")]
84    pub prompt_cache_key: Option<String>,
85}
86
87/// The `thinking` fields Moonshot takes beside the type the generation
88/// options map.
89#[non_exhaustive]
90#[derive(Clone, Debug, Default, PartialEq, Serialize)]
91pub struct MoonshotThinking {
92    /// Which earlier turns' thinking the model keeps in context.
93    #[serde(skip_serializing_if = "Option::is_none")]
94    pub keep: Option<ThinkingKeep>,
95}
96
97impl MoonshotThinking {
98    fn is_empty(&self) -> bool {
99        self.keep.is_none()
100    }
101}
102
103/// Which earlier turns' thinking the model keeps in context.
104#[non_exhaustive]
105#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
106#[serde(rename_all = "lowercase")]
107pub enum ThinkingKeep {
108    /// Every turn's.
109    All,
110}
111
112impl MoonshotChat {
113    /// No option set.
114    pub fn new() -> Self {
115        Self::default()
116    }
117
118    /// Keep `keep` earlier turns' thinking in context.
119    pub fn thinking_keep(mut self, keep: ThinkingKeep) -> Self {
120        self.thinking.keep = Some(keep);
121        self
122    }
123
124    /// Route the prompt cache by `key`.
125    pub fn prompt_cache_key(mut self, key: impl Into<String>) -> Self {
126        self.prompt_cache_key = Some(key.into());
127        self
128    }
129}
130
131/// Moonshot's reply fields. Each is `None` when the reply lacks it.
132#[non_exhaustive]
133#[derive(Clone, Debug, Default, PartialEq)]
134pub struct MoonshotExtras {
135    /// The first choice's own usage. OpenAI-format API only.
136    pub choice_usage: Option<Value>,
137    /// Prompt token details. OpenAI-format API only.
138    pub prompt_tokens_details: Option<Value>,
139    /// `stop_reason`, verbatim. Messages route.
140    pub stop_reason: Option<String>,
141    /// `stop_sequence`: the stop sequence the turn ended on. Messages route.
142    pub stop_sequence: Option<String>,
143}
144
145impl ReplyExtras for MoonshotExtras {
146    fn from_reply(api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
147        if api.as_str() == MESSAGES_API {
148            let MessagesStop {
149                stop_reason,
150                stop_sequence,
151            } = MessagesStop::read("Moonshot", api, raw)?;
152            return Ok(Self {
153                stop_reason,
154                stop_sequence,
155                ..Self::default()
156            });
157        }
158        if api.as_str() != "openai.chat" {
159            return Ok(Self::default());
160        }
161        Ok(Self {
162            choice_usage: reply_field(raw, "/choices/0/usage")?,
163            prompt_tokens_details: reply_field(raw, "/usage/prompt_tokens_details")?,
164            ..Self::default()
165        })
166    }
167}
168
169#[cfg(test)]
170mod tests;
rig_core/providers/moonshot/extension.rs

rig_core/providers/moonshot/
extension.rs
