wasm-agent 0.2.0

// SPDX-License-Identifier: MIT
//! Native tool calling: the JSON formats OpenAI, Anthropic and Ollama use.
//!
//! Text ReAct (`Action: tool(input)`) works with any model, but current
//! models are trained to call tools through the API instead. This module
//! turns a [`ToolRegistry`](crate::ToolRegistry) into the `tools` array each
//! API expects, reads the tool calls out of a raw response body, and formats
//! the results to send back. No HTTP client: you send and receive the JSON.
//!
//! ```
//! use wasm_agent::{Agent, AgentConfig, NativeStep, ToolSpec};
//! use wasm_agent::types::ToolResult;
//!
//! let mut agent = Agent::new(AgentConfig::default());
//! agent.tools_mut().register(
//!     ToolSpec::new("get_weather", "Weather for a city",
//!         r#"{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}"#),
//!     Box::new(|args: &str| ToolResult { tool_name: "get_weather".into(), output: format!("sunny ({args})"), success: true }),
//! )?;
//! let tools = agent.tools().openai_tools()?; // put this in the request body
//! assert_eq!(tools[0]["function"]["name"], "get_weather");
//!
//! // The model's reply (OpenAI Chat Completions):
//! let reply = r#"{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[
//!     {"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Paris\"}"}}]}}]}"#;
//! let NativeStep::ToolCalls(results) = agent.step_response(reply)? else { panic!() };
//! assert_eq!(results[0].output, r#"sunny ({"city":"Paris"})"#);
//! let follow_up = wasm_agent::native::tool_results_openai(&results); // append to messages
//! assert_eq!(follow_up[0]["tool_call_id"], "call_1");
//! # Ok::<(), wasm_agent::AgentError>(())
//! ```

use crate::error::AgentError;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

/// One tool call requested by the model.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCall {
    /// The provider's call id (needed when sending the result back), if any.
    pub id: Option<String>,
    /// Tool name.
    pub name: String,
    /// Arguments as a JSON string (an object for object schemas).
    pub arguments: String,
}

/// The result of one tool call, ready to send back.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolOutput {
    /// The call id this answers.
    pub id: Option<String>,
    /// Tool name.
    pub name: String,
    /// Whether the tool succeeded.
    pub success: bool,
    /// The tool's output (or error text).
    pub output: String,
}

/// What a model response contained.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ModelTurn {
    /// Text the model wrote (may accompany tool calls), if any.
    pub text: Option<String>,
    /// Tool calls, in order. Empty means a final answer.
    pub calls: Vec<ToolCall>,
}

fn args_string(v: Option<&Value>) -> String {
    match v {
        // OpenAI sends arguments as a JSON string
        Some(Value::String(s)) => s.clone(),
        // Anthropic and Ollama send an object
        Some(other) => other.to_string(),
        None => "{}".into(),
    }
}

fn push_text(text: &mut Option<String>, piece: &str) {
    if piece.is_empty() {
        return;
    }
    match text {
        Some(t) => t.push_str(piece),
        None => *text = Some(piece.to_string()),
    }
}

