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

rig/providers/
moonshot.rs
