1use async_trait::async_trait;
2use reqwest::Client as HttpClient;
3use serde_json::{Map, Value};
4
5use crate::config::TimeoutsConfig;
6use crate::config::constants::{env_vars, models};
7use crate::config::core::{AnthropicConfig, ModelConfig, PromptCachingConfig};
8use crate::config::models::{MiMoAuthMethod, detect_mimo_auth_method};
9use crate::llm::error_display;
10use crate::llm::provider::{LLMError, LLMProvider, LLMRequest, LLMResponse, LLMStream};
11
12use super::{
13 common::{
14 ensure_model, extract_prompt_cache_settings_default, impl_llm_client, override_base_url,
15 parse_json_response, parse_response_openai_format, resolve_model,
16 serialize_messages_openai_format, serialize_tools_openai_format,
17 spawn_openai_compatible_stream, validate_supported_models,
18 },
19 error_handling::handle_openai_http_error,
20 extract_reasoning_trace,
21};
22
23const PROVIDER_NAME: &str = "Xiaomi MiMo";
24const PROVIDER_KEY: &str = "mimo";
25
26pub struct MiMoProvider {
27 api_key: String,
28 http_client: HttpClient,
29 base_url: String,
30 model: String,
31 prompt_cache_enabled: bool,
32 model_behavior: Option<ModelConfig>,
33 auth_method: MiMoAuthMethod,
34}
35
36impl MiMoProvider {
37 pub fn new(api_key: String) -> Self {
38 let auth_method = detect_mimo_auth_method(&api_key, None);
39 Self::with_model_internal(
40 api_key,
41 models::mimo::DEFAULT_MODEL.to_string(),
42 None,
43 None,
44 TimeoutsConfig::default(),
45 None,
46 auth_method,
47 )
48 }
49
50 pub fn with_model(api_key: String, model: String) -> Self {
51 let auth_method = detect_mimo_auth_method(&api_key, None);
52 Self::with_model_internal(
53 api_key,
54 model,
55 None,
56 None,
57 TimeoutsConfig::default(),
58 None,
59 auth_method,
60 )
61 }
62
63 pub fn new_with_client(
64 api_key: String,
65 model: String,
66 http_client: reqwest::Client,
67 base_url: String,
68 _timeouts: TimeoutsConfig,
69 ) -> Self {
70 let auth_method = detect_mimo_auth_method(&api_key, Some(&base_url));
71 Self {
72 api_key,
73 http_client,
74 base_url,
75 model,
76 prompt_cache_enabled: false,
77 model_behavior: None,
78 auth_method,
79 }
80 }
81
82 pub fn from_config(
83 api_key: Option<String>,
84 model: Option<String>,
85 base_url: Option<String>,
86 prompt_cache: Option<PromptCachingConfig>,
87 timeouts: Option<TimeoutsConfig>,
88 _anthropic: Option<AnthropicConfig>,
89 model_behavior: Option<ModelConfig>,
90 ) -> Self {
91 let api_key_value = api_key.unwrap_or_default();
92 let auth_method = detect_mimo_auth_method(&api_key_value, base_url.as_deref());
93
94 Self::with_model_internal(
95 api_key_value,
96 resolve_model(model, models::mimo::DEFAULT_MODEL),
97 prompt_cache,
98 base_url,
99 timeouts.unwrap_or_default(),
100 model_behavior,
101 auth_method,
102 )
103 }
104
105 fn with_model_internal(
106 api_key: String,
107 model: String,
108 prompt_cache: Option<PromptCachingConfig>,
109 base_url: Option<String>,
110 timeouts: TimeoutsConfig,
111 model_behavior: Option<ModelConfig>,
112 auth_method: MiMoAuthMethod,
113 ) -> Self {
114 use crate::llm::http_client::HttpClientFactory;
115
116 let (prompt_cache_enabled, _) =
117 extract_prompt_cache_settings_default(prompt_cache, PROVIDER_KEY);
118
119 let default_base = auth_method.api_base();
120 let env_var = match auth_method {
121 MiMoAuthMethod::PayAsYouGo => env_vars::MIMO_BASE_URL,
122 MiMoAuthMethod::TokenPlan => env_vars::MIMO_TOKEN_PLAN_BASE_URL,
123 };
124
125 Self {
126 api_key,
127 http_client: HttpClientFactory::for_llm(&timeouts),
128 base_url: override_base_url(default_base, base_url, Some(env_var)),
129 model,
130 prompt_cache_enabled,
131 model_behavior,
132 auth_method,
133 }
134 }
135
136 #[must_use]
137 #[inline]
138 fn is_thinking_enabled(request: &LLMRequest) -> bool {
139 request
140 .reasoning_effort
141 .is_some_and(|e| e != crate::config::types::ReasoningEffortLevel::None)
142 }
143
144 fn convert_to_mimo_format(&self, request: &LLMRequest) -> Result<Value, LLMError> {
145 let mut payload = Map::with_capacity(12);
146
147 payload.insert("model".to_owned(), Value::String(request.model.clone()));
148
149 let mut messages = self.serialize_messages(request)?;
150
151 if let Some(system_prompt) = &request.system_prompt {
152 let trimmed = system_prompt.trim();
153 if !trimmed.is_empty() {
154 messages.insert(0, serde_json::json!({"role": "system", "content": trimmed}));
155 }
156 }
157
158 payload.insert("messages".to_owned(), Value::Array(messages));
159
160 if let Some(max_tokens) = request.max_tokens {
161 payload.insert(
162 "max_completion_tokens".to_owned(),
163 Value::Number(serde_json::Number::from(max_tokens as u64)),
164 );
165 }
166
167 let thinking_enabled = Self::is_thinking_enabled(request);
168
169 if !thinking_enabled {
170 if let Some(temperature) = request.temperature {
171 payload.insert(
172 "temperature".to_owned(),
173 Value::Number(super::common::float_to_json_number(temperature)?),
174 );
175 }
176
177 if let Some(top_p) = request.top_p {
178 payload.insert(
179 "top_p".to_owned(),
180 Value::Number(super::common::float_to_json_number(top_p)?),
181 );
182 }
183 }
184
185 if request.stream {
186 payload.insert("stream".to_string(), Value::Bool(true));
187 payload.insert(
188 "stream_options".to_string(),
189 serde_json::json!({"include_usage": true}),
190 );
191 }
192
193 if let Some(tools) = &request.tools
194 && let Some(serialized_tools) = serialize_tools_openai_format(tools)
195 {
196 payload.insert("tools".to_string(), Value::Array(serialized_tools));
197 }
198
199 if let Some(choice) = &request.tool_choice {
200 payload.insert(
201 "tool_choice".to_string(),
202 choice.to_provider_format(PROVIDER_KEY),
203 );
204 }
205
206 if let Some(effort) = request.reasoning_effort {
207 if effort == crate::config::types::ReasoningEffortLevel::None {
208 payload.insert(
209 "thinking".to_owned(),
210 serde_json::json!({"type": "disabled"}),
211 );
212 } else {
213 payload.insert(
214 "thinking".to_owned(),
215 serde_json::json!({"type": "enabled"}),
216 );
217 }
218 }
219
220 if let Some(meta) = &request.metadata
221 && let Some(user_id) = meta.get("user_id").and_then(|v| v.as_str())
222 {
223 payload.insert("user_id".to_owned(), Value::String(user_id.to_owned()));
224 }
225
226 Ok(Value::Object(payload))
227 }
228
229 async fn send_request(&self, payload: &Value) -> Result<reqwest::Response, LLMError> {
230 let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
231
232 let mut request = self.http_client.post(&url).json(payload);
233
234 request = match self.auth_method {
236 MiMoAuthMethod::PayAsYouGo => request.header("api-key", &self.api_key),
237 MiMoAuthMethod::TokenPlan => {
238 request.header("Authorization", format!("Bearer {}", &self.api_key))
239 }
240 };
241
242 request.send().await.map_err(|e| LLMError::Network {
243 message: error_display::format_llm_error(
244 PROVIDER_NAME,
245 &format!("network error: {}", e),
246 ),
247 metadata: None,
248 })
249 }
250
251 fn serialize_messages(&self, request: &LLMRequest) -> Result<Vec<Value>, LLMError> {
252 serialize_messages_openai_format(request, PROVIDER_KEY)
253 }
254
255 fn parse_response(&self, response_json: Value, model: String) -> Result<LLMResponse, LLMError> {
256 let reasoning_extractor = |message: &Value, choice: &Value| {
257 message
258 .get("reasoning_content")
259 .and_then(extract_reasoning_trace)
260 .or_else(|| {
261 choice
262 .get("reasoning_content")
263 .and_then(extract_reasoning_trace)
264 })
265 };
266
267 parse_response_openai_format(
268 response_json,
269 PROVIDER_NAME,
270 model,
271 self.prompt_cache_enabled,
272 Some(reasoning_extractor),
273 )
274 }
275}
276
277#[async_trait]
278impl LLMProvider for MiMoProvider {
279 fn name(&self) -> &str {
280 PROVIDER_KEY
281 }
282
283 fn supports_streaming(&self) -> bool {
284 true
285 }
286
287 fn supports_tools(&self, _model: &str) -> bool {
288 true
289 }
290
291 fn supports_structured_output(&self, _model: &str) -> bool {
292 true
293 }
294
295 fn supports_vision(&self, model: &str) -> bool {
296 model == models::mimo::MIMO_V2_5
297 }
298
299 fn supports_reasoning(&self, model: &str) -> bool {
300 let requested = if model.trim().is_empty() {
301 &self.model
302 } else {
303 model
304 };
305
306 self.model_behavior
307 .as_ref()
308 .and_then(|b| b.model_supports_reasoning)
309 .unwrap_or(false)
310 || requested == models::mimo::MIMO_V2_5_PRO
311 || requested == models::mimo::MIMO_V2_5
312 }
313
314 fn supports_reasoning_effort(&self, _model: &str) -> bool {
315 self.model_behavior
316 .as_ref()
317 .and_then(|b| b.model_supports_reasoning_effort)
318 .unwrap_or(false)
319 }
320
321 fn effective_context_size(&self, model: &str) -> usize {
322 let requested = if model.trim().is_empty() {
323 &self.model
324 } else {
325 model
326 };
327 match requested {
328 models::mimo::MIMO_V2_5_PRO | models::mimo::MIMO_V2_5 => 1_048_576,
329 _ => 128_000,
330 }
331 }
332
333 async fn generate(&self, mut request: LLMRequest) -> Result<LLMResponse, LLMError> {
334 let model = ensure_model(&mut request, &self.model);
335
336 let payload = self.convert_to_mimo_format(&request)?;
337 let response = self.send_request(&payload).await?;
338 let env_key = self.auth_method.env_key();
339 let response = handle_openai_http_error(response, PROVIDER_NAME, env_key).await?;
340
341 let response_json = parse_json_response(response, PROVIDER_NAME).await?;
342 self.parse_response(response_json, model)
343 }
344
345 async fn stream(&self, mut request: LLMRequest) -> Result<LLMStream, LLMError> {
346 ensure_model(&mut request, &self.model);
347 self.validate_request(&request)?;
348 request.stream = true;
349 let model = request.model.clone();
350
351 let payload = self.convert_to_mimo_format(&request)?;
352 let response = self.send_request(&payload).await?;
353 let env_key = self.auth_method.env_key();
354 let response = handle_openai_http_error(response, PROVIDER_NAME, env_key).await?;
355
356 Ok(spawn_openai_compatible_stream(
357 response,
358 PROVIDER_NAME,
359 model,
360 Some("reasoning_content"),
361 super::shared::OpenAiDeltaOrder::ReasoningFirst,
362 ))
363 }
364
365 fn supported_models(&self) -> Vec<String> {
366 let model_list = match self.auth_method {
367 MiMoAuthMethod::PayAsYouGo => models::mimo::PAYG_MODELS,
368 MiMoAuthMethod::TokenPlan => models::mimo::TOKEN_PLAN_MODELS,
369 };
370 model_list.iter().map(|model| model.to_string()).collect()
371 }
372
373 fn validate_request(&self, request: &LLMRequest) -> Result<(), LLMError> {
374 let model_list = match self.auth_method {
375 MiMoAuthMethod::PayAsYouGo => models::mimo::PAYG_MODELS,
376 MiMoAuthMethod::TokenPlan => models::mimo::TOKEN_PLAN_MODELS,
377 };
378 validate_supported_models(request, PROVIDER_NAME, PROVIDER_KEY, model_list)
379 }
380}
381
382impl_llm_client!(MiMoProvider);