/// Read the text and tool calls out of a raw response body from OpenAI Chat
/// Completions (`choices[0].message`), the OpenAI Responses API (`output`),
/// Anthropic Messages (`content` blocks), or Ollama's `/api/chat` (`message`).
///
/// # Errors
/// [`AgentError::ParseError`] if the body is not JSON or matches none of them.
pub fn parse_model_response(body: &str) -> Result<ModelTurn, AgentError> {
    let v: Value = serde_json::from_str(body).map_err(|e| AgentError::ParseError(format!("response is not JSON: {e}")))?;
    let mut turn = ModelTurn::default();

    // OpenAI Chat Completions and Ollama: a single message object
    let message = v.pointer("/choices/0/message").or_else(|| v.get("message"));
    if let Some(m) = message {
        if let Some(t) = m.get("content").and_then(Value::as_str) {
            push_text(&mut turn.text, t);
        }
        for c in m.get("tool_calls").and_then(Value::as_array).into_iter().flatten() {
            let f = c.get("function").unwrap_or(c);
            let name = f.get("name").and_then(Value::as_str).unwrap_or_default().to_string();
            if name.is_empty() {
                return Err(AgentError::ParseError("tool call without a name".into()));
            }
            turn.calls.push(ToolCall {
                id: c.get("id").and_then(Value::as_str).map(str::to_string),
                name,
                arguments: args_string(f.get("arguments")),
            });
        }
        return Ok(turn);
    }

    // Anthropic Messages: content blocks
    if let Some(blocks) = v.get("content").and_then(Value::as_array) {
        for b in blocks {
            match b.get("type").and_then(Value::as_str) {
                Some("text") => push_text(&mut turn.text, b.get("text").and_then(Value::as_str).unwrap_or("")),
                Some("tool_use") => turn.calls.push(ToolCall {
                    id: b.get("id").and_then(Value::as_str).map(str::to_string),
                    name: b.get("name").and_then(Value::as_str).unwrap_or_default().to_string(),
                    arguments: args_string(b.get("input")),
                }),
                _ => {}
            }
        }
        return Ok(turn);
    }

    // OpenAI Responses API: output items
    if let Some(items) = v.get("output").and_then(Value::as_array) {
        for it in items {
            match it.get("type").and_then(Value::as_str) {
                Some("function_call") => turn.calls.push(ToolCall {
                    id: it.get("call_id").or_else(|| it.get("id")).and_then(Value::as_str).map(str::to_string),
                    name: it.get("name").and_then(Value::as_str).unwrap_or_default().to_string(),
                    arguments: args_string(it.get("arguments")),
                }),
                Some("message") => {
                    for c in it.get("content").and_then(Value::as_array).into_iter().flatten() {
                        push_text(&mut turn.text, c.get("text").and_then(Value::as_str).unwrap_or(""));
                    }
                }
                _ => {}
            }
        }
        return Ok(turn);
    }

    Err(AgentError::ParseError("not an OpenAI, Anthropic or Ollama chat response".into()))
}

fn shown(o: &ToolOutput) -> String {
    if o.success { o.output.clone() } else { format!("Error: {}", o.output) }
}

/// Tool results as OpenAI Chat Completions `tool` messages, to append after
/// the assistant message that made the calls. Failed calls start with
/// `Error:` (the format has no error flag).
pub fn tool_results_openai(outputs: &[ToolOutput]) -> Vec<Value> {
    outputs
        .iter()
        .map(|o| json!({"role": "tool", "tool_call_id": o.id.clone().unwrap_or_default(), "content": shown(o)}))
        .collect()
}

/// Tool results as one Anthropic `user` message of `tool_result` blocks.
pub fn tool_results_anthropic(outputs: &[ToolOutput]) -> Value {
    let blocks: Vec<Value> = outputs
        .iter()
        .map(|o| {
            let mut b = json!({"type": "tool_result", "tool_use_id": o.id.clone().unwrap_or_default(), "content": o.output});
            if !o.success {
                b["is_error"] = json!(true);
            }
            b
        })
        .collect();
    json!({"role": "user", "content": blocks})
}

/// Tool results as OpenAI Responses API `function_call_output` input items.
/// Failed calls start with `Error:`.
pub fn tool_results_responses(outputs: &[ToolOutput]) -> Vec<Value> {
    outputs
        .iter()
        .map(|o| json!({"type": "function_call_output", "call_id": o.id.clone().unwrap_or_default(), "output": shown(o)}))
        .collect()
}

