Skip to main content

vtcode_core/llm/providers/
mimo.rs

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