rig/providers/
moonshot.rs

1//! Moonshot API client and Rig integration
2//!
3//! # Example
4//! ```
5//! use rig::providers::moonshot;
6//!
7//! let client = moonshot::Client::new("YOUR_API_KEY");
8//!
9//! let moonshot_model = client.completion_model(moonshot::MOONSHOT_CHAT);
10//! ```
11use crate::client::{ClientBuilderError, CompletionClient, ProviderClient};
12use crate::json_utils::merge;
13use crate::providers::openai::send_compatible_streaming_request;
14use crate::streaming::StreamingCompletionResponse;
15use crate::{
16    completion::{self, CompletionError, CompletionRequest},
17    json_utils,
18    providers::openai,
19};
20use crate::{impl_conversion_traits, message};
21use serde::Deserialize;
22use serde_json::{Value, json};
23
24// ================================================================
25// Main Moonshot Client
26// ================================================================
27const MOONSHOT_API_BASE_URL: &str = "https://api.moonshot.cn/v1";
28
29pub struct ClientBuilder<'a> {
30    api_key: &'a str,
31    base_url: &'a str,
32    http_client: Option<reqwest::Client>,
33}
34
35impl<'a> ClientBuilder<'a> {
36    pub fn new(api_key: &'a str) -> Self {
37        Self {
38            api_key,
39            base_url: MOONSHOT_API_BASE_URL,
40            http_client: None,
41        }
42    }
43
44    pub fn base_url(mut self, base_url: &'a str) -> Self {
45        self.base_url = base_url;
46        self
47    }
48
49    pub fn custom_client(mut self, client: reqwest::Client) -> Self {
50        self.http_client = Some(client);
51        self
52    }
53
54    pub fn build(self) -> Result<Client, ClientBuilderError> {
55        let http_client = if let Some(http_client) = self.http_client {
56            http_client
57        } else {
58            reqwest::Client::builder().build()?
59        };
60
61        Ok(Client {
62            base_url: self.base_url.to_string(),
63            api_key: self.api_key.to_string(),
64            http_client,
65        })
66    }
67}
68
69#[derive(Clone)]
70pub struct Client {
71    base_url: String,
72    api_key: String,
73    http_client: reqwest::Client,
74}
75
76impl std::fmt::Debug for Client {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("Client")
79            .field("base_url", &self.base_url)
80            .field("http_client", &self.http_client)
81            .field("api_key", &"<REDACTED>")
82            .finish()
83    }
84}
85
86impl Client {
87    /// Create a new Moonshot client builder.
88    ///
89    /// # Example
90    /// ```
91    /// use rig::providers::moonshot::{ClientBuilder, self};
92    ///
93    /// // Initialize the Moonshot client
94    /// let moonshot = Client::builder("your-moonshot-api-key")
95    ///    .build()
96    /// ```
97    pub fn builder(api_key: &str) -> ClientBuilder<'_> {
98        ClientBuilder::new(api_key)
99    }
100
101    /// Create a new Moonshot client. For more control, use the `builder` method.
102    ///
103    /// # Panics
104    /// - If the reqwest client cannot be built (if the TLS backend cannot be initialized).
105    pub fn new(api_key: &str) -> Self {
106        Self::builder(api_key)
107            .build()
108            .expect("Moonshot client should build")
109    }
110
111    pub(crate) fn post(&self, path: &str) -> reqwest::RequestBuilder {
112        let url = format!("{}/{}", self.base_url, path).replace("//", "/");
113        self.http_client.post(url).bearer_auth(&self.api_key)
114    }
115}
116
117impl ProviderClient for Client {
118    /// Create a new Moonshot client from the `MOONSHOT_API_KEY` environment variable.
119    /// Panics if the environment variable is not set.
120    fn from_env() -> Self {
121        let api_key = std::env::var("MOONSHOT_API_KEY").expect("MOONSHOT_API_KEY not set");
122        Self::new(&api_key)
123    }
124
125    fn from_val(input: crate::client::ProviderValue) -> Self {
126        let crate::client::ProviderValue::Simple(api_key) = input else {
127            panic!("Incorrect provider value type")
128        };
129        Self::new(&api_key)
130    }
131}
132
133impl CompletionClient for Client {
134    type CompletionModel = CompletionModel;
135
136    /// Create a completion model with the given name.
137    ///
138    /// # Example
139    /// ```
140    /// use rig::providers::moonshot::{Client, self};
141    ///
142    /// // Initialize the Moonshot client
143    /// let moonshot = Client::new("your-moonshot-api-key");
144    ///
145    /// let completion_model = moonshot.completion_model(moonshot::MOONSHOT_CHAT);
146    /// ```
147    fn completion_model(&self, model: &str) -> CompletionModel {
148        CompletionModel::new(self.clone(), model)
149    }
150}
151
152impl_conversion_traits!(
153    AsEmbeddings,
154    AsTranscription,
155    AsImageGeneration,
156    AsAudioGeneration for Client
157);
158
159#[derive(Debug, Deserialize)]
160struct ApiErrorResponse {
161    error: MoonshotError,
162}
163
164#[derive(Debug, Deserialize)]
165struct MoonshotError {
166    message: String,
167}
168
169#[derive(Debug, Deserialize)]
170#[serde(untagged)]
171enum ApiResponse<T> {
172    Ok(T),
173    Err(ApiErrorResponse),
174}
175
176// ================================================================
177// Moonshot Completion API
178// ================================================================
179pub const MOONSHOT_CHAT: &str = "moonshot-v1-128k";
180
181#[derive(Clone)]
182pub struct CompletionModel {
183    client: Client,
184    pub model: String,
185}
186
187impl CompletionModel {
188    pub fn new(client: Client, model: &str) -> Self {
189        Self {
190            client,
191            model: model.to_string(),
192        }
193    }
194
195    fn create_completion_request(
196        &self,
197        completion_request: CompletionRequest,
198    ) -> Result<Value, CompletionError> {
199        // Build up the order of messages (context, chat_history)
200        let mut partial_history = vec![];
201        if let Some(docs) = completion_request.normalized_documents() {
202            partial_history.push(docs);
203        }
204        partial_history.extend(completion_request.chat_history);
205
206        // Initialize full history with preamble (or empty if non-existent)
207        let mut full_history: Vec<openai::Message> = completion_request
208            .preamble
209            .map_or_else(Vec::new, |preamble| {
210                vec![openai::Message::system(&preamble)]
211            });
212
213        // Convert and extend the rest of the history
214        full_history.extend(
215            partial_history
216                .into_iter()
217                .map(message::Message::try_into)
218                .collect::<Result<Vec<Vec<openai::Message>>, _>>()?
219                .into_iter()
220                .flatten()
221                .collect::<Vec<_>>(),
222        );
223
224        let request = if completion_request.tools.is_empty() {
225            json!({
226                "model": self.model,
227                "messages": full_history,
228                "temperature": completion_request.temperature,
229            })
230        } else {
231            json!({
232                "model": self.model,
233                "messages": full_history,
234                "temperature": completion_request.temperature,
235                "tools": completion_request.tools.into_iter().map(openai::ToolDefinition::from).collect::<Vec<_>>(),
236                "tool_choice": "auto",
237            })
238        };
239
240        let request = if let Some(params) = completion_request.additional_params {
241            json_utils::merge(request, params)
242        } else {
243            request
244        };
245
246        Ok(request)
247    }
248}
249
250impl completion::CompletionModel for CompletionModel {
251    type Response = openai::CompletionResponse;
252    type StreamingResponse = openai::StreamingCompletionResponse;
253
254    #[cfg_attr(feature = "worker", worker::send)]
255    async fn completion(
256        &self,
257        completion_request: CompletionRequest,
258    ) -> Result<completion::CompletionResponse<openai::CompletionResponse>, CompletionError> {
259        let request = self.create_completion_request(completion_request)?;
260
261        let response = self
262            .client
263            .post("/chat/completions")
264            .json(&request)
265            .send()
266            .await?;
267
268        if response.status().is_success() {
269            let t = response.text().await?;
270            tracing::debug!(target: "rig", "MoonShot completion error: {}", t);
271
272            match serde_json::from_str::<ApiResponse<openai::CompletionResponse>>(&t)? {
273                ApiResponse::Ok(response) => {
274                    tracing::info!(target: "rig",
275                        "MoonShot completion token usage: {:?}",
276                        response.usage.clone().map(|usage| format!("{usage}")).unwrap_or("N/A".to_string())
277                    );
278                    response.try_into()
279                }
280                ApiResponse::Err(err) => Err(CompletionError::ProviderError(err.error.message)),
281            }
282        } else {
283            Err(CompletionError::ProviderError(response.text().await?))
284        }
285    }
286
287    #[cfg_attr(feature = "worker", worker::send)]
288    async fn stream(
289        &self,
290        request: CompletionRequest,
291    ) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
292        let mut request = self.create_completion_request(request)?;
293
294        request = merge(
295            request,
296            json!({"stream": true, "stream_options": {"include_usage": true}}),
297        );
298
299        let builder = self.client.post("/chat/completions").json(&request);
300
301        send_compatible_streaming_request(builder).await
302    }
303}
rig/providers/moonshot.rs

rig/providers/
moonshot.rs
