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

rig/providers/
moonshot.rs
