edgequake_llm/
stream_tools.rs

1//! Streaming tool-call assembly helpers used by EdgeCrab's provider call path.
2
3use std::collections::BTreeMap;
4
5use crate::error::{LlmError, Result};
6use crate::traits::{FunctionCall, ToolCall};
7
8/// Partial tool call accumulated from [`crate::traits::StreamChunk::ToolCallDelta`] events.
9#[derive(Debug, Clone, Default)]
10pub struct PartialStreamToolCall {
11    pub id: Option<String>,
12    pub function_name: Option<String>,
13    pub arguments: String,
14    pub thought_signature: Option<String>,
15}
16
17/// Options for [`finalize_streamed_tool_calls_with_repair`].
18#[derive(Debug, Clone, Copy, Default)]
19pub struct FinalizeStreamToolCallsOptions {
20    /// When true, empty / null-like argument payloads become `{}`.
21    pub empty_args_fallback: bool,
22}
23
24/// Finalize indexed partial tool calls into wire-ready [`ToolCall`] values.
25///
26/// `repair` is an optional hook (EdgeCrab argument pipeline) applied to each
27/// argument string before the call is emitted.
28pub fn finalize_streamed_tool_calls_with_repair(
29    partials: BTreeMap<usize, PartialStreamToolCall>,
30    options: FinalizeStreamToolCallsOptions,
31    repair: Option<fn(&str) -> String>,
32) -> Result<Vec<ToolCall>> {
33    let mut out = Vec::with_capacity(partials.len());
34    for (index, partial) in partials {
35        let name = partial.function_name.unwrap_or_default();
36        if name.trim().is_empty() {
37            return Err(LlmError::InvalidRequest(format!(
38                "streamed tool call at index {index} missing function name"
39            )));
40        }
41        let id = partial
42            .id
43            .filter(|s| !s.trim().is_empty())
44            .unwrap_or_else(|| format!("call_{index}"));
45
46        let mut args = partial.arguments;
47        if let Some(repair_fn) = repair {
48            args = repair_fn(&args);
49        } else if options.empty_args_fallback {
50            let trimmed = args.trim();
51            if trimmed.is_empty() || trimmed == "null" || trimmed == "None" {
52                args = "{}".to_string();
53            }
54        }
55
56        let trimmed = args.trim();
57        if trimmed.is_empty() || trimmed == "null" || trimmed == "None" {
58            if options.empty_args_fallback {
59                args = "{}".to_string();
60            } else {
61                return Err(LlmError::InvalidRequest(format!(
62                    "streamed tool call at index {index} finished without arguments"
63                )));
64            }
65        } else if !options.empty_args_fallback
66            && serde_json::from_str::<serde_json::Value>(trimmed).is_err()
67        {
68            return Err(LlmError::InvalidRequest(format!(
69                "streamed tool call at index {index} has invalid JSON arguments"
70            )));
71        }
72
73        out.push(ToolCall {
74            id,
75            call_type: "function".into(),
76            function: FunctionCall {
77                name,
78                arguments: args,
79            },
80            thought_signature: partial.thought_signature,
81        });
82    }
83    Ok(out)
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn finalize_applies_repair_and_default_id() {
92        let mut map = BTreeMap::new();
93        map.insert(
94            0,
95            PartialStreamToolCall {
96                id: None,
97                function_name: Some("read_file".into()),
98                arguments: "null".into(),
99                thought_signature: None,
100            },
101        );
102        let calls = finalize_streamed_tool_calls_with_repair(
103            map,
104            FinalizeStreamToolCallsOptions {
105                empty_args_fallback: true,
106            },
107            Some(|s| {
108                if s.trim() == "null" {
109                    "{}".into()
110                } else {
111                    s.to_string()
112                }
113            }),
114        )
115        .expect("ok");
116        assert_eq!(calls.len(), 1);
117        assert_eq!(calls[0].id, "call_0");
118        assert_eq!(calls[0].function.arguments, "{}");
119    }
120}
edgequake_llm/stream_tools.rs

edgequake_llm/
stream_tools.rs
