Skip to main content

wasm_agent/
native.rs

1// SPDX-License-Identifier: MIT
2//! Native tool calling: the JSON formats OpenAI, Anthropic and Ollama use.
3//!
4//! Text ReAct (`Action: tool(input)`) works with any model, but current
5//! models are trained to call tools through the API instead. This module
6//! turns a [`ToolRegistry`](crate::ToolRegistry) into the `tools` array each
7//! API expects, reads the tool calls out of a raw response body, and formats
8//! the results to send back. No HTTP client: you send and receive the JSON.
9//!
10//! ```
11//! use wasm_agent::{Agent, AgentConfig, NativeStep, ToolSpec};
12//! use wasm_agent::types::ToolResult;
13//!
14//! let mut agent = Agent::new(AgentConfig::default());
15//! agent.tools_mut().register(
16//!     ToolSpec::new("get_weather", "Weather for a city",
17//!         r#"{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}"#),
18//!     Box::new(|args: &str| ToolResult { tool_name: "get_weather".into(), output: format!("sunny ({args})"), success: true }),
19//! )?;
20//! let tools = agent.tools().openai_tools()?; // put this in the request body
21//! assert_eq!(tools[0]["function"]["name"], "get_weather");
22//!
23//! // The model's reply (OpenAI Chat Completions):
24//! let reply = r#"{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[
25//!     {"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Paris\"}"}}]}}]}"#;
26//! let NativeStep::ToolCalls(results) = agent.step_response(reply)? else { panic!() };
27//! assert_eq!(results[0].output, r#"sunny ({"city":"Paris"})"#);
28//! let follow_up = wasm_agent::native::tool_results_openai(&results); // append to messages
29//! assert_eq!(follow_up[0]["tool_call_id"], "call_1");
30//! # Ok::<(), wasm_agent::AgentError>(())
31//! ```
32
33use crate::error::AgentError;
34use serde::{Deserialize, Serialize};
35use serde_json::{json, Value};
36
37/// One tool call requested by the model.
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct ToolCall {
40    /// The provider's call id (needed when sending the result back), if any.
41    pub id: Option<String>,
42    /// Tool name.
43    pub name: String,
44    /// Arguments as a JSON string (an object for object schemas).
45    pub arguments: String,
46}
47
48/// The result of one tool call, ready to send back.
49#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50pub struct ToolOutput {
51    /// The call id this answers.
52    pub id: Option<String>,
53    /// Tool name.
54    pub name: String,
55    /// Whether the tool succeeded.
56    pub success: bool,
57    /// The tool's output (or error text).
58    pub output: String,
59}
60
61/// What a model response contained.
62#[derive(Debug, Clone, Default, PartialEq, Eq)]
63pub struct ModelTurn {
64    /// Text the model wrote (may accompany tool calls), if any.
65    pub text: Option<String>,
66    /// Tool calls, in order. Empty means a final answer.
67    pub calls: Vec<ToolCall>,
68}
69
70fn args_string(v: Option<&Value>) -> String {
71    match v {
72        // OpenAI sends arguments as a JSON string
73        Some(Value::String(s)) => s.clone(),
74        // Anthropic and Ollama send an object
75        Some(other) => other.to_string(),
76        None => "{}".into(),
77    }
78}
79
80fn push_text(text: &mut Option<String>, piece: &str) {
81    if piece.is_empty() {
82        return;
83    }
84    match text {
85        Some(t) => t.push_str(piece),
86        None => *text = Some(piece.to_string()),
87    }
88}
89
90/// Read the text and tool calls out of a raw response body from OpenAI Chat
91/// Completions (`choices[0].message`), the OpenAI Responses API (`output`),
92/// Anthropic Messages (`content` blocks), or Ollama's `/api/chat` (`message`).
93///
94/// # Errors
95/// [`AgentError::ParseError`] if the body is not JSON or matches none of them.
96pub fn parse_model_response(body: &str) -> Result<ModelTurn, AgentError> {
97    let v: Value = serde_json::from_str(body).map_err(|e| AgentError::ParseError(format!("response is not JSON: {e}")))?;
98    let mut turn = ModelTurn::default();
99
100    // OpenAI Chat Completions and Ollama: a single message object
101    let message = v.pointer("/choices/0/message").or_else(|| v.get("message"));
102    if let Some(m) = message {
103        if let Some(t) = m.get("content").and_then(Value::as_str) {
104            push_text(&mut turn.text, t);
105        }
106        for c in m.get("tool_calls").and_then(Value::as_array).into_iter().flatten() {
107            let f = c.get("function").unwrap_or(c);
108            let name = f.get("name").and_then(Value::as_str).unwrap_or_default().to_string();
109            if name.is_empty() {
110                return Err(AgentError::ParseError("tool call without a name".into()));
111            }
112            turn.calls.push(ToolCall {
113                id: c.get("id").and_then(Value::as_str).map(str::to_string),
114                name,
115                arguments: args_string(f.get("arguments")),
116            });
117        }
118        return Ok(turn);
119    }
120
121    // Anthropic Messages: content blocks
122    if let Some(blocks) = v.get("content").and_then(Value::as_array) {
123        for b in blocks {
124            match b.get("type").and_then(Value::as_str) {
125                Some("text") => push_text(&mut turn.text, b.get("text").and_then(Value::as_str).unwrap_or("")),
126                Some("tool_use") => turn.calls.push(ToolCall {
127                    id: b.get("id").and_then(Value::as_str).map(str::to_string),
128                    name: b.get("name").and_then(Value::as_str).unwrap_or_default().to_string(),
129                    arguments: args_string(b.get("input")),
130                }),
131                _ => {}
132            }
133        }
134        return Ok(turn);
135    }
136
137    // OpenAI Responses API: output items
138    if let Some(items) = v.get("output").and_then(Value::as_array) {
139        for it in items {
140            match it.get("type").and_then(Value::as_str) {
141                Some("function_call") => turn.calls.push(ToolCall {
142                    id: it.get("call_id").or_else(|| it.get("id")).and_then(Value::as_str).map(str::to_string),
143                    name: it.get("name").and_then(Value::as_str).unwrap_or_default().to_string(),
144                    arguments: args_string(it.get("arguments")),
145                }),
146                Some("message") => {
147                    for c in it.get("content").and_then(Value::as_array).into_iter().flatten() {
148                        push_text(&mut turn.text, c.get("text").and_then(Value::as_str).unwrap_or(""));
149                    }
150                }
151                _ => {}
152            }
153        }
154        return Ok(turn);
155    }
156
157    Err(AgentError::ParseError("not an OpenAI, Anthropic or Ollama chat response".into()))
158}
159
160fn shown(o: &ToolOutput) -> String {
161    if o.success { o.output.clone() } else { format!("Error: {}", o.output) }
162}
163
164/// Tool results as OpenAI Chat Completions `tool` messages, to append after
165/// the assistant message that made the calls. Failed calls start with
166/// `Error:` (the format has no error flag).
167pub fn tool_results_openai(outputs: &[ToolOutput]) -> Vec<Value> {
168    outputs
169        .iter()
170        .map(|o| json!({"role": "tool", "tool_call_id": o.id.clone().unwrap_or_default(), "content": shown(o)}))
171        .collect()
172}
173
174/// Tool results as one Anthropic `user` message of `tool_result` blocks.
175pub fn tool_results_anthropic(outputs: &[ToolOutput]) -> Value {
176    let blocks: Vec<Value> = outputs
177        .iter()
178        .map(|o| {
179            let mut b = json!({"type": "tool_result", "tool_use_id": o.id.clone().unwrap_or_default(), "content": o.output});
180            if !o.success {
181                b["is_error"] = json!(true);
182            }
183            b
184        })
185        .collect();
186    json!({"role": "user", "content": blocks})
187}
188
189/// Tool results as OpenAI Responses API `function_call_output` input items.
190/// Failed calls start with `Error:`.
191pub fn tool_results_responses(outputs: &[ToolOutput]) -> Vec<Value> {
192    outputs
193        .iter()
194        .map(|o| json!({"type": "function_call_output", "call_id": o.id.clone().unwrap_or_default(), "output": shown(o)}))
195        .collect()
196}
197
198/// The parameters schema a native API needs, and whether the tool's own
199/// schema had to be wrapped. APIs require an object schema; a tool declared
200/// with a non-object schema (for example `{"type":"string"}`) is exposed as
201/// `{"input": <schema>}` and unwrapped again on dispatch.
202pub(crate) fn api_schema(schema: &str) -> Result<(Value, bool), AgentError> {
203    let trimmed = schema.trim();
204    let v: Value = if trimmed.is_empty() || trimmed == "{}" {
205        json!({"type": "string"})
206    } else {
207        serde_json::from_str(trimmed).map_err(|e| AgentError::InvalidToolSignature(format!("input schema is not JSON: {e}")))?
208    };
209    if v.get("type").and_then(Value::as_str) == Some("object") {
210        Ok((v, false))
211    } else {
212        Ok((json!({"type": "object", "properties": {"input": v}, "required": ["input"]}), true))
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn openai_chat() {
222        let t = parse_model_response(
223            r#"{"choices":[{"message":{"role":"assistant","content":"Let me check.","tool_calls":[
224              {"id":"call_a","type":"function","function":{"name":"x","arguments":"{\"q\":1}"}},
225              {"id":"call_b","type":"function","function":{"name":"y","arguments":"{}"}}]}}]}"#,
226        )
227        .unwrap();
228        assert_eq!(t.text.as_deref(), Some("Let me check."));
229        assert_eq!(t.calls.len(), 2);
230        assert_eq!(t.calls[0], ToolCall { id: Some("call_a".into()), name: "x".into(), arguments: r#"{"q":1}"#.into() });
231        let fin = parse_model_response(r#"{"choices":[{"message":{"role":"assistant","content":"Paris."}}]}"#).unwrap();
232        assert!(fin.calls.is_empty());
233        assert_eq!(fin.text.as_deref(), Some("Paris."));
234    }
235
236    #[test]
237    fn anthropic_messages() {
238        let t = parse_model_response(
239            r#"{"id":"msg","type":"message","role":"assistant","stop_reason":"tool_use","content":[
240              {"type":"text","text":"Checking."},
241              {"type":"tool_use","id":"toolu_1","name":"get_weather","input":{"city":"Paris"}}]}"#,
242        )
243        .unwrap();
244        assert_eq!(t.calls[0].id.as_deref(), Some("toolu_1"));
245        assert_eq!(t.calls[0].arguments, r#"{"city":"Paris"}"#);
246        assert_eq!(t.text.as_deref(), Some("Checking."));
247    }
248
249    #[test]
250    fn openai_responses_and_ollama() {
251        let t = parse_model_response(
252            r#"{"output":[{"type":"function_call","id":"fc_1","call_id":"call_9","name":"x","arguments":"{\"a\":true}"}]}"#,
253        )
254        .unwrap();
255        assert_eq!(t.calls[0].id.as_deref(), Some("call_9"));
256        let o = parse_model_response(
257            r#"{"model":"llama3.2","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"x","arguments":{"a":1}}}]},"done":true}"#,
258        )
259        .unwrap();
260        assert_eq!(o.calls[0], ToolCall { id: None, name: "x".into(), arguments: r#"{"a":1}"#.into() });
261        assert_eq!(o.text, None);
262    }
263
264    #[test]
265    fn rejects_unknown_shapes() {
266        assert!(parse_model_response("not json").is_err());
267        assert!(parse_model_response(r#"{"hello":1}"#).is_err());
268        assert!(parse_model_response(r#"{"choices":[{"message":{"tool_calls":[{"function":{"arguments":"{}"}}]}}]}"#).is_err());
269    }
270
271    #[test]
272    fn result_formats() {
273        let outs = vec![
274            ToolOutput { id: Some("c1".into()), name: "x".into(), success: true, output: "ok".into() },
275            ToolOutput { id: Some("c2".into()), name: "y".into(), success: false, output: "boom".into() },
276        ];
277        assert_eq!(tool_results_openai(&outs)[1], json!({"role": "tool", "tool_call_id": "c2", "content": "Error: boom"}));
278        assert_eq!(tool_results_openai(&outs)[0]["content"], "ok");
279        let a = tool_results_anthropic(&outs);
280        assert_eq!(a["content"][1]["is_error"], json!(true));
281        assert!(a["content"][0].get("is_error").is_none());
282        assert_eq!(tool_results_responses(&outs)[0]["call_id"], "c1");
283    }
284
285    #[test]
286    fn schema_wrapping() {
287        let (s, wrapped) = api_schema(r#"{"type":"object","properties":{}}"#).unwrap();
288        assert!(!wrapped);
289        assert_eq!(s["type"], "object");
290        let (s, wrapped) = api_schema(r#"{"type":"string"}"#).unwrap();
291        assert!(wrapped);
292        assert_eq!(s["properties"]["input"]["type"], "string");
293        assert!(api_schema("{nope").is_err());
294        assert!(api_schema("").unwrap().1);
295    }
296}