wasm-agent 0.2.0

// SPDX-License-Identifier: MIT
//! JavaScript bindings (feature `js`): an `Agent` class whose tools are
//! JavaScript functions.
//!
//! ```js
//! import init, { Agent } from "./wasm_agent.js";
//! await init();
//! const agent = new Agent(10, 8192, true);            // max steps, context tokens, recover from errors
//! agent.registerTool("add", "Adds two numbers",
//!   '{"type":"object","properties":{"a":{"type":"number"},"b":{"type":"number"}}}',
//!   (args) => String(JSON.parse(args).a + JSON.parse(args).b));
//! const body = { model: "gpt-4o-mini", tools: agent.openaiTools(), messages: [...] };
//! const step = agent.stepResponse(await (await fetch(url, ...)).text());
//! // step = { kind: "tool_calls", results: [...] } or { kind: "final", text }
//! ```
//!
//! A tool function receives its input as a string and returns a string (other
//! values are JSON-encoded). If it throws, the model sees `Error: <message>`.

use wasm_bindgen::prelude::*;

use crate::loop_runner::NativeStep;
use crate::types::{AgentConfig, ReActStep, ToolResult};
use crate::ToolSpec;

fn js_err(e: crate::AgentError) -> JsError {
    JsError::new(&e.to_string())
}

fn to_js(v: &serde_json::Value) -> Result<JsValue, JsError> {
    js_sys::JSON::parse(&v.to_string()).map_err(|_| JsError::new("could not convert value to JavaScript"))
}

/// A ReAct / tool-calling agent for JavaScript.
#[wasm_bindgen(js_name = Agent)]
pub struct JsAgent {
    inner: crate::Agent,
}

#[wasm_bindgen(js_class = Agent)]
impl JsAgent {
    /// `new Agent(maxIterations = 10, contextTokens = 8192, recoverFromErrors = false)`.
    #[wasm_bindgen(constructor)]
    pub fn new(max_iterations: Option<u32>, context_tokens: Option<usize>, recover_from_errors: Option<bool>) -> JsAgent {
        let d = AgentConfig::default();
        let config = AgentConfig {
            max_iterations: max_iterations.unwrap_or(d.max_iterations),
            context_token_limit: context_tokens.unwrap_or(d.context_token_limit),
            recover_from_errors: recover_from_errors.unwrap_or(false),
            ..d
        };
        JsAgent { inner: crate::Agent::new(config) }
    }

    /// Register a tool backed by a JavaScript function `(input: string) => string`.
    #[wasm_bindgen(js_name = registerTool)]
    pub fn register_tool(&mut self, name: &str, description: &str, schema: &str, f: js_sys::Function) -> Result<(), JsError> {
        let tool = name.to_string();
        self.inner
            .tools_mut()
            .register(
                ToolSpec::new(name, description, schema),
                Box::new(move |input: &str| match f.call1(&JsValue::NULL, &JsValue::from_str(input)) {
                    Ok(v) => {
                        let output = v.as_string().unwrap_or_else(|| {
                            js_sys::JSON::stringify(&v).ok().and_then(|s| s.as_string()).unwrap_or_default()
                        });
                        ToolResult { tool_name: tool.clone(), output, success: true }
                    }
                    Err(e) => {
                        let msg = js_sys::Reflect::get(&e, &JsValue::from_str("message"))
                            .ok()
                            .and_then(|m| m.as_string())
                            .or_else(|| e.as_string())
                            .unwrap_or_else(|| "tool threw an exception".into());
                        ToolResult { tool_name: tool.clone(), output: msg, success: false }
                    }
                }),
            )
            .map_err(js_err)
    }

    /// Add the standard ReAct system prompt (with the tool list).
    #[wasm_bindgen(js_name = addSystemPrompt)]
    pub fn add_system_prompt(&mut self) {
        self.inner.add_system_prompt();
    }

    /// Add the user's question.
    pub fn ask(&mut self, question: &str) {
        self.inner.ask(question);
    }

    /// Process one text reply. Returns `{kind, content, tool?, input?}` where
    /// kind is `"thought"`, `"action"`, `"observation"` or `"final"`.
    pub fn step(&mut self, reply: &str) -> Result<JsValue, JsError> {
        let step = self.inner.step(reply).map_err(js_err)?;
        let v = match step {
            ReActStep::Thought(t) => serde_json::json!({"kind": "thought", "content": t}),
            ReActStep::Action { tool, input } => serde_json::json!({"kind": "action", "content": input, "tool": tool, "input": input}),
            ReActStep::Observation(o) => serde_json::json!({"kind": "observation", "content": o}),
            ReActStep::FinalAnswer(a) => serde_json::json!({"kind": "final", "content": a}),
        };
        to_js(&v)
    }

    /// Process one raw response body from OpenAI, Anthropic or Ollama.
    /// Returns `{kind: "tool_calls", results: [{id, name, success, output}]}`
    /// or `{kind: "final", text}`.
    #[wasm_bindgen(js_name = stepResponse)]
    pub fn step_response(&mut self, body: &str) -> Result<JsValue, JsError> {
        let v = match self.inner.step_response(body).map_err(js_err)? {
            NativeStep::ToolCalls(results) => serde_json::json!({"kind": "tool_calls", "results": results}),
            NativeStep::Final(text) => serde_json::json!({"kind": "final", "text": text}),
        };
        to_js(&v)
    }

    /// The `tools` array for OpenAI-style APIs.
    #[wasm_bindgen(js_name = openaiTools)]
    pub fn openai_tools(&self) -> Result<JsValue, JsError> {
        to_js(&serde_json::Value::Array(self.inner.tools().openai_tools().map_err(js_err)?))
    }

    /// The `tools` array for Anthropic.
    #[wasm_bindgen(js_name = anthropicTools)]
    pub fn anthropic_tools(&self) -> Result<JsValue, JsError> {
        to_js(&serde_json::Value::Array(self.inner.tools().anthropic_tools().map_err(js_err)?))
    }

    /// The conversation as one prompt string (text ReAct).
    pub fn transcript(&self) -> String {
        self.inner.transcript()
    }

    /// The conversation as `[{role, content}]` chat messages (text ReAct).
    #[wasm_bindgen(js_name = chatMessages)]
    pub fn chat_messages(&self) -> Result<JsValue, JsError> {
        to_js(&serde_json::Value::Array(self.inner.history().to_chat_messages()))
    }

    /// Number of replies processed.
    #[wasm_bindgen(getter)]
    pub fn steps(&self) -> usize {
        self.inner.iteration_count()
    }
}