microclaw_engine/llm/
stream.rs

1use super::*;
2
3#[derive(Default)]
4pub(crate) struct StreamToolUseBlock {
5    pub(crate) id: String,
6    pub(crate) name: String,
7    pub(crate) input_json: String,
8    pub(crate) thought_signature: Option<String>,
9}
10
11pub(crate) fn usage_from_json(v: &serde_json::Value) -> Option<Usage> {
12    let input = v
13        .get("input_tokens")
14        .and_then(json_u64)
15        .or_else(|| v.get("prompt_tokens").and_then(json_u64))?;
16    let output = v
17        .get("output_tokens")
18        .and_then(json_u64)
19        .or_else(|| v.get("completion_tokens").and_then(json_u64))
20        .unwrap_or(0);
21    // Cache accounting: OpenAI-compatible providers nest it under
22    // `prompt_tokens_details`/`input_tokens_details`, Anthropic-style providers
23    // report it as top-level fields.
24    let cache_read = v
25        .get("cache_read_input_tokens")
26        .and_then(json_u64)
27        .or_else(|| {
28            v.get("prompt_tokens_details")
29                .and_then(|d| d.get("cached_tokens"))
30                .and_then(json_u64)
31        })
32        .or_else(|| {
33            v.get("input_tokens_details")
34                .and_then(|d| d.get("cached_tokens"))
35                .and_then(json_u64)
36        })
37        .unwrap_or(0);
38    let cache_creation = v
39        .get("cache_creation_input_tokens")
40        .and_then(json_u64)
41        .unwrap_or(0);
42    Some(Usage {
43        input_tokens: u32::try_from(input).unwrap_or(u32::MAX),
44        output_tokens: u32::try_from(output).unwrap_or(u32::MAX),
45        cache_read_input_tokens: u32::try_from(cache_read).unwrap_or(u32::MAX),
46        cache_creation_input_tokens: u32::try_from(cache_creation).unwrap_or(u32::MAX),
47    })
48}
49
50pub(crate) fn json_u64(v: &serde_json::Value) -> Option<u64> {
51    match v {
52        serde_json::Value::Number(n) => n.as_u64(),
53        serde_json::Value::String(s) => s.parse::<u64>().ok(),
54        _ => None,
55    }
56}
57
58// OpenAI-compatible streaming usage is cumulative (running total), not per-chunk delta.
59// We therefore keep the max seen values instead of summing to avoid double counting.
60pub(crate) fn merge_usage_max(slot: &mut Option<Usage>, incoming: Usage) {
61    match slot {
62        Some(current) => {
63            current.input_tokens = current.input_tokens.max(incoming.input_tokens);
64            current.output_tokens = current.output_tokens.max(incoming.output_tokens);
65        }
66        None => {
67            *slot = Some(incoming);
68        }
69    }
70}
71
72pub(crate) fn normalize_stop_reason(reason: Option<String>) -> Option<String> {
73    match reason.as_deref() {
74        Some("tool_use") | Some("tool_calls") => Some("tool_use".into()),
75        Some("max_tokens") | Some("length") => Some("max_tokens".into()),
76        Some("stop") | Some("end_turn") | None => Some("end_turn".into()),
77        Some(other) => Some(other.to_string()),
78    }
79}
80
81pub(crate) fn parse_tool_input(input_json: &str) -> serde_json::Value {
82    let trimmed = input_json.trim();
83    if trimmed.is_empty() {
84        return json!({});
85    }
86    serde_json::from_str(trimmed).unwrap_or_else(|_| json!({}))
87}
88
89pub(crate) fn normalize_tool_input_for_request(input: &serde_json::Value) -> serde_json::Value {
90    match input {
91        serde_json::Value::Object(_) => input.clone(),
92        serde_json::Value::String(text) => {
93            let trimmed = text.trim();
94            if trimmed.is_empty() {
95                return json!({});
96            }
97            match serde_json::from_str::<serde_json::Value>(trimmed) {
98                Ok(serde_json::Value::Object(map)) => serde_json::Value::Object(map),
99                _ => json!({}),
100            }
101        }
102        _ => json!({}),
103    }
104}
105
106pub(crate) fn minimax_tool_wrapper_regex() -> &'static Regex {
107    static RE: OnceLock<Regex> = OnceLock::new();
108    RE.get_or_init(|| {
109        Regex::new(r"</?(?:minimax:tool_call|invoke|parameter)>")
110            .expect("MiniMax tool wrapper regex must compile")
111    })
112}
113
114pub(crate) fn strip_minimax_tool_wrappers(text: &str) -> String {
115    minimax_tool_wrapper_regex()
116        .replace_all(text, " ")
117        .into_owned()
118}
119
120pub(crate) fn parse_raw_tool_use_block(
121    input: &str,
122    call_number: usize,
123) -> Option<(StreamToolUseBlock, &str)> {
124    let rest = input.trim_start();
125    let prefix = "[tool_use:";
126    if !rest.starts_with(prefix) {
127        return None;
128    }
129
130    let mut cursor = prefix.len();
131    let after_prefix = &rest[cursor..];
132    let name_and_args = after_prefix.trim_start();
133    cursor += after_prefix.len().saturating_sub(name_and_args.len());
134
135    let open_paren_rel = name_and_args.find('(')?;
136    let name = name_and_args[..open_paren_rel].trim();
137    if name.is_empty() {
138        return None;
139    }
140    cursor += open_paren_rel + 1;
141
142    let mut depth = 1usize;
143    let mut in_string = false;
144    let mut escaping = false;
145    let mut close_paren_at: Option<usize> = None;
146
147    for (offset, ch) in rest[cursor..].char_indices() {
148        if in_string {
149            if escaping {
150                escaping = false;
151                continue;
152            }
153            match ch {
154                '\\' => escaping = true,
155                '"' => in_string = false,
156                _ => {}
157            }
158            continue;
159        }
160
161        match ch {
162            '"' => in_string = true,
163            '(' => depth += 1,
164            ')' => {
165                depth = depth.saturating_sub(1);
166                if depth == 0 {
167                    close_paren_at = Some(cursor + offset);
168                    break;
169                }
170            }
171            _ => {}
172        }
173    }
174
175    let close_paren_at = close_paren_at?;
176    let args = rest[cursor..close_paren_at].trim();
177    let mut tail = &rest[close_paren_at + 1..];
178    tail = tail.trim_start();
179    if !tail.starts_with(']') {
180        return None;
181    }
182
183    Some((
184        StreamToolUseBlock {
185            id: format!("raw_tool_call_{call_number}"),
186            name: name.to_string(),
187            input_json: if args.is_empty() {
188                "{}".to_string()
189            } else {
190                args.to_string()
191            },
192            thought_signature: None,
193        },
194        &tail[1..],
195    ))
196}
197
198pub(crate) fn extract_raw_tool_use_blocks(text: &str) -> Option<Vec<StreamToolUseBlock>> {
199    let normalized = strip_minimax_tool_wrappers(text);
200    if !normalized.contains("[tool_use:") {
201        return None;
202    }
203
204    let mut calls = Vec::new();
205    let mut rest = normalized.as_str();
206    loop {
207        let trimmed = rest.trim();
208        if trimmed.is_empty() {
209            break;
210        }
211        let (call, tail) = parse_raw_tool_use_block(trimmed, calls.len() + 1)?;
212        calls.push(call);
213        rest = tail;
214    }
215
216    if calls.is_empty() {
217        None
218    } else {
219        Some(calls)
220    }
221}
222
223pub(crate) fn has_tool_use_block(content: &[ResponseContentBlock]) -> bool {
224    content
225        .iter()
226        .any(|b| matches!(b, ResponseContentBlock::ToolUse { .. }))
227}
228
229pub(crate) fn combine_visible_and_reasoning_text(visible: &str, reasoning: &str) -> String {
230    let visible = visible.trim();
231    let reasoning = reasoning.trim();
232    if reasoning.is_empty() {
233        return visible.to_string();
234    }
235    if visible.is_empty() {
236        return format!("<thought>\n{}\n</thought>", reasoning);
237    }
238    format!("<thought>\n{}\n</thought>\n\n{}", reasoning, visible)
239}
240
241pub(crate) fn combine_response_text_for_display(
242    visible: &str,
243    reasoning: &str,
244    show_thinking: bool,
245) -> String {
246    if show_thinking {
247        combine_visible_and_reasoning_text(visible, reasoning)
248    } else {
249        visible.trim().to_string()
250    }
251}
252
253pub(crate) fn build_stream_response(
254    ordered_indexes: Vec<usize>,
255    text_blocks: std::collections::HashMap<usize, String>,
256    tool_blocks: std::collections::HashMap<usize, StreamToolUseBlock>,
257    stop_reason: Option<String>,
258    usage: Option<Usage>,
259) -> MessagesResponse {
260    let mut content = Vec::new();
261    for index in ordered_indexes {
262        if let Some(text) = text_blocks.get(&index) {
263            if !text.is_empty() {
264                content.push(ResponseContentBlock::Text { text: text.clone() });
265            }
266        }
267        if let Some(tool) = tool_blocks.get(&index) {
268            content.push(ResponseContentBlock::ToolUse {
269                id: sanitize_tool_id(&tool.id),
270                name: tool.name.clone(),
271                input: parse_tool_input(&tool.input_json),
272                thought_signature: tool.thought_signature.clone(),
273            });
274        }
275    }
276
277    if content.is_empty() {
278        content.push(ResponseContentBlock::Text {
279            text: String::new(),
280        });
281    }
282
283    let mut normalized_stop_reason = normalize_stop_reason(stop_reason);
284    if !tool_blocks.is_empty() {
285        normalized_stop_reason = Some("tool_use".to_string());
286    } else if normalized_stop_reason.as_deref() == Some("tool_use") && !has_tool_use_block(&content)
287    {
288        warn!("Downgrading stop_reason=tool_use to end_turn because no tool_calls were parsed");
289        normalized_stop_reason = Some("end_turn".into());
290    }
291
292    MessagesResponse {
293        content,
294        stop_reason: normalized_stop_reason,
295        usage,
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302    #[allow(unused_imports)]
303    use crate::llm::test_prelude::*;
304
305    #[test]
306    fn test_normalize_stop_reason_stream_variants() {
307        assert_eq!(
308            normalize_stop_reason(Some("tool_calls".into())).as_deref(),
309            Some("tool_use")
310        );
311        assert_eq!(
312            normalize_stop_reason(Some("length".into())).as_deref(),
313            Some("max_tokens")
314        );
315        assert_eq!(
316            normalize_stop_reason(Some("stop".into())).as_deref(),
317            Some("end_turn")
318        );
319    }
320
321    #[test]
322    fn test_build_stream_response_tool_calls_without_blocks_downgrades_to_end_turn() {
323        let resp = build_stream_response(
324            vec![],
325            std::collections::HashMap::new(),
326            std::collections::HashMap::new(),
327            Some("tool_calls".into()),
328            None,
329        );
330        assert_eq!(resp.stop_reason.as_deref(), Some("end_turn"));
331        assert!(!resp
332            .content
333            .iter()
334            .any(|b| matches!(b, ResponseContentBlock::ToolUse { .. })));
335    }
336
337    #[test]
338    fn test_extract_raw_tool_use_blocks_from_minimax_wrappers() {
339        let text = "<minimax:tool_call>\n<invoke>\n<parameter>\n[tool_use: bash({\"command\":\"uptime\"})]\n</parameter>\n</invoke>\n</minimax:tool_call>";
340        let calls = extract_raw_tool_use_blocks(text).expect("should parse raw tool calls");
341        assert_eq!(calls.len(), 1);
342        assert_eq!(calls[0].name, "bash");
343        assert_eq!(calls[0].input_json, "{\"command\":\"uptime\"}");
344    }
345
346    #[test]
347    fn test_usage_from_json_supports_openai_prompt_completion_tokens() {
348        let v = json!({
349            "prompt_tokens": 12,
350            "completion_tokens": 34
351        });
352        let usage = usage_from_json(&v).expect("usage should parse");
353        assert_eq!(usage.input_tokens, 12);
354        assert_eq!(usage.output_tokens, 34);
355    }
356
357    #[test]
358    fn test_usage_from_json_supports_numeric_strings() {
359        let v = json!({
360            "input_tokens": "56",
361            "output_tokens": "78"
362        });
363        let usage = usage_from_json(&v).expect("usage should parse");
364        assert_eq!(usage.input_tokens, 56);
365        assert_eq!(usage.output_tokens, 78);
366    }
367
368    #[test]
369    fn test_usage_from_json_reads_openai_prompt_cache_details() {
370        let v = json!({
371            "prompt_tokens": 100,
372            "completion_tokens": 10,
373            "prompt_tokens_details": { "cached_tokens": 64 }
374        });
375        let usage = usage_from_json(&v).expect("usage should parse");
376        assert_eq!(usage.cache_read_input_tokens, 64);
377        assert_eq!(usage.cache_creation_input_tokens, 0);
378    }
379
380    #[test]
381    fn test_usage_from_json_reads_responses_api_input_token_details() {
382        let v = json!({
383            "input_tokens": 100,
384            "output_tokens": 10,
385            "input_tokens_details": { "cached_tokens": 32 }
386        });
387        let usage = usage_from_json(&v).expect("usage should parse");
388        assert_eq!(usage.cache_read_input_tokens, 32);
389    }
390
391    #[test]
392    fn test_usage_from_json_reads_anthropic_style_cache_fields() {
393        let v = json!({
394            "input_tokens": 100,
395            "output_tokens": 10,
396            "cache_read_input_tokens": 70,
397            "cache_creation_input_tokens": 30
398        });
399        let usage = usage_from_json(&v).expect("usage should parse");
400        assert_eq!(usage.cache_read_input_tokens, 70);
401        assert_eq!(usage.cache_creation_input_tokens, 30);
402    }
403
404    #[test]
405    fn test_usage_from_json_defaults_cache_fields_to_zero() {
406        let v = json!({ "prompt_tokens": 5, "completion_tokens": 6 });
407        let usage = usage_from_json(&v).expect("usage should parse");
408        assert_eq!(usage.cache_read_input_tokens, 0);
409        assert_eq!(usage.cache_creation_input_tokens, 0);
410    }
411
412    #[test]
413    fn test_build_stream_response_tool_json_parsing() {
414        let mut tool_blocks = std::collections::HashMap::new();
415        tool_blocks.insert(
416            0,
417            StreamToolUseBlock {
418                id: "call_1".into(),
419                name: "bash".into(),
420                input_json: r#"{"command":"ls","cwd":"/tmp"}"#.into(),
421                thought_signature: None,
422            },
423        );
424        let resp = build_stream_response(
425            vec![0],
426            std::collections::HashMap::new(),
427            tool_blocks,
428            Some("tool_use".into()),
429            None,
430        );
431        assert_eq!(resp.stop_reason.as_deref(), Some("tool_use"));
432        match &resp.content[0] {
433            ResponseContentBlock::ToolUse {
434                id,
435                name,
436                input,
437                thought_signature,
438            } => {
439                assert_eq!(id, "call_1");
440                assert_eq!(name, "bash");
441                assert_eq!(input["command"], "ls");
442                assert_eq!(input["cwd"], "/tmp");
443                assert!(thought_signature.is_none());
444            }
445            _ => panic!("Expected ToolUse"),
446        }
447    }
448
449    // -----------------------------------------------------------------------
450    // create_provider
451    // -----------------------------------------------------------------------
452}
microclaw_engine/llm/stream.rs

microclaw_engine/llm/
stream.rs