/// The parameters schema a native API needs, and whether the tool's own
/// schema had to be wrapped. APIs require an object schema; a tool declared
/// with a non-object schema (for example `{"type":"string"}`) is exposed as
/// `{"input": <schema>}` and unwrapped again on dispatch.
pub(crate) fn api_schema(schema: &str) -> Result<(Value, bool), AgentError> {
    let trimmed = schema.trim();
    let v: Value = if trimmed.is_empty() || trimmed == "{}" {
        json!({"type": "string"})
    } else {
        serde_json::from_str(trimmed).map_err(|e| AgentError::InvalidToolSignature(format!("input schema is not JSON: {e}")))?
    };
    if v.get("type").and_then(Value::as_str) == Some("object") {
        Ok((v, false))
    } else {
        Ok((json!({"type": "object", "properties": {"input": v}, "required": ["input"]}), true))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn openai_chat() {
        let t = parse_model_response(
            r#"{"choices":[{"message":{"role":"assistant","content":"Let me check.","tool_calls":[
              {"id":"call_a","type":"function","function":{"name":"x","arguments":"{\"q\":1}"}},
              {"id":"call_b","type":"function","function":{"name":"y","arguments":"{}"}}]}}]}"#,
        )
        .unwrap();
        assert_eq!(t.text.as_deref(), Some("Let me check."));
        assert_eq!(t.calls.len(), 2);
        assert_eq!(t.calls[0], ToolCall { id: Some("call_a".into()), name: "x".into(), arguments: r#"{"q":1}"#.into() });
        let fin = parse_model_response(r#"{"choices":[{"message":{"role":"assistant","content":"Paris."}}]}"#).unwrap();
        assert!(fin.calls.is_empty());
        assert_eq!(fin.text.as_deref(), Some("Paris."));
    }

    #[test]
    fn anthropic_messages() {
        let t = parse_model_response(
            r#"{"id":"msg","type":"message","role":"assistant","stop_reason":"tool_use","content":[
              {"type":"text","text":"Checking."},
              {"type":"tool_use","id":"toolu_1","name":"get_weather","input":{"city":"Paris"}}]}"#,
        )
        .unwrap();
        assert_eq!(t.calls[0].id.as_deref(), Some("toolu_1"));
        assert_eq!(t.calls[0].arguments, r#"{"city":"Paris"}"#);
        assert_eq!(t.text.as_deref(), Some("Checking."));
    }

    #[test]
    fn openai_responses_and_ollama() {
        let t = parse_model_response(
            r#"{"output":[{"type":"function_call","id":"fc_1","call_id":"call_9","name":"x","arguments":"{\"a\":true}"}]}"#,
        )
        .unwrap();
        assert_eq!(t.calls[0].id.as_deref(), Some("call_9"));
        let o = parse_model_response(
            r#"{"model":"llama3.2","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"x","arguments":{"a":1}}}]},"done":true}"#,
        )
        .unwrap();
        assert_eq!(o.calls[0], ToolCall { id: None, name: "x".into(), arguments: r#"{"a":1}"#.into() });
        assert_eq!(o.text, None);
    }

    #[test]
    fn rejects_unknown_shapes() {
        assert!(parse_model_response("not json").is_err());
        assert!(parse_model_response(r#"{"hello":1}"#).is_err());
        assert!(parse_model_response(r#"{"choices":[{"message":{"tool_calls":[{"function":{"arguments":"{}"}}]}}]}"#).is_err());
    }

    #[test]
    fn result_formats() {
        let outs = vec![
            ToolOutput { id: Some("c1".into()), name: "x".into(), success: true, output: "ok".into() },
            ToolOutput { id: Some("c2".into()), name: "y".into(), success: false, output: "boom".into() },
        ];
        assert_eq!(tool_results_openai(&outs)[1], json!({"role": "tool", "tool_call_id": "c2", "content": "Error: boom"}));
        assert_eq!(tool_results_openai(&outs)[0]["content"], "ok");
        let a = tool_results_anthropic(&outs);
        assert_eq!(a["content"][1]["is_error"], json!(true));
        assert!(a["content"][0].get("is_error").is_none());
        assert_eq!(tool_results_responses(&outs)[0]["call_id"], "c1");
    }

    #[test]
    fn schema_wrapping() {
        let (s, wrapped) = api_schema(r#"{"type":"object","properties":{}}"#).unwrap();
        assert!(!wrapped);
        assert_eq!(s["type"], "object");
        let (s, wrapped) = api_schema(r#"{"type":"string"}"#).unwrap();
        assert!(wrapped);
        assert_eq!(s["properties"]["input"]["type"], "string");
        assert!(api_schema("{nope").is_err());
        assert!(api_schema("").unwrap().1);
    }
}