atomr_agents_tool/
parser.rs

1//! Tool-call parser.
2//!
3//! Consumes the opaque `tool_call_delta` JSON values that
4//! `atomr-infer-core` streams as part of `TokenChunk`. Producers
5//! differ by provider, so the parser dispatches on a `Provider`
6//! discriminant and accumulates partial JSON-string arguments
7//! across chunks.
8
9use atomr_agents_core::{AgentError, Result, Value};
10use serde::{Deserialize, Serialize};
11use std::collections::BTreeMap;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum Provider {
16    OpenAi,
17    Anthropic,
18}
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct ParsedToolCall {
22    pub id: String,
23    pub name: String,
24    /// JSON string of arguments (assembled from streaming deltas).
25    pub arguments_raw: String,
26}
27
28impl ParsedToolCall {
29    pub fn arguments(&self) -> Result<Value> {
30        if self.arguments_raw.trim().is_empty() {
31            return Ok(Value::Null);
32        }
33        serde_json::from_str::<Value>(&self.arguments_raw)
34            .map_err(|e| AgentError::Tool(format!("tool args parse: {e}")))
35    }
36}
37
38/// Stateful streaming parser. Feed each `tool_call_delta` value as it
39/// arrives; call `finish` to drain the accumulated calls.
40pub struct ToolCallParser {
41    provider: Provider,
42    /// keyed by tool-call index (OpenAI) or content-block index (Anthropic).
43    calls: BTreeMap<u32, Partial>,
44}
45
46#[derive(Debug, Default)]
47struct Partial {
48    id: String,
49    name: String,
50    args: String,
51}
52
53impl ToolCallParser {
54    pub fn new(provider: Provider) -> Self {
55        Self {
56            provider,
57            calls: BTreeMap::new(),
58        }
59    }
60
61    pub fn feed(&mut self, delta: &Value) -> Result<()> {
62        match self.provider {
63            Provider::OpenAi => self.feed_openai(delta),
64            Provider::Anthropic => self.feed_anthropic(delta),
65        }
66    }
67
68    pub fn finish(self) -> Vec<ParsedToolCall> {
69        self.calls
70            .into_values()
71            .map(|p| ParsedToolCall {
72                id: p.id,
73                name: p.name,
74                arguments_raw: p.args,
75            })
76            .collect()
77    }
78
79    fn feed_openai(&mut self, delta: &Value) -> Result<()> {
80        // OpenAI streams an array of tool-call deltas under
81        // `delta.tool_calls`. Each element has `index`, optional `id`,
82        // optional `type`, and `function: { name?, arguments? }`.
83        let arr = delta
84            .get("tool_calls")
85            .and_then(|v| v.as_array())
86            .ok_or_else(|| AgentError::Tool("openai: missing tool_calls".into()))?;
87        for item in arr {
88            let idx = item
89                .get("index")
90                .and_then(|v| v.as_u64())
91                .ok_or_else(|| AgentError::Tool("openai: tool_call missing index".into()))?
92                as u32;
93            let entry = self.calls.entry(idx).or_default();
94            if let Some(id) = item.get("id").and_then(|v| v.as_str()) {
95                entry.id = id.to_string();
96            }
97            if let Some(func) = item.get("function") {
98                if let Some(name) = func.get("name").and_then(|v| v.as_str()) {
99                    entry.name.push_str(name);
100                }
101                if let Some(args) = func.get("arguments").and_then(|v| v.as_str()) {
102                    entry.args.push_str(args);
103                }
104            }
105        }
106        Ok(())
107    }
108
109    fn feed_anthropic(&mut self, delta: &Value) -> Result<()> {
110        // Anthropic SSE event shapes:
111        //   content_block_start { index, content_block: { type:"tool_use", id, name, input: {} } }
112        //   content_block_delta { index, delta: { type:"input_json_delta", partial_json: "..." } }
113        // We accept either shape, identified by which keys are
114        // present.
115        let idx = delta
116            .get("index")
117            .and_then(|v| v.as_u64())
118            .ok_or_else(|| AgentError::Tool("anthropic: delta missing index".into()))?
119            as u32;
120        let entry = self.calls.entry(idx).or_default();
121        if let Some(block) = delta.get("content_block") {
122            if block.get("type").and_then(|v| v.as_str()) == Some("tool_use") {
123                if let Some(id) = block.get("id").and_then(|v| v.as_str()) {
124                    entry.id = id.to_string();
125                }
126                if let Some(name) = block.get("name").and_then(|v| v.as_str()) {
127                    entry.name = name.to_string();
128                }
129                // Sometimes the full input arrives at start. Skip
130                // empty objects since deltas will follow.
131                if let Some(input) = block.get("input") {
132                    let is_empty_object = input.as_object().map(|m| m.is_empty()).unwrap_or(false);
133                    if !input.is_null() && !is_empty_object {
134                        entry.args = serde_json::to_string(input).unwrap_or_default();
135                    }
136                }
137            }
138        }
139        if let Some(d) = delta.get("delta") {
140            if d.get("type").and_then(|v| v.as_str()) == Some("input_json_delta") {
141                if let Some(pj) = d.get("partial_json").and_then(|v| v.as_str()) {
142                    entry.args.push_str(pj);
143                }
144            }
145        }
146        Ok(())
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use serde_json::json;
154
155    #[test]
156    fn openai_streaming_assembly() {
157        let mut p = ToolCallParser::new(Provider::OpenAi);
158        p.feed(&json!({
159            "tool_calls": [{
160                "index": 0,
161                "id": "call_abc",
162                "type": "function",
163                "function": {"name": "get_weather", "arguments": "{\"city\":"}
164            }]
165        }))
166        .unwrap();
167        p.feed(&json!({
168            "tool_calls": [{
169                "index": 0,
170                "function": {"arguments": "\"NYC\"}"}
171            }]
172        }))
173        .unwrap();
174        let calls = p.finish();
175        assert_eq!(calls.len(), 1);
176        assert_eq!(calls[0].id, "call_abc");
177        assert_eq!(calls[0].name, "get_weather");
178        let args = calls[0].arguments().unwrap();
179        assert_eq!(args, json!({"city": "NYC"}));
180    }
181
182    #[test]
183    fn anthropic_streaming_assembly() {
184        let mut p = ToolCallParser::new(Provider::Anthropic);
185        p.feed(&json!({
186            "index": 1,
187            "content_block": {
188                "type": "tool_use",
189                "id": "toolu_xyz",
190                "name": "search",
191                "input": {}
192            }
193        }))
194        .unwrap();
195        p.feed(&json!({
196            "index": 1,
197            "delta": {"type": "input_json_delta", "partial_json": "{\"q\":"}
198        }))
199        .unwrap();
200        p.feed(&json!({
201            "index": 1,
202            "delta": {"type": "input_json_delta", "partial_json": "\"rust\"}"}
203        }))
204        .unwrap();
205        let calls = p.finish();
206        assert_eq!(calls.len(), 1);
207        assert_eq!(calls[0].id, "toolu_xyz");
208        assert_eq!(calls[0].name, "search");
209        assert_eq!(calls[0].arguments().unwrap(), json!({"q": "rust"}));
210    }
211}
atomr_agents_tool/parser.rs

atomr_agents_tool/
parser.rs
