vtcode_llm/providers/
mimo.rs

1use reqwest::RequestBuilder;
2use serde_json::{Map, Value};
3use vtcode_config::constants::{env_vars, models, urls};
4use vtcode_config::models::{MiMoAuthMethod, detect_mimo_auth_method};
5
6use super::common::override_base_url;
7use super::extract_reasoning_trace;
8use super::openai_compat::{OpenAiCompatCore, OpenAiCompatSpec, impl_openai_compat_provider};
9use crate::provider::{LLMError, LLMRequest};
10
11pub struct MimoSpec;
12
13fn mimo_reasoning(message: &Value, choice: &Value) -> Option<String> {
14    message
15        .get("reasoning_content")
16        .and_then(extract_reasoning_trace)
17        .or_else(|| choice.get("reasoning_content").and_then(extract_reasoning_trace))
18}
19
20/// Auth method is fully derivable from the key prefix and base URL, so it is
21/// recomputed where needed instead of being stored on the provider.
22fn auth_method(core: &OpenAiCompatCore<MimoSpec>) -> MiMoAuthMethod {
23    detect_mimo_auth_method(&core.api_key, Some(&core.base_url))
24}
25
26impl OpenAiCompatSpec for MimoSpec {
27    const NAME: &'static str = "Xiaomi MiMo";
28    const KEY: &'static str = "mimo";
29    const API_KEY_ENV: &'static str = "MIMO_API_KEY";
30    const DEFAULT_MODEL: &'static str = models::mimo::DEFAULT_MODEL;
31    const DEFAULT_BASE_URL: &'static str = urls::MIMO_API_BASE;
32    const BASE_URL_ENV: Option<&'static str> = Some(env_vars::MIMO_BASE_URL);
33    const LISTED_MODELS: &'static [&'static str] = models::mimo::PAYG_MODELS;
34    const VALIDATION_ALLOWLIST: Option<&'static [&'static str]> = None;
35
36    const MAX_TOKENS_KEY: &'static str = "max_completion_tokens";
37    const STREAM_OPTIONS_INCLUDE_USAGE: bool = true;
38    const INCLUDE_USER_ID: bool = true;
39    const RESPONSE_REASONING_EXTRACTOR: Option<super::openai_compat::ReasoningExtractor> = Some(mimo_reasoning);
40
41    fn resolve_base_url(api_key: &str, base_url: Option<String>) -> String {
42        let auth = detect_mimo_auth_method(api_key, base_url.as_deref());
43        let env_var = match auth {
44            MiMoAuthMethod::PayAsYouGo | MiMoAuthMethod::Unknown => env_vars::MIMO_BASE_URL,
45            MiMoAuthMethod::TokenPlan => env_vars::MIMO_TOKEN_PLAN_BASE_URL,
46        };
47        override_base_url(auth.api_base(), base_url, Some(env_var))
48    }
49
50    fn response_cache_metrics(core: &OpenAiCompatCore<Self>) -> bool {
51        core.prompt_cache_enabled
52    }
53
54    fn stream_cache_metrics(_core: &OpenAiCompatCore<Self>) -> bool {
55        true
56    }
57
58    fn insert_reasoning(
59        _core: &OpenAiCompatCore<Self>,
60        request: &LLMRequest,
61        payload: &mut Map<String, Value>,
62    ) -> Result<(), LLMError> {
63        if let Some(effort) = request.reasoning_effort {
64            let kind = if effort == vtcode_config::types::ReasoningEffortLevel::None {
65                "disabled"
66            } else {
67                "enabled"
68            };
69            payload.insert("thinking".to_owned(), serde_json::json!({"type": kind}));
70        }
71        Ok(())
72    }
73
74    fn apply_auth(core: &OpenAiCompatCore<Self>, builder: RequestBuilder) -> RequestBuilder {
75        match auth_method(core) {
76            MiMoAuthMethod::PayAsYouGo | MiMoAuthMethod::Unknown => builder.header("api-key", &core.api_key),
77            MiMoAuthMethod::TokenPlan => builder.bearer_auth(&core.api_key),
78        }
79    }
80
81    fn api_key_env(core: &OpenAiCompatCore<Self>) -> &'static str {
82        auth_method(core).env_key()
83    }
84
85    fn listed_models(core: &OpenAiCompatCore<Self>) -> &'static [&'static str] {
86        match auth_method(core) {
87            MiMoAuthMethod::PayAsYouGo | MiMoAuthMethod::Unknown => models::mimo::PAYG_MODELS,
88            MiMoAuthMethod::TokenPlan => models::mimo::TOKEN_PLAN_MODELS,
89        }
90    }
91
92    fn validate(core: &OpenAiCompatCore<Self>, request: &LLMRequest) -> Result<(), LLMError> {
93        super::common::validate_supported_models(request, Self::NAME, Self::KEY, Self::listed_models(core))
94    }
95}
96
97impl_openai_compat_provider!(MiMoProvider, MimoSpec, {
98    fn supports_streaming(&self) -> bool {
99        true
100    }
101
102    fn supports_structured_output(&self, _model: &str) -> bool {
103        true
104    }
105
106    fn supports_vision(&self, model: &str) -> bool {
107        model == models::mimo::MIMO_V2_6_PRO
108            || model == models::mimo::MIMO_V2_6_FLASH
109            || model == models::mimo::MIMO_V2_6_PRO_ULTRASPEED
110    }
111
112    fn supports_reasoning(&self, model: &str) -> bool {
113        let requested = if model.trim().is_empty() {
114            &self.core.model
115        } else {
116            model
117        };
118
119        self.core
120            .model_behavior
121            .as_ref()
122            .and_then(|b| b.model_supports_reasoning)
123            .unwrap_or(false)
124            || requested == models::mimo::MIMO_V2_6_PRO
125            || requested == models::mimo::MIMO_V2_6_FLASH
126            || requested == models::mimo::MIMO_V2_6_PRO_ULTRASPEED
127    }
128
129    fn supports_reasoning_effort(&self, _model: &str) -> bool {
130        self.core
131            .model_behavior
132            .as_ref()
133            .and_then(|b| b.model_supports_reasoning_effort)
134            .unwrap_or(false)
135    }
136
137    fn effective_context_size(&self, model: &str) -> usize {
138        let requested = if model.trim().is_empty() {
139            &self.core.model
140        } else {
141            model
142        };
143        match requested {
144            models::mimo::MIMO_V2_6_PRO | models::mimo::MIMO_V2_6_FLASH | models::mimo::MIMO_V2_6_PRO_ULTRASPEED => {
145                1_048_576
146            }
147            _ => 128_000,
148        }
149    }
150});
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::provider::{LLMProvider, Message, ToolChoice};
156    use std::sync::Arc;
157    use vtcode_config::types::ReasoningEffortLevel;
158
159    fn provider() -> MiMoProvider {
160        MiMoProvider::from_config(
161            Some("sk-test-key".to_string()),
162            Some("mimo-v2.6-pro".to_string()),
163            Some("https://example.test/v1".to_string()),
164            None,
165            None,
166            None,
167            None,
168        )
169    }
170
171    fn base_request() -> LLMRequest {
172        LLMRequest {
173            messages: vec![Message::user("hello".to_string())].into(),
174            system_prompt: Some(Arc::from("system guidance")),
175            model: "mimo-v2.6-pro".to_string(),
176            max_tokens: Some(512),
177            temperature: Some(0.5),
178            top_p: Some(0.25),
179            stream: true,
180            tool_choice: Some(ToolChoice::Auto),
181            metadata: Some(serde_json::json!({"user_id": "user-42"})),
182            ..Default::default()
183        }
184    }
185
186    #[test]
187    fn golden_payload_basic_shape() {
188        let payload = provider().core.convert_request(&base_request()).unwrap();
189
190        assert_eq!(payload["model"], "mimo-v2.6-pro");
191        let messages = payload["messages"].as_array().unwrap();
192        assert_eq!(messages.len(), 2);
193        assert_eq!(messages[0]["role"], "system");
194        assert_eq!(messages[0]["content"], "system guidance");
195        assert_eq!(payload["max_completion_tokens"], 512);
196        assert!(payload.get("max_tokens").is_none());
197        assert_eq!(payload["temperature"], 0.5);
198        assert_eq!(payload["top_p"], 0.25);
199        assert_eq!(payload["stream"], true);
200        assert_eq!(payload["stream_options"]["include_usage"], true);
201        assert_eq!(payload["tool_choice"], "auto");
202        assert_eq!(payload["user_id"], "user-42");
203        assert!(payload.get("thinking").is_none());
204    }
205
206    #[test]
207    fn golden_payload_thinking_object_and_sampling_suppression() {
208        let mut request = base_request();
209        request.reasoning_effort = Some(ReasoningEffortLevel::High);
210        let payload = provider().core.convert_request(&request).unwrap();
211        assert_eq!(payload["thinking"]["type"], "enabled");
212        assert!(payload.get("temperature").is_none());
213        assert!(payload.get("top_p").is_none());
214
215        let mut request = base_request();
216        request.reasoning_effort = Some(ReasoningEffortLevel::None);
217        let payload = provider().core.convert_request(&request).unwrap();
218        assert_eq!(payload["thinking"]["type"], "disabled");
219        assert_eq!(payload["temperature"], 0.5);
220    }
221
222    #[test]
223    fn auth_method_drives_supported_models() {
224        let payg = MiMoProvider::from_config(Some("sk-test-key".to_string()), None, None, None, None, None, None);
225        assert_eq!(
226            payg.supported_models(),
227            models::mimo::PAYG_MODELS.iter().map(|m| m.to_string()).collect::<Vec<_>>()
228        );
229
230        let token_plan = MiMoProvider::from_config(Some("tp-test-key".to_string()), None, None, None, None, None, None);
231        assert_eq!(
232            token_plan.supported_models(),
233            models::mimo::TOKEN_PLAN_MODELS
234                .iter()
235                .map(|m| m.to_string())
236                .collect::<Vec<_>>()
237        );
238    }
239}
vtcode_llm/providers/mimo.rs

vtcode_llm/providers/
mimo.rs
