Skip to main content

va_ai_api_bridge/providers/
mimo.rs

1use std::collections::HashMap;
2
3use serde_json::{json, Value};
4
5use super::deepseek::{
6    collect_reasoning_from_anthropic_input, collect_reasoning_from_gemini_input,
7    collect_reasoning_from_responses_input, collect_tool_outputs_from_chat_request,
8    collect_tool_outputs_from_responses_input, inject_reasoning_content, repair_tool_call_history,
9    strip_anthropic_reasoning_content_blocks, RequestReasoning,
10};
11use super::ProviderRequestSource;
12
13const MISSING_REASONING_CONTENT_FALLBACK: &str =
14    "Previous MiMo reasoning_content is unavailable from the local bridge.";
15
16#[derive(Debug, Clone, Default)]
17pub struct MimoBridgeAdapter;
18
19impl MimoBridgeAdapter {
20    pub fn prepare_chat_request(
21        &mut self,
22        source: ProviderRequestSource,
23        original_request: &Value,
24        chat_request: &mut Value,
25    ) {
26        if source == ProviderRequestSource::AnthropicMessages {
27            strip_anthropic_reasoning_content_blocks(chat_request);
28        }
29
30        let mut tool_outputs = HashMap::new();
31        collect_tool_outputs_from_responses_input(original_request, &mut tool_outputs);
32        collect_tool_outputs_from_chat_request(chat_request, &mut tool_outputs);
33        repair_tool_call_history(&tool_outputs, chat_request);
34
35        let mut reasoning = RequestReasoning::default();
36        match source {
37            ProviderRequestSource::OpenAiResponses => {
38                collect_reasoning_from_responses_input(&mut reasoning, original_request);
39            }
40            ProviderRequestSource::AnthropicMessages => {
41                collect_reasoning_from_anthropic_input(&mut reasoning, original_request);
42            }
43            ProviderRequestSource::GeminiGenerateContent => {
44                collect_reasoning_from_gemini_input(&mut reasoning, original_request);
45            }
46            ProviderRequestSource::OpenAiChat => {}
47        }
48        inject_reasoning_content(&reasoning, chat_request, MISSING_REASONING_CONTENT_FALLBACK);
49
50        if let Some(object) = chat_request.as_object_mut() {
51            object.insert("thinking".to_string(), json!({ "type": "enabled" }));
52        }
53    }
54
55    pub fn normalize_chat_response(&mut self, response: &mut Value) {
56        normalize_null_tool_calls(response);
57    }
58}
59
60fn normalize_null_tool_calls(value: &mut Value) {
61    let Some(choices) = value.get_mut("choices").and_then(Value::as_array_mut) else {
62        return;
63    };
64
65    for choice in choices {
66        normalize_message_tool_calls(choice, "message");
67        normalize_message_tool_calls(choice, "delta");
68    }
69}
70
71fn normalize_message_tool_calls(choice: &mut Value, key: &str) {
72    let Some(message) = choice.get_mut(key).and_then(Value::as_object_mut) else {
73        return;
74    };
75    if matches!(message.get("tool_calls"), Some(Value::Null)) {
76        message.insert("tool_calls".to_string(), Value::Array(Vec::new()));
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use crate::{DecodeState, OpenAiChatTranslator, WireTranslator};
83    use serde_json::json;
84
85    use super::super::reasoning_blob::encode_reasoning_content;
86    use super::super::ProviderRequestSource;
87    use super::MimoBridgeAdapter;
88
89    #[test]
90    fn enables_mimo_thinking_and_replays_responses_reasoning_content() {
91        let original_request = json!({
92            "input": [
93                {
94                    "type": "reasoning",
95                    "encrypted_content": encode_reasoning_content("Call pwd, then answer.")
96                },
97                {
98                    "type": "function_call",
99                    "call_id": "call_pwd",
100                    "name": "exec_command",
101                    "arguments": "{\"cmd\":\"pwd\"}"
102                },
103                {
104                    "type": "function_call_output",
105                    "call_id": "call_pwd",
106                    "output": "/tmp/project"
107                }
108            ]
109        });
110        let mut chat_request = json!({
111            "model": "mimo-v2.5-pro",
112            "messages": [{
113                "role": "assistant",
114                "content": null,
115                "tool_calls": [{
116                    "id": "call_pwd",
117                    "type": "function",
118                    "function": { "name": "exec_command", "arguments": "{\"cmd\":\"pwd\"}" }
119                }]
120            }, {
121                "role": "tool",
122                "tool_call_id": "call_pwd",
123                "content": "/tmp/project"
124            }]
125        });
126        let mut adapter = MimoBridgeAdapter;
127
128        adapter.prepare_chat_request(
129            ProviderRequestSource::OpenAiResponses,
130            &original_request,
131            &mut chat_request,
132        );
133
134        assert_eq!(chat_request["thinking"], json!({ "type": "enabled" }));
135        assert_eq!(
136            chat_request["messages"][0]["reasoning_content"],
137            "Call pwd, then answer."
138        );
139    }
140
141    #[test]
142    fn fills_missing_mimo_reasoning_content_for_tool_history() {
143        let mut chat_request = json!({
144            "model": "mimo-v2.5-pro",
145            "messages": [{
146                "role": "assistant",
147                "content": null,
148                "tool_calls": [{
149                    "id": "call_old",
150                    "type": "function",
151                    "function": { "name": "exec_command", "arguments": "{\"cmd\":\"pwd\"}" }
152                }]
153            }]
154        });
155        let mut adapter = MimoBridgeAdapter;
156
157        adapter.prepare_chat_request(
158            ProviderRequestSource::OpenAiChat,
159            &chat_request.clone(),
160            &mut chat_request,
161        );
162
163        assert_eq!(
164            chat_request["messages"][0]["reasoning_content"],
165            super::MISSING_REASONING_CONTENT_FALLBACK
166        );
167    }
168
169    #[test]
170    fn normalizes_mimo_null_tool_calls_for_chat_completion_response() {
171        let mut response = json!({
172            "id": "chatcmpl_mimo",
173            "model": "mimo-v2.5-pro",
174            "choices": [{
175                "index": 0,
176                "finish_reason": "stop",
177                "message": {
178                    "role": "assistant",
179                    "content": "OK",
180                    "tool_calls": null,
181                    "reasoning_content": "Answer briefly."
182                }
183            }]
184        });
185        let mut adapter = MimoBridgeAdapter;
186
187        adapter.normalize_chat_response(&mut response);
188
189        assert_eq!(response["choices"][0]["message"]["tool_calls"], json!([]));
190        OpenAiChatTranslator
191            .decode_response(response)
192            .expect("normalized MiMo response decodes");
193    }
194
195    #[test]
196    fn normalizes_mimo_null_tool_calls_for_chat_stream_chunk() {
197        let mut chunk = json!({
198            "id": "chatcmpl_mimo",
199            "model": "mimo-v2.5-pro",
200            "choices": [{
201                "index": 0,
202                "finish_reason": null,
203                "delta": {
204                    "role": "assistant",
205                    "content": "",
206                    "tool_calls": null,
207                    "reasoning_content": null
208                }
209            }]
210        });
211        let mut adapter = MimoBridgeAdapter;
212        let mut state = DecodeState::default();
213
214        adapter.normalize_chat_response(&mut chunk);
215
216        assert_eq!(chunk["choices"][0]["delta"]["tool_calls"], json!([]));
217        OpenAiChatTranslator
218            .decode_stream_chunk(chunk, &mut state)
219            .expect("normalized MiMo stream chunk decodes");
220    }
221}