diff --git a/crates/agent-wasm/src/lib.rs b/crates/agent-wasm/src/lib.rs index 4dcabf7..aca3846 100644 --- a/crates/agent-wasm/src/lib.rs +++ b/crates/agent-wasm/src/lib.rs @@ -8,12 +8,15 @@ //! one JSON "outcome" object the page inspects to decide whether to call //! tools or show the final answer: //! -//! - `{"type": "needTools", "calls": [{"callId", "name", "arguments"}]}` -//! - `{"type": "final", "text": "..."}` +//! - `{"type": "needTools", "calls": [{"callId", "name", "arguments"}], "timing": T}` +//! - `{"type": "final", "text": "...", "timing": T}` //! - `{"type": "contextOverflow", "needed": n, "available": n}` //! - `{"type": "aborted"}` //! - `{"type": "error", "message": "..."}` //! +//! `T` is `{"promptTokens", "thinkTokens", "answerTokens"}`, the token +//! counts of the step. +//! //! Run it in a Web Worker. While `start` or `provideToolResults` is //! running, the agent is busy: any other call that needs it (a second //! `start`, `provideToolResults`, `setOnText`, `reset`) throws an "agent is @@ -24,20 +27,33 @@ use std::cell::RefCell; use wasm_bindgen::prelude::*; -use agent::agent::{AbortFlag, Agent, PendingToolCall, StepOutcome, Thinking}; +use agent::agent::{AbortFlag, Agent, PendingToolCall, Step, StepOutcome, Thinking}; use agent::lean_generator::LeanGenerator; use anyhow::Result; use lean::{Backend, GgufSource, LoadOptions, TokenizerSource}; use serde_json::Value; +/// Token counts of a step, as the `timing` field of `needTools` and +/// `final` outcomes. Times are not reported: wasm32 has no monotonic clock +/// in the agent, so the page times the step itself. +fn timing_json(step: &Step) -> Value { + serde_json::json!({ + "promptTokens": step.timing.prompt_tokens, + "thinkTokens": step.timing.think_tokens, + "answerTokens": step.timing.answer_tokens, + }) +} + fn outcome_to_json(outcome: StepOutcome) -> String { match outcome { - StepOutcome::NeedTools { calls, .. } => { + StepOutcome::NeedTools { calls, step } => { let calls: Vec = calls.iter().map(|c: &PendingToolCall| serde_json::json!({"callId": c.call_id, "name": c.name, "arguments": c.arguments})).collect(); - serde_json::json!({"type": "needTools", "calls": calls}).to_string() + serde_json::json!({"type": "needTools", "calls": calls, "timing": timing_json(&step)}) + .to_string() } - StepOutcome::Final { text, .. } => { - serde_json::json!({"type": "final", "text": text}).to_string() + StepOutcome::Final { text, step } => { + serde_json::json!({"type": "final", "text": text, "timing": timing_json(&step)}) + .to_string() } StepOutcome::ContextOverflow { needed, available } => { serde_json::json!({"type": "contextOverflow", "needed": needed, "available": available}) @@ -205,3 +221,51 @@ impl AgentJs { *self.agent.borrow_mut() = Some(agent); } } + +#[cfg(test)] +mod tests { + use super::*; + use agent::agent::{ParsedOutput, StepTiming}; + + fn step(timing: StepTiming) -> Step { + Step { + prompt: String::new(), + generated_text: String::new(), + parsed: ParsedOutput::Text(String::new()), + retries: 0, + tool_errors: 0, + oversized_results: 0, + repeat_guard: false, + forced_text_answer: false, + timing, + } + } + + #[test] + fn need_tools_and_final_carry_token_counts() { + let timing = StepTiming { + prompt_tokens: 812, + think_tokens: 40, + answer_tokens: 17, + ..Default::default() + }; + let expected = + serde_json::json!({"promptTokens": 812, "thinkTokens": 40, "answerTokens": 17}); + + let v: Value = serde_json::from_str(&outcome_to_json(StepOutcome::NeedTools { + calls: vec![], + step: step(timing), + })) + .unwrap(); + assert_eq!(v["type"], "needTools"); + assert_eq!(v["timing"], expected); + + let v: Value = serde_json::from_str(&outcome_to_json(StepOutcome::Final { + text: "ok".into(), + step: step(timing), + })) + .unwrap(); + assert_eq!(v["type"], "final"); + assert_eq!(v["timing"], expected); + } +}