zesdex/src/app/runtime/stream/tools/mod.rs

use super::turn::ParsedToolCall;
use serde_json::{json, Value};

pub struct ToolCallAccumulator {
    calls: Vec<ParsedToolCall>,
}

impl ToolCallAccumulator {
    pub fn new() -> Self {
        ToolCallAccumulator { calls: Vec::new() }
    }

    pub fn add_delta(
        &mut self,
        index: usize,
        id: Option<&str>,
        name: Option<&str>,
        arguments_delta: &str,
    ) {
        while self.calls.len() <= index {
            self.calls.push(ParsedToolCall {
                id: String::new(),
                name: String::new(),
                arguments: String::new(),
                is_complete: false,
            });
        }
        let tc = &mut self.calls[index];
        if let Some(new_id) = id {
            if !new_id.is_empty() {
                tc.id = new_id.to_string();
            }
        }
        if let Some(new_name) = name {
            if !new_name.is_empty() {
                tc.name = new_name.to_string();
            }
        }
        tc.arguments.push_str(arguments_delta);
    }

    pub fn calls(&self) -> &[ParsedToolCall] {
        &self.calls
    }

    pub fn is_complete(&self) -> bool {
        !self.calls.is_empty() && self.calls.iter().all(|tc| !tc.name.is_empty() && !tc.arguments.is_empty())
    }

    pub fn reset(&mut self) {
        self.calls.clear();
    }

    pub fn pending_args(&self) -> Vec<Value> {
        self.calls
            .iter()
            .filter(|tc| !tc.name.is_empty())
            .map(|tc| {
                json!({
                    "tool_call_id": tc.id,
                    "name": tc.name,
                    "arguments": tc.arguments,
                })
            })
            .collect()
    }
}

impl Default for ToolCallAccumulator {
    fn default() -> Self {
        Self::new()
    }
}