Skip to main content

ai_lib_rust/utils/
tool_call_assembler.rs

1use crate::types::tool::ToolCall;
2
3/// Collects tool call events (started + argument fragments) into final ToolCall objects.
4/// This is intentionally tolerant: if JSON parsing fails, it keeps the raw string.
5#[derive(Default)]
6pub struct ToolCallAssembler {
7    tool_calls: Vec<ToolCall>,
8}
9
10impl ToolCallAssembler {
11    pub fn new() -> Self {
12        Self::default()
13    }
14
15    pub fn on_started(&mut self, id: String, name: String) {
16        if self.tool_calls.iter().any(|t| t.id == id) {
17            return;
18        }
19        self.tool_calls.push(ToolCall {
20            id,
21            name,
22            arguments: serde_json::Value::String(String::new()),
23        });
24    }
25
26    pub fn on_partial(&mut self, id: &str, fragment: &str) {
27        if let Some(tc) = self.tool_calls.iter_mut().find(|t| t.id == id) {
28            match &mut tc.arguments {
29                serde_json::Value::String(s) => s.push_str(fragment),
30                _ => tc.arguments = serde_json::Value::String(fragment.to_string()),
31            }
32        }
33    }
34
35    pub fn finalize(mut self) -> Vec<ToolCall> {
36        for tc in &mut self.tool_calls {
37            if let serde_json::Value::String(s) = &tc.arguments {
38                let trimmed = s.trim();
39                if !trimmed.is_empty() {
40                    if let Ok(v) = serde_json::from_str::<serde_json::Value>(trimmed) {
41                        tc.arguments = v;
42                    }
43                }
44            }
45        }
46        self.tool_calls
47    }
48}