diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..acff135 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,5 @@ +# wasm32 SIMD128 in every wasm build of this workspace (lean's CPU forward +# pass uses SIMD128 kernels when it is compiled in). Every browser with +# WebGPU or a recent wasm engine supports it. +[target.wasm32-unknown-unknown] +rustflags = ["-C", "target-feature=+simd128"] diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index d0bc154..bc6657e 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -3,13 +3,15 @@ name: pages on: push: branches: [main] + pull_request: + branches: [main] workflow_dispatch: permissions: contents: read concurrency: - group: pages + group: pages-${{ github.ref }} cancel-in-progress: false jobs: @@ -30,10 +32,14 @@ jobs: run: ENGINE_BUILD="${{ github.sha }}" scripts/build.sh - uses: actions/upload-pages-artifact@v3 + if: github.event_name != 'pull_request' with: path: _site + # Off until the repo has a page to publish: _site holds only pkg/, with + # no index.html. deploy: + if: false needs: build runs-on: ubuntu-latest permissions: diff --git a/.gitignore b/.gitignore index 308bfda..3fc50ad 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ **/pkg/ **/pkg-mt/ crates/agent/reference/.venv +.claude/ diff --git a/Cargo.lock b/Cargo.lock index b56eb07..e962c21 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9,8 +9,7 @@ dependencies = [ "anyhow", "clap", "lean", - "minijinja", - "minijinja-contrib", + "lean-chat", "pollster", "serde", "serde_json", @@ -26,10 +25,10 @@ dependencies = [ "console_error_panic_hook", "js-sys", "lean", + "lean-chat", "serde_json", - "tokenizers", "wasm-bindgen", - "web-sys", + "wasm-bindgen-futures", ] [[package]] @@ -787,20 +786,15 @@ checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" [[package]] name = "lean" version = "0.1.0" -source = "git+https://github.com/idle-intelligence/llm-web?rev=b321c16117cc2aef425c6471cbcb4cd7257fe518#b321c16117cc2aef425c6471cbcb4cd7257fe518" +source = "git+https://github.com/idle-intelligence/llm-web?rev=1fce23d866c23fb6f3635abb51db52c50ff87808#1fce23d866c23fb6f3635abb51db52c50ff87808" dependencies = [ "anyhow", - "bytemuck", - "byteorder", - "clap", "console_error_panic_hook", - "futures-channel", "getrandom", - "half", "js-sys", - "minijinja", - "minijinja-contrib", - "pollster", + "lean-chat", + "lean-core", + "rayon", "serde", "serde_json", "tokenizers", @@ -810,6 +804,36 @@ dependencies = [ "wgpu", ] +[[package]] +name = "lean-chat" +version = "0.1.0" +source = "git+https://github.com/idle-intelligence/llm-web?rev=1fce23d866c23fb6f3635abb51db52c50ff87808#1fce23d866c23fb6f3635abb51db52c50ff87808" +dependencies = [ + "anyhow", + "js-sys", + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "wasm-bindgen", +] + +[[package]] +name = "lean-core" +version = "0.1.0" +source = "git+https://github.com/idle-intelligence/llm-web?rev=1fce23d866c23fb6f3635abb51db52c50ff87808#1fce23d866c23fb6f3635abb51db52c50ff87808" +dependencies = [ + "anyhow", + "bytemuck", + "byteorder", + "futures-channel", + "half", + "pollster", + "rayon", + "wasm-bindgen", + "wgpu", +] + [[package]] name = "libc" version = "0.2.190" diff --git a/Cargo.toml b/Cargo.toml index 2e92bb5..836f38d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,23 +13,22 @@ edition = "2021" anyhow = "1.0" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["preserve_order"] } -minijinja = { version = "2", features = ["json", "preserve_order"] } -minijinja-contrib = { version = "2.24.0", features = ["pycompat"] } tokenizers = { version = "0.22", default-features = false } clap = { version = "4.5", features = ["derive", "env"] } +pollster = "0.3" wasm-bindgen = "0.2" +wasm-bindgen-futures = "0.4" js-sys = "0.3" -web-sys = { version = "0.3", features = ["console"] } console_error_panic_hook = "0.1" -getrandom = { version = "0.3", default-features = false } -# Pinned to an exact commit of origin/main so a clean clone always resolves -# the same lean API, with no path dependency anywhere in this workspace. -# Default features (native-bin) pull in the pollster/clap lean's own -# native Engine::new()/CLI need; agent-wasm overrides to default-features -# = false, features = ["web"] for the wasm32 build. -lean = { git = "https://github.com/idle-intelligence/llm-web", rev = "b321c16117cc2aef425c6471cbcb4cd7257fe518" } +# lean's engine (text in, text out) and its chat layer (messages and tools +# to prompt text with the model's own template), pinned to one commit of +# llm-web so a clean clone always resolves the same API. No default +# features: the agent crate turns on `threads` natively and agent-wasm +# turns on `web`. +lean = { git = "https://github.com/idle-intelligence/llm-web", rev = "1fce23d866c23fb6f3635abb51db52c50ff87808", default-features = false } +lean-chat = { git = "https://github.com/idle-intelligence/llm-web", rev = "1fce23d866c23fb6f3635abb51db52c50ff87808" } [profile.release] lto = true diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..06c18a8 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Thomas Cambau + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index 0174c94..b93db3d 100644 --- a/README.md +++ b/README.md @@ -1,49 +1,79 @@ # agent-web -A tool-calling agent for Qwen2.5-Instruct models, running on the [lean](https://github.com/idle-intelligence/llm-web) inference engine. +A tool-calling agent loop for Qwen2.5-Instruct models, running on the [lean](https://github.com/idle-intelligence/llm-web) inference engine. -lean is inference only: it loads weights, tokenizes, runs the model forward pass, and picks a backend (WebGPU, then CPU threads, then a single CPU thread) by capability. Everything an agent needs on top of that, the tool-calling chat template, output parsing, the step-wise loop, and an eval harness, lives here instead. +lean runs the model: text in, text out, on a backend it picks by capability (WebGPU, then CPU threads, then a single CPU thread). Its chat layer, lean-chat, renders messages and tools into prompt text with the model's own chat template. This repo holds what an agent adds on top: parsing the model's tool calls, the step-wise loop, and an offline eval harness. The loop reaches the model through one small trait (prompt text in, streamed text out, abort, a typed context-overflow error), so any engine that completes text can drive it. ## What it does -- Renders Qwen2.5-Instruct's own tools chat template (`...` in the system turn, `{"name":..., "arguments":...}` in the model's output), token-exact against `transformers.apply_chat_template(..., tools=...)`. -- Parses the model's tool calls out of its raw output, with a retry policy for malformed or repeated calls. -- Drives a step-wise loop: render prompt, generate, parse, hand back pending tool calls, accept results, repeat until a final text answer. -- Scores transcripts against a fixed set of utterance cases (a Sonos speaker-control tool set) under several acceptance modes (exact call sequence, required subset, no mutation, any valid play call). +- Renders each step's prompt with the model's own chat template and the tools, through lean-chat. For Qwen2.5-Instruct that is `...` in the system turn and `{"name":..., "arguments":...}` in the model's output. The `fixture_parity_agent_tools` test checks the prompts byte for byte against `transformers.apply_chat_template(..., tools=...)`. +- Parses the model's tool calls out of its raw output, and regenerates when a reply is malformed, calls an unknown tool, or repeats a call. +- Drives a step-wise loop: render the prompt, generate, parse, hand back the pending tool calls, take their results, and repeat until a final text answer. +- Scores runs against a fixed set of utterance cases for a Sonos speaker-control tool set, under four acceptance modes: exact call sequence, required calls in order, no mutating call, any valid play call. ## Crates -- `crates/agent`: the template, parser, loop, and eval library. Depends on lean through a pinned git revision. -- `crates/agent-wasm`: wasm-bindgen bindings for the agent, running on lean's native CPU engine. -- `examples/agent_cli.rs`: a native CLI that runs one utterance through the loop, answering tool calls from fixture JSON, and prints each step. -- `examples/agent_eval.rs`: a native CLI that runs the full eval set and prints a markdown results table. +- `crates/agent`: the parser, the loop and the eval harness, plus `LeanGenerator`, the loop's generator on lean's engine. Its examples are two native CLIs: `agent_cli` runs one utterance through the loop and prints each step, answering tool calls from fixture files; `agent_eval` runs the eval set and prints a markdown results table. +- `crates/agent-wasm`: wasm-bindgen bindings for the agent on lean's engine, for a Web Worker. lean uses WebGPU when the browser grants an adapter, otherwise the CPU. -## CLI examples +lean and lean-chat come from a pinned git revision of [idle-intelligence/llm-web](https://github.com/idle-intelligence/llm-web). + +## Model + +The examples and the model tests use [Qwen2.5-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-3B-Instruct): the Q4_0 GGUF from [Qwen/Qwen2.5-3B-Instruct-GGUF](https://huggingface.co/Qwen/Qwen2.5-3B-Instruct-GGUF), and the tokenizer files from the model repo. + +```sh +hf download Qwen/Qwen2.5-3B-Instruct-GGUF qwen2.5-3b-instruct-q4_0.gguf --local-dir models/Qwen2.5-3B-Instruct-GGUF +hf download Qwen/Qwen2.5-3B-Instruct tokenizer.json tokenizer_config.json --local-dir models/Qwen2.5-3B-Instruct +``` + +## Build and test + +```sh +cargo build --release +cargo test --workspace +``` + +The tests that need the model files are ignored by default. With the files above: + +```sh +LEAN_GGUF_QWEN25_3B=$PWD/models/Qwen2.5-3B-Instruct-GGUF/qwen2.5-3b-instruct-q4_0.gguf \ +LEAN_TOKENIZER_DIR_QWEN25_3B=$PWD/models/Qwen2.5-3B-Instruct \ +cargo test -p agent --release -- --ignored +``` + +The wasm package (needs [wasm-pack](https://rustwasm.github.io/wasm-pack/)) is built into `_site/pkg` by: + +```sh +scripts/build.sh +``` + +## Run + +One utterance, with the tool calls answered from `crates/agent/fixtures/sonos/results/`: ```sh cargo run -p agent --example agent_cli --release -- \ - --model /path/to/qwen2.5-3b-instruct-q4_0.gguf \ - --tokenizer /path/to/Qwen2.5-3B-Instruct \ - --tools crates/agent/fixtures/sonos/tools-12.json \ + --model models/Qwen2.5-3B-Instruct-GGUF/qwen2.5-3b-instruct-q4_0.gguf \ + --tokenizer models/Qwen2.5-3B-Instruct \ + --tools crates/agent/fixtures/sonos/tools-subset.json \ --utterance "pause the kitchen" ``` +The eval set, as a markdown table on stdout (`AGENT_EVAL_MAX_CTX` sets the context size in tokens, default 4096; `AGENT_EVAL_DATE` the date in the table's title): + ```sh -LEAN_GGUF_QWEN25_3B=/path/to/qwen2.5-3b-instruct-q4_0.gguf \ -LEAN_TOKENIZER_DIR_QWEN25_3B=/path/to/Qwen2.5-3B-Instruct \ +LEAN_GGUF_QWEN25_3B=models/Qwen2.5-3B-Instruct-GGUF/qwen2.5-3b-instruct-q4_0.gguf \ +LEAN_TOKENIZER_DIR_QWEN25_3B=models/Qwen2.5-3B-Instruct \ cargo run -p agent --example agent_eval --release ``` -Both pick lean's GPU engine when a wgpu adapter is available and fall back to the CPU engine otherwise, the same capability-based choice lean's own CLI makes. - -## Demo - -A browser demo on the repo's GitHub Pages is planned; link to follow once it is built. +Both pick the backend by capability: the GPU when a wgpu adapter is available, otherwise the CPU forward pass on every core. Each step reuses the KV cache of the previous one for the prompt prefix they share, so the system prompt and the tools are prefilled once per turn. ## Eval -43 utterance cases, a 12-tool Sonos speaker-control set, Qwen2.5-3B-Instruct (Q4_0 GGUF), lean's CPU engine, unconstrained decoding (no grammar-constrained output): 9.1% correct, run on 2026-10-07. See `docs/runs/2026-10-07-mcp-agent-lean.md` for the full run log, including a 0.5B-model follow-up and the baseline this was compared against. +`eval/utterances.json` holds 46 cases. With the 13-tool subset (`crates/agent/fixtures/sonos/tools-subset.json`), 33 of them are solvable and scored; the other 13 need a tool outside the subset and are skipped. Qwen2.5-3B-Instruct Q4_0, run on the CPU with unconstrained decoding (no grammar-constrained output) on 2026-10-07, got 3 of the 33 right (9.1%). The per-case results are in `docs/runs/2026-10-07-sonos-eval.md`. ## License -Not yet decided; no license file is included. +The code is under the MIT license, see `LICENSE`. The model weights are not part of this repo and come under their own license, given on each model's page. diff --git a/crates/agent-wasm/Cargo.toml b/crates/agent-wasm/Cargo.toml index e31d103..0d1b862 100644 --- a/crates/agent-wasm/Cargo.toml +++ b/crates/agent-wasm/Cargo.toml @@ -11,8 +11,8 @@ agent = { path = "../agent" } anyhow = { workspace = true } serde_json = { workspace = true } lean = { workspace = true, features = ["web"] } +lean-chat = { workspace = true } wasm-bindgen = { workspace = true } +wasm-bindgen-futures = { workspace = true } js-sys = { workspace = true } -web-sys = { workspace = true } console_error_panic_hook = { workspace = true } -tokenizers = { workspace = true, features = ["unstable_wasm", "fancy-regex"] } diff --git a/crates/agent-wasm/src/lib.rs b/crates/agent-wasm/src/lib.rs index 67c71f4..75b35ab 100644 --- a/crates/agent-wasm/src/lib.rs +++ b/crates/agent-wasm/src/lib.rs @@ -1,61 +1,35 @@ -//! wasm-bindgen surface for the tool-calling agent, built on lean's public -//! native CPU engine (`lean::cpu`). A browser can't call an MCP tool -//! synchronously (it's an async round trip through the page), so -//! `start`/`provideToolResults` each return one JSON "outcome" object the -//! page inspects to decide whether to call tools or show the final -//! answer: +//! wasm-bindgen surface for the tool-calling agent, on lean's engine. +//! lean picks its backend by capability when the model loads: WebGPU when +//! the browser grants an adapter and the model loads on it, otherwise the +//! CPU forward pass (SIMD128, one thread). +//! +//! A browser can't call an MCP tool synchronously (it's an async round +//! trip through the page), so `start`/`provideToolResults` each resolve to +//! 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": "contextOverflow", "needed": n, "available": n}` +//! - `{"type": "aborted"}` //! - `{"type": "error", "message": "..."}` //! -//! Only the CPU engine is wired here. +//! 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 +//! busy" error. `backend` and `abortHandle()` always work; stop a running +//! step through the handle. -use std::io::Cursor; -use std::rc::Rc; +use std::cell::RefCell; use wasm_bindgen::prelude::*; -use agent::agent::{Agent, AgentChatTemplate, Generator, PendingToolCall, StepOutcome, ToolCaller}; +use agent::agent::{AbortFlag, Agent, PendingToolCall, StepOutcome}; +use agent::lean_generator::LeanGenerator; use anyhow::Result; +use lean::{Backend, GgufSource, LoadOptions, TokenizerSource}; use serde_json::Value; -/// Greedy CPU generator over a shared (`Rc`) `CpuModel`. `Rc`, not a -/// borrow, so the `Agent` that owns one of these can live inside -/// `AgentCpu` directly and persist across `start`/`provideToolResults` -/// calls. -struct RcCpuGenerator { - model: Rc, -} - -impl Generator for RcCpuGenerator { - fn generate(&mut self, prompt_ids: &[u32], max_new_tokens: usize) -> Result> { - let max_ctx = prompt_ids.len() + max_new_tokens; - let mut cache = lean::cpu::CpuKvCache::new(&self.model.config, max_ctx); - let logits = lean::cpu::forward_prefill(&self.model, &mut cache, prompt_ids); - let mut next_id = lean::model::argmax(&logits); - let mut generated = Vec::with_capacity(max_new_tokens); - for _ in 0..max_new_tokens { - if self.model.config.eos_token_ids.contains(&next_id) { - break; - } - generated.push(next_id); - next_id = lean::cpu::forward_decode_step_argmax(&self.model, &mut cache, next_id); - } - Ok(generated) - } -} - -/// No-op `ToolCaller`: the step-wise wasm API never calls tools itself, -/// `start`/`provide_tool_results` only ever reach `StepOutcome::NeedTools` -/// or `StepOutcome::Final`, never `Agent::run`. -struct NoCaller; -impl ToolCaller for NoCaller { - fn call(&mut self, name: &str, _args: &Value) -> Result { - anyhow::bail!("AgentCpu never calls tools itself (tool: {name}); the page does, via provideToolResults") - } -} - fn outcome_to_json(outcome: StepOutcome) -> String { match outcome { StepOutcome::NeedTools { calls, .. } => { @@ -65,98 +39,145 @@ fn outcome_to_json(outcome: StepOutcome) -> String { StepOutcome::Final { text, .. } => { serde_json::json!({"type": "final", "text": text}).to_string() } + StepOutcome::ContextOverflow { needed, available } => { + serde_json::json!({"type": "contextOverflow", "needed": needed, "available": available}) + .to_string() + } + StepOutcome::Aborted => serde_json::json!({"type": "aborted"}).to_string(), StepOutcome::Error { message, .. } => { serde_json::json!({"type": "error", "message": message}).to_string() } } } +/// Stops the agent's running step; see the module doc. #[wasm_bindgen] -pub struct AgentCpu { - agent: Option>, -} +pub struct AgentAbort(AbortFlag); #[wasm_bindgen] -impl AgentCpu { - #[wasm_bindgen(js_name = create)] - pub fn create() -> AgentCpu { - console_error_panic_hook::set_once(); - AgentCpu { agent: None } +impl AgentAbort { + pub fn abort(&self) { + self.0.abort(); } +} + +#[wasm_bindgen(js_name = Agent)] +pub struct AgentJs { + /// `None` while a step runs. + agent: RefCell>>, + backend: &'static str, + abort: AbortFlag, +} - /// `gguf_bytes` is copied out of the `Uint8Array` into a `Vec` and - /// wrapped in a `std::io::Cursor`, matching `lean::cpu::CpuModel`'s own - /// doc comment on how its wasm callers are expected to load a GGUF - /// (`load_from_reader` is generic over any `Read + Seek`, no - /// lean-internal wasm-byte-reader type is needed here). - #[wasm_bindgen(js_name = load)] - pub fn load( - &mut self, +#[wasm_bindgen(js_class = Agent)] +impl AgentJs { + /// `gguf_bytes` is the GGUF file, `tokenizer_json` the text of + /// `tokenizer.json`, `tokenizer_config_json` the text of + /// `tokenizer_config.json` (its `chat_template` renders every prompt). + /// `max_ctx` is the KV cache size in tokens: each step's prompt plus + /// `max_new_tokens` must fit. + pub async fn load( gguf_bytes: js_sys::Uint8Array, tokenizer_json: String, tokenizer_config_json: String, system_prompt: String, max_new_tokens: u32, - ) -> Result<(), JsError> { - let bytes = gguf_bytes.to_vec(); - let model = lean::cpu::CpuModel::load_from_reader(Cursor::new(bytes)) - .map_err(|e| JsError::new(&format!("failed to load model: {e}")))?; - let tokenizer = tokenizers::Tokenizer::from_bytes(tokenizer_json.as_bytes()) - .map_err(|e| JsError::new(&format!("failed to load tokenizer.json: {e}")))?; - let template = AgentChatTemplate::from_tokenizer_config_json(&tokenizer_config_json) - .map_err(|e| JsError::new(&format!("{e}")))?; - let generator = RcCpuGenerator { - model: Rc::new(model), - }; - self.agent = Some(Agent::new( - template, - tokenizer, - generator, - NoCaller, - system_prompt, - max_new_tokens as usize, - )); + max_ctx: u32, + ) -> Result { + console_error_panic_hook::set_once(); + let template = lean_chat::template_from_tokenizer_config(&tokenizer_config_json) + .map_err(|e| JsError::new(&format!("{e:#}")))?; + let generator = LeanGenerator::load(LoadOptions { + gguf: GgufSource::Bytes(gguf_bytes.to_vec()), + tokenizer: TokenizerSource::Json(tokenizer_json), + max_ctx, + backend: Backend::Auto, + }) + .await + .map_err(|e| JsError::new(&format!("failed to load model: {e}")))?; + let backend = generator.backend(); + let agent = Agent::new(template, generator, system_prompt, max_new_tokens as usize); + Ok(AgentJs { + abort: agent.abort_flag(), + agent: RefCell::new(Some(agent)), + backend, + }) + } + + /// The backend lean picked: `"webgpu"`, `"threads"` or `"single"`. + #[wasm_bindgen(getter)] + pub fn backend(&self) -> String { + self.backend.to_string() + } + + /// `on_text(text)` receives the model's raw output as it is generated, + /// every step. It must not call this agent. + #[wasm_bindgen(js_name = setOnText)] + pub fn set_on_text(&self, on_text: js_sys::Function) -> Result<(), JsError> { + let mut agent = self.take()?; + agent.set_on_text(move |text| { + let _ = on_text.call1(&JsValue::NULL, &JsValue::from_str(text)); + }); + self.put(agent); Ok(()) } + /// A handle whose `abort()` stops the running step (it resolves to an + /// `aborted` outcome). Take it before calling `start`; it stays valid + /// for every later turn. + #[wasm_bindgen(js_name = abortHandle)] + pub fn abort_handle(&self) -> AgentAbort { + AgentAbort(self.abort.clone()) + } + /// Begin a new turn. `tools_json` is a JSON array of MCP `tools/list` /// entries (`{"name", "description", "inputSchema"}`); pass `"[]"` for - /// no tools. Returns one outcome JSON object, see the module doc. - #[wasm_bindgen(js_name = start)] - pub fn start(&mut self, utterance: String, tools_json: String) -> Result { + /// no tools. Resolves to one outcome JSON object, see the module doc. + pub async fn start(&self, utterance: String, tools_json: String) -> Result { let raw_tools: Vec = serde_json::from_str(&tools_json) .map_err(|e| JsError::new(&format!("invalid tools JSON: {e}")))?; - let agent = self - .agent - .as_mut() - .ok_or_else(|| JsError::new("load() must be called first"))?; - Ok(outcome_to_json( - agent.start_from_mcp(&utterance, &raw_tools), - )) + let mut agent = self.take()?; + let outcome = agent.start_from_mcp(&utterance, &raw_tools).await; + self.put(agent); + Ok(outcome_to_json(outcome)) } /// Continue the turn with tool results. `results_json` is a JSON array /// of `[callId, result]` pairs, `result` being whatever JSON value the /// tool call returned (an MCP-shaped `{"isError": true, ...}` is read - /// as an error). Returns one outcome JSON object, same shape as + /// as an error). Resolves to one outcome JSON object, same shape as /// `start`. #[wasm_bindgen(js_name = provideToolResults)] - pub fn provide_tool_results(&mut self, results_json: String) -> Result { + pub async fn provide_tool_results(&self, results_json: String) -> Result { let results: Vec<(String, Value)> = serde_json::from_str(&results_json) .map_err(|e| JsError::new(&format!("invalid results JSON: {e}")))?; - let agent = self - .agent - .as_mut() - .ok_or_else(|| JsError::new("load() must be called first"))?; - Ok(outcome_to_json(agent.provide_tool_results(results))) + let mut agent = self.take()?; + let outcome = agent.provide_tool_results(results).await; + self.put(agent); + Ok(outcome_to_json(outcome)) } /// Abandon the current turn (e.g. after an `error` outcome) without /// reloading the model. - #[wasm_bindgen(js_name = reset)] - pub fn reset(&mut self) { - if let Some(agent) = self.agent.as_mut() { - agent.reset(); - } + pub fn reset(&self) -> Result<(), JsError> { + let mut agent = self.take()?; + agent.reset(); + self.put(agent); + Ok(()) + } +} + +impl AgentJs { + /// The agent, out of its cell for the duration of one call; an error + /// while another call holds it. + fn take(&self) -> Result, JsError> { + self.agent + .borrow_mut() + .take() + .ok_or_else(|| JsError::new("agent is busy: a step is still running")) + } + + fn put(&self, agent: Agent) { + *self.agent.borrow_mut() = Some(agent); } } diff --git a/crates/agent/Cargo.toml b/crates/agent/Cargo.toml index 5396bf6..5db5d1a 100644 --- a/crates/agent/Cargo.toml +++ b/crates/agent/Cargo.toml @@ -3,21 +3,38 @@ name = "agent" version.workspace = true edition.workspace = true +[features] +default = ["lean"] +# `LeanGenerator`: the agent's `Generator` on lean's engine. +lean = ["dep:lean"] + [dependencies] anyhow = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -minijinja = { workspace = true } -minijinja-contrib = { workspace = true } -lean = { workspace = true } +lean-chat = { workspace = true } +lean = { workspace = true, optional = true } +# Natively, lean's CPU forward pass can use every core when no GPU adapter +# is available. [target.'cfg(not(target_arch = "wasm32"))'.dependencies] -tokenizers = { workspace = true, features = ["fancy-regex"] } -pollster = "0.3" - -[target.'cfg(target_arch = "wasm32")'.dependencies] -tokenizers = { workspace = true, features = ["unstable_wasm", "fancy-regex"] } +lean = { workspace = true, optional = true, features = ["threads"] } [dev-dependencies] -serde = { workspace = true } clap = { workspace = true } +pollster = { workspace = true } + +[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] +tokenizers = { workspace = true, features = ["fancy-regex"] } + +[[example]] +name = "agent_cli" +required-features = ["lean"] + +[[example]] +name = "agent_eval" +required-features = ["lean"] + +[[test]] +name = "agent_live" +required-features = ["lean"] diff --git a/crates/agent/examples/agent_cli.rs b/crates/agent/examples/agent_cli.rs index d009262..c2787b4 100644 --- a/crates/agent/examples/agent_cli.rs +++ b/crates/agent/examples/agent_cli.rs @@ -5,18 +5,27 @@ //! cargo run -p agent --example agent_cli --release -- \ //! --model /path/to/model.gguf \ //! --tokenizer /path/to/tokenizer_dir \ -//! --tools fixtures/sonos/tools-12.json \ +//! --tools crates/agent/fixtures/sonos/tools-subset.json \ //! --utterance "pause the kitchen" //! ``` //! //! Tool calls are answered from `/results/.json` fixture -//! files next to the tools file (see `fixtures/sonos/results/`), same +//! files next to the tools file (see `crates/agent/fixtures/sonos/results/`), same //! shape the eval harness uses; there is no live MCP connection here. +//! +//! The model runs on lean's engine, which picks its backend by capability +//! (GPU when a wgpu adapter is available, else the CPU forward pass on +//! every core). The model's raw output is streamed to stderr as it is +//! generated. + +use std::io::Write; -use agent::agent::{Agent, AgentChatTemplate, FixtureCaller, StepOutcome, Tool, ToolCaller}; -use agent::backend::best_generator; +use agent::agent::{call_or_error, tools_from_mcp, Agent, StepOutcome}; +use agent::eval::FixtureCaller; +use agent::lean_generator::LeanGenerator; use anyhow::{Context, Result}; use clap::Parser; +use lean::{Backend, GgufSource, LoadOptions, TokenizerSource}; use serde_json::Value; #[derive(Parser)] @@ -37,6 +46,10 @@ struct Args { max_new_tokens: usize, #[arg(long, default_value_t = 6)] max_steps: usize, + /// KV cache size in tokens: each step's prompt plus `max_new_tokens` + /// must fit. + #[arg(long, default_value_t = 4096)] + max_ctx: u32, #[arg( long, default_value = "You are a helpful assistant with access to tools." @@ -47,67 +60,58 @@ struct Args { fn main() -> Result<()> { let args = Args::parse(); - let tokenizer = tokenizers::Tokenizer::from_file(format!("{}/tokenizer.json", args.tokenizer)) - .map_err(|e| anyhow::anyhow!("loading tokenizer.json: {e}"))?; - let cfg_json = std::fs::read_to_string(format!("{}/tokenizer_config.json", args.tokenizer)) - .context("reading tokenizer_config.json")?; - let template = AgentChatTemplate::from_tokenizer_config_json(&cfg_json)?; + let template = + lean_chat::read_chat_template(&format!("{}/tokenizer_config.json", args.tokenizer))?; let raw_tools: Vec = serde_json::from_str( &std::fs::read_to_string(&args.tools).with_context(|| format!("reading {}", args.tools))?, )?; - let tools: Vec = raw_tools - .iter() - .filter_map(|t| { - Some(Tool::from_mcp( - t.get("name")?.as_str()?, - t.get("description")?.as_str()?, - t.get("inputSchema")?.clone(), - )) - }) - .collect(); + let tools = tools_from_mcp(&raw_tools).with_context(|| format!("reading {}", args.tools))?; let results_dir = std::path::Path::new(&args.tools) .parent() .unwrap_or_else(|| std::path::Path::new(".")) .join("results"); - // One FixtureCaller to satisfy `Agent`'s generic `ToolCaller` type - // parameter (unused: this CLI drives the loop step-wise and never - // calls `Agent::run`, so `self.caller` inside `Agent` is never - // invoked) and a second, actually used, to answer tool calls below. - let mut tool_runner = FixtureCaller::new(results_dir.clone()); + let mut tool_runner = FixtureCaller::new(results_dir); - let generator = best_generator(&args.model)?; - let mut agent = Agent::new( - template, - tokenizer, - generator, - FixtureCaller::new(results_dir), - args.system_prompt, - args.max_new_tokens, - ); + let generator = pollster::block_on(LeanGenerator::load(LoadOptions { + gguf: GgufSource::Path(args.model.clone()), + tokenizer: TokenizerSource::Path(format!("{}/tokenizer.json", args.tokenizer)), + max_ctx: args.max_ctx, + backend: Backend::Auto, + }))?; + eprintln!("backend: {}", generator.backend()); + let mut agent = Agent::new(template, generator, args.system_prompt, args.max_new_tokens); agent.set_max_steps(args.max_steps); + agent.set_on_text(|text| { + eprint!("{text}"); + let _ = std::io::stderr().flush(); + }); - let mut outcome = agent.start(&args.utterance, &tools); + let mut outcome = pollster::block_on(agent.start(&args.utterance, &tools)); loop { match outcome { StepOutcome::Final { text, .. } => { + eprintln!(); println!("step: final answer"); println!(" text: {text}"); break; } StepOutcome::NeedTools { calls, .. } => { + eprintln!(); let mut results = Vec::with_capacity(calls.len()); for call in &calls { println!("step: tool call {}({})", call.name, call.arguments); - let result = tool_runner - .call(&call.name, &call.arguments) - .unwrap_or(Value::Null); + let result = call_or_error(&mut tool_runner, call); println!(" result: {result}"); results.push((call.call_id.clone(), result)); } - outcome = agent.provide_tool_results(results); + outcome = pollster::block_on(agent.provide_tool_results(results)); + } + StepOutcome::ContextOverflow { needed, available } => { + anyhow::bail!("context overflow: needed {needed} tokens, {available} available (raise --max-ctx)"); } + StepOutcome::Aborted => anyhow::bail!("aborted"), StepOutcome::Error { message, .. } => { anyhow::bail!("agent error: {message}"); } diff --git a/crates/agent/examples/agent_eval.rs b/crates/agent/examples/agent_eval.rs index a796b6e..b049a0b 100644 --- a/crates/agent/examples/agent_eval.rs +++ b/crates/agent/examples/agent_eval.rs @@ -7,58 +7,61 @@ //! LEAN_TOKENIZER_DIR_QWEN25_3B=/path/to/Qwen2.5-3B-Instruct \ //! cargo run -p agent --example agent_eval --release //! ``` +//! +//! The model runs on lean's engine, which picks its backend by capability. +//! `AGENT_EVAL_MAX_CTX` sets the KV cache size in tokens (default 4096). -use agent::agent::{Agent, AgentChatTemplate, FixtureCaller}; -use agent::backend::best_generator; +use agent::agent::Agent; use agent::eval::{load_all_tools, load_cases, render_markdown, run_all, select_tools, ToolSet}; +use agent::lean_generator::LeanGenerator; +use lean::{Backend, GgufSource, LoadOptions, TokenizerSource}; fn main() -> anyhow::Result<()> { let gguf_path = std::env::var("LEAN_GGUF_QWEN25_3B").expect("set LEAN_GGUF_QWEN25_3B"); let tokenizer_dir = std::env::var("LEAN_TOKENIZER_DIR_QWEN25_3B").expect("set LEAN_TOKENIZER_DIR_QWEN25_3B"); let repo_root = std::env::var("AGENT_EVAL_REPO_ROOT") - .unwrap_or_else(|_| concat!(env!("CARGO_MANIFEST_DIR"), "/..").to_string()); + .unwrap_or_else(|_| concat!(env!("CARGO_MANIFEST_DIR"), "/../..").to_string()); - let tokenizer = tokenizers::Tokenizer::from_file(format!("{tokenizer_dir}/tokenizer.json")) - .expect("loading tokenizer.json"); - let cfg_json = std::fs::read_to_string(format!("{tokenizer_dir}/tokenizer_config.json")) - .expect("reading tokenizer_config.json"); + let max_ctx: u32 = std::env::var("AGENT_EVAL_MAX_CTX") + .map(|s| s.parse().expect("AGENT_EVAL_MAX_CTX must be a number")) + .unwrap_or(4096); let template = - AgentChatTemplate::from_tokenizer_config_json(&cfg_json).expect("compiling chat_template"); + lean_chat::read_chat_template(&format!("{tokenizer_dir}/tokenizer_config.json"))?; let fixtures_dir = format!("{repo_root}/crates/agent/fixtures/sonos"); let results_dir = format!("{fixtures_dir}/results"); let all_tools = load_all_tools(&fixtures_dir)?; - let tools = select_tools(&all_tools, ToolSet::Twelve, &fixtures_dir)?; + let tools = select_tools(&all_tools, ToolSet::Subset, &fixtures_dir)?; let cases = load_cases(format!("{repo_root}/eval/utterances.json"))?; let system_prompt = "You are a helpful home assistant with access to Sonos speaker controls."; let max_new_tokens = 256usize; let max_steps = 6usize; - let generator = best_generator(&gguf_path)?; - let caller = FixtureCaller::new(results_dir.clone()); // replaced per case by run_all - let mut agent = Agent::new( - template, - tokenizer, - generator, - caller, - system_prompt, - max_new_tokens, - ); + let generator = pollster::block_on(LeanGenerator::load(LoadOptions { + gguf: GgufSource::Path(gguf_path), + tokenizer: TokenizerSource::Path(format!("{tokenizer_dir}/tokenizer.json")), + max_ctx, + backend: Backend::Auto, + }))?; + let backend = generator.backend(); + eprintln!("backend: {backend}"); + let mut agent = Agent::new(template, generator, system_prompt, max_new_tokens); agent.set_max_steps(max_steps); let date = std::env::var("AGENT_EVAL_DATE").unwrap_or_else(|_| "2026-10-08".to_string()); - let mut report = run_all( + let mut report = pollster::block_on(run_all( &mut agent, &tools, &cases, - ToolSet::Twelve, + ToolSet::Subset, &results_dir, "Qwen2.5-3B-Instruct-q4_0.gguf (lean, capability-selected engine)", date, - ); - report.label = "base43".to_string(); + )); + report.backend = format!("native, {backend}"); + report.label = "sonos-subset".to_string(); report.system_prompt = system_prompt.to_string(); report.max_new_tokens = max_new_tokens; report.max_steps = max_steps; diff --git a/crates/agent/fixtures/sonos/README.md b/crates/agent/fixtures/sonos/README.md index 23eb80f..a638e8c 100644 --- a/crates/agent/fixtures/sonos/README.md +++ b/crates/agent/fixtures/sonos/README.md @@ -1,17 +1,20 @@ # Sonos MCP tool schemas -Tool schemas of the Sonos MCP server (34 tools) in `tools/list` shape: name, description and inputSchema. +Tool schemas of the Sonos MCP server, 34 tools, with each tool's name, +description and input schema copied verbatim, in MCP `tools/list` shape. +They were not captured from a raw `tools/list` response, so wire-level +details (extra fields, key names) may differ from what a server sends. ## Files -- `tools.json` , all 34 tools, MCP `tools/list`-shaped: +- `tools.json`: all 34 tools, MCP `tools/list`-shaped: `[{"name", "description", "inputSchema"}, ...]`. -- `tools-12.json` , the 12-tool subset used for the tool-count experiment +- `tools-subset.json`: the 13-tool subset the eval runs with (`get_households_and_groups_and_players`, `get_now_playing`, `pause`, `resume`, `skip_to_next_track`, `skip_to_previous_track`, `set_group_volume`, `adjust_group_volume`, `get_group_volume`, - `set_group_mute`, `get_sonos_favorites`, `play_sonos_favorite`) , same - entries as in `tools.json`, verbatim. + `set_group_mute`, `get_sonos_favorites`, `play_sonos_favorite`), the + same entries as in `tools.json`, verbatim, in the order the eval uses. ## Tool name → required parameters @@ -55,19 +58,13 @@ Tool schemas of the Sonos MCP server (34 tools) in `tools/list` shape: name, des Note: `move_audio_to_players` uses `source_group_id` (not `group_id`), and several music-service play tools (`play_radio`, `play_station`, `play_sonos_favorite` via `favorite_id`, etc.) take extra optional params -(`music_service`, `artist`, `album`, `shuffle`, etc.) , see `tools.json` for +(`music_service`, `artist`, `album`, `shuffle`, etc.); see `tools.json` for the full schemas. ## `results/` -One canned JSON file per tool, `results/.json`, the bare return -value for that call against the household above (not per-argument). These -files stay in the bare/parsed shape shown by `cat`ing them , but -`agent::FixtureCaller::call` (`crates/agent/src/agent.rs`) doesn't hand -that value back as-is: it wraps it as a real MCP `tools/call` result, -`{"content":[{"type":"text","text":""}]}`, with the payload -pretty-printed (2-space indent) the way Sonos's server sends it, and (for -a fixture carrying a top-level `error` field) `isError: true` instead. -Fixtures are served MCP-shaped; the native eval exercises the same -parsing path as the browser, rather than a shortcut that only works -against a pre-parsed value. +One canned JSON file per tool, `results/.json`: the bare value that +call returns against the fixture household (one Kitchen, Living Room and +Bedroom group each), the same for any arguments. `agent::eval::FixtureCaller` +returns it as is; the agent reads a result with a top-level non-null +`error` field, or an MCP result with `isError: true`, as a tool error. diff --git a/crates/agent/fixtures/sonos/tools-12.json b/crates/agent/fixtures/sonos/tools-subset.json similarity index 100% rename from crates/agent/fixtures/sonos/tools-12.json rename to crates/agent/fixtures/sonos/tools-subset.json diff --git a/crates/agent/reference/gen_fixture_agent_tools.py b/crates/agent/reference/gen_fixture_agent_tools.py index 3c09f1b..81e1667 100644 --- a/crates/agent/reference/gen_fixture_agent_tools.py +++ b/crates/agent/reference/gen_fixture_agent_tools.py @@ -1,14 +1,16 @@ -"""Reference fixture for `agent.rs`'s Qwen2.5-Instruct tools-template port. +"""Reference fixture for the agent's Qwen2.5-Instruct tools-template prompts. Renders `tokenizer.apply_chat_template(messages, tools=..., add_generation_prompt=True)` -through HF `transformers` itself (tokenizer only -- no model forward pass -needed for a template/tokenization check), for the input cases in agent_inputs/ (01/02/04: OpenAI-shaped messages, which Qwen2.5's own template also accepts) plus one new multi-turn case this port's agent loop -actually produces (object-shaped `tool_calls[].function.arguments`, two -tool round trips). +through HF `transformers` itself (tokenizer only, no model forward pass), +for the input cases in `agent_inputs/*.json` (OpenAI-shaped messages, which +Qwen2.5's own template accepts) plus one multi-turn case generated here, +with object-shaped `tool_calls[].function.arguments` and two tool round +trips, the shape the agent loop produces. LEAN_TOKENIZER_DIR_QWEN25_3B directory with tokenizer.json + tokenizer_config.json -venv: crates/agent/reference/.venv (gitignored; tokenizer-only, no torch needed). +Run it in a venv with `transformers` installed (tokenizer only, no torch +needed). """ import json import os diff --git a/crates/agent/src/agent.rs b/crates/agent/src/agent.rs index c2bc197..a6859d9 100644 --- a/crates/agent/src/agent.rs +++ b/crates/agent/src/agent.rs @@ -1,103 +1,25 @@ -//! Tool-calling chat template, output parser and step-wise agent loop for -//! Qwen2.5-Instruct's own tools format: `...` in the system -//! turn, `{"name":..., "arguments":...}` in the +//! Tool-calling agent loop for Qwen2.5-Instruct's own tools format: +//! `...` in the system turn (rendered by the model's chat +//! template), `{"name":..., "arguments":...}` in the //! model's output. //! -//! This module has no dependency on any particular inference engine: the -//! loop is driven by the [`Generator`] trait, implemented against lean's -//! public CPU and GPU APIs in `crate::backend`. +//! Prompts are rendered with lean-chat (the model's own chat template, byte +//! for byte what transformers' `apply_chat_template` renders). The loop has +//! no dependency on any inference engine: it is driven by the +//! [`Generator`] trait, prompt text in and reply text out. use anyhow::{bail, Context, Result}; -use minijinja::value::{Kwargs, Value as MjValue}; -use minijinja::{context, Environment}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::io; +pub use crate::generator::{AbortFlag, GenerateError, Generator}; +pub use lean_chat::{ChatTemplate, Message}; + // --------------------------------------------------------------------- -// Message / tool-call wire types +// Tool definitions // --------------------------------------------------------------------- -/// A single chat message, matching the fields Qwen2.5-Instruct's tools -/// template inspects (`role`, `content`, `tool_calls`, `name`). -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct Message { - pub role: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub content: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub name: Option, -} - -impl Message { - pub fn system(content: impl Into) -> Self { - Self { - role: "system".into(), - content: Some(Value::String(content.into())), - tool_calls: None, - name: None, - } - } - - pub fn user(content: impl Into) -> Self { - Self { - role: "user".into(), - content: Some(Value::String(content.into())), - tool_calls: None, - name: None, - } - } - - pub fn assistant_text(content: impl Into) -> Self { - Self { - role: "assistant".into(), - content: Some(Value::String(content.into())), - tool_calls: None, - name: None, - } - } - - pub fn assistant_tool_calls(tool_calls: Vec) -> Self { - Self { - role: "assistant".into(), - content: None, - tool_calls: Some(tool_calls), - name: None, - } - } - - /// A `tool`-role message carrying one call's result (Qwen2.5's - /// template wraps every `tool`-role message's content in - /// `...`). - pub fn tool_result(name: impl Into, content: impl Into) -> Self { - Self { - role: "tool".into(), - content: Some(Value::String(content.into())), - tool_calls: None, - name: Some(name.into()), - } - } -} - -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolCallEntry { - #[serde(rename = "type", default = "default_function_kind")] - pub kind: String, - pub function: ToolCallFunction, -} - -fn default_function_kind() -> String { - "function".to_string() -} - -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolCallFunction { - pub name: String, - pub arguments: Value, -} - /// One entry of the `tools` list, transformers/OpenAI shape: /// `{"type": "function", "function": {"name", "description", "parameters"}}`. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -107,6 +29,10 @@ pub struct Tool { pub function: ToolFunction, } +fn default_function_kind() -> String { + "function".to_string() +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ToolFunction { pub name: String, @@ -131,12 +57,39 @@ impl Tool { }, } } + + /// Reads one MCP `tools/list` entry. `name` is required; a missing or + /// `null` `description` reads as `""` and a missing `inputSchema` as + /// an object schema with no properties (both are optional in MCP). + pub fn try_from_mcp(entry: &Value) -> Result { + let name = entry + .get("name") + .and_then(Value::as_str) + .with_context(|| format!("MCP tool entry has no string `name`: {entry}"))?; + let description = match entry.get("description") { + None | Some(Value::Null) => "", + Some(Value::String(s)) => s, + Some(other) => bail!("MCP tool `{name}`: `description` is not a string: {other}"), + }; + let input_schema = match entry.get("inputSchema") { + None | Some(Value::Null) => serde_json::json!({"type": "object", "properties": {}}), + Some(schema @ Value::Object(_)) => schema.clone(), + Some(other) => bail!("MCP tool `{name}`: `inputSchema` is not an object: {other}"), + }; + Ok(Tool::from_mcp(name, description, input_schema)) + } +} + +/// Reads an MCP `tools/list` array with [`Tool::try_from_mcp`]; the first +/// invalid entry is an error naming it. +pub fn tools_from_mcp(entries: &[Value]) -> Result> { + entries.iter().map(Tool::try_from_mcp).collect() } // --------------------------------------------------------------------- -// py_tojson: byte-exact `tojson` against Python's `json.dumps` (minijinja's -// built-in `tojson` HTML-escapes and uses different compact-mode separator -// spacing, which breaks parity against `transformers.apply_chat_template`). +// Python `json.dumps` formatting for tool results (the same separators +// the chat template's `tojson` uses for tool calls), so a tool result +// reads like the rest of the prompt. // --------------------------------------------------------------------- #[derive(Default)] @@ -173,79 +126,29 @@ pub(crate) fn to_json_compact_py(value: &T) -> Result minijinja::Error { - minijinja::Error::new( - minijinja::ErrorKind::InvalidOperation, - format!("cannot serialize to JSON: {err}"), - ) -} - -fn py_tojson(value: MjValue, kwargs: Kwargs) -> std::result::Result { - let indent: Option = kwargs.get("indent").ok(); - kwargs.assert_all_used()?; - - let rendered = match indent { - Some(indent) => { - let mut out = Vec::new(); - let indentation = " ".repeat(indent); - let formatter = serde_json::ser::PrettyFormatter::with_indent(indentation.as_bytes()); - let mut ser = serde_json::Serializer::with_formatter(&mut out, formatter); - serde::Serialize::serialize(&value, &mut ser).map_err(json_error)?; - String::from_utf8(out).expect("serde_json only emits valid UTF-8") - } - None => to_json_compact_py(&value).map_err(json_error)?, - }; - - Ok(MjValue::from_safe_string(rendered)) -} - // --------------------------------------------------------------------- -// Template: Qwen2.5-Instruct's own `chat_template`, tools-enabled. +// Prompt rendering // --------------------------------------------------------------------- -pub struct AgentChatTemplate { - env: Environment<'static>, -} - -impl AgentChatTemplate { - pub fn new(template_source: impl Into) -> Result { - let mut env = Environment::new(); - env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback); - env.add_filter("tojson", py_tojson); - env.add_template_owned("chat", template_source.into()) - .context("failed to compile chat_template")?; - Ok(Self { env }) - } - - pub fn from_tokenizer_config_json(tokenizer_config_json: &str) -> Result { - let value: Value = serde_json::from_str(tokenizer_config_json)?; - let template_source = value - .get("chat_template") - .and_then(Value::as_str) - .context("tokenizer_config.json has no string 'chat_template' field")?; - Self::new(template_source) - } - - #[cfg(not(target_arch = "wasm32"))] - pub fn from_tokenizer_config_path(path: &str) -> Result { - let text = std::fs::read_to_string(path).with_context(|| format!("reading {path}"))?; - Self::from_tokenizer_config_json(&text) - } - - /// `tools` renders as Jinja `none` when empty, matching - /// `transformers.apply_chat_template([...], tools=[])`'s behavior of - /// skipping the tool preamble. - pub fn render( - &self, - messages: &[Message], - tools: &[Tool], - add_generation_prompt: bool, - ) -> Result { - let tmpl = self.env.get_template("chat").unwrap(); - let tools_ctx = if tools.is_empty() { None } else { Some(tools) }; - let rendered = tmpl.render(context! { messages => messages, tools => tools_ctx, add_generation_prompt => add_generation_prompt })?; - Ok(rendered) - } +/// The prompt for the next model turn: `messages` and `tools` through the +/// model's chat template (lean-chat, byte-exact with transformers' +/// `apply_chat_template(messages, tools=tools, add_generation_prompt=True)`). +/// No tools renders like `tools=None`, which skips the tools preamble. +pub fn render_prompt( + template: &ChatTemplate, + messages: &[Message], + tools: &[Tool], +) -> Result { + let tools: Vec = tools + .iter() + .map(serde_json::to_value) + .collect::>()?; + let tools = if tools.is_empty() { + None + } else { + Some(tools.as_slice()) + }; + lean_chat::render(template, messages, tools, true) } // --------------------------------------------------------------------- @@ -268,7 +171,26 @@ pub enum ParsedOutput { #[derive(Debug, Clone, PartialEq, Deserialize)] struct RawToolCall { name: String, - arguments: Value, + #[serde(default)] + arguments: Option, +} + +/// A call's `arguments` as a JSON object: missing or `null` is `{}`, and a +/// string holding a JSON object (double-encoded) is that object. +fn call_arguments(name: &str, arguments: Option) -> Result { + match arguments { + None | Some(Value::Null) => Ok(Value::Object(Default::default())), + Some(args @ Value::Object(_)) => Ok(args), + Some(Value::String(s)) => match serde_json::from_str::(&s) { + Ok(args @ Value::Object(_)) => Ok(args), + _ => Err(format!( + "arguments of `{name}` are not a JSON object: {s:?}" + )), + }, + Some(other) => Err(format!( + "arguments of `{name}` are not a JSON object: {other}" + )), + } } /// Parse one turn of raw (decoded) model output. @@ -276,7 +198,7 @@ struct RawToolCall { /// - A trailing `<|im_end|>` (and whitespace before it) is stripped first. /// - Every `...` block found is parsed as one call; /// a block whose contents don't parse as `{"name", "arguments"}` is an -/// error, not a panic. +/// error, not a panic. A call with no `arguments` has empty arguments. /// - With no `` block at all, the (trimmed) text is the plain /// answer. pub fn parse_output(raw: &str) -> Result { @@ -300,17 +222,20 @@ pub fn parse_output(raw: &str) -> Result { let body = after_open[..end].trim(); let raw_call: RawToolCall = serde_json::from_str(body) .map_err(|e| format!("malformed tool_call JSON {body:?}: {e}"))?; + let arguments = call_arguments(&raw_call.name, raw_call.arguments)?; calls.push(ToolCall { name: raw_call.name, - arguments: raw_call.arguments, + arguments, }); rest = &after_open[end + CLOSE.len()..]; } Ok(ParsedOutput::ToolCalls(calls)) } -/// Above this many characters, a compacted tool result is kept in full -/// (never truncated) but is worth flagging; see `format_tool_result`. +/// The default limit on one tool result's formatted content, in bytes. The +/// agent feeds a larger result back as a tool error saying it was too +/// large, so one oversized result does not overflow the context; see +/// [`Agent::set_max_result_chars`]. pub const MAX_RESULT_CHARS: usize = 8000; /// If `value` is a string that itself parses as JSON, replace it with a @@ -338,20 +263,31 @@ fn compact_embedded_json(value: &Value) -> Value { /// Build the `tool`-role `Message` carrying a tool's result: `content` is /// a JSON-encoded string, Python-`json.dumps`-formatted, after first -/// compacting any pretty-printed JSON embedded in `result`. +/// compacting any pretty-printed JSON embedded in `result`. The content is +/// kept whole, whatever its size. pub fn format_tool_result(name: &str, result: &Value) -> Message { let compacted = compact_embedded_json(result); let content = to_json_compact_py(&compacted).unwrap_or_else(|_| result.to_string()); - if content.len() > MAX_RESULT_CHARS { - eprintln!("warning: tool result for {name} is {} chars, over MAX_RESULT_CHARS ({MAX_RESULT_CHARS}); keeping it uncut", content.len()); + tool_message(name, content) +} + +/// A `tool`-role message carrying one call's result (Qwen2.5's template +/// wraps its content in `...`). +fn tool_message(name: &str, content: String) -> Message { + Message { + name: Some(name.to_string()), + ..Message::tool(content) } - Message::tool_result(name, content) } /// If `result` is an MCP-shaped error (`isError: true` or a top-level -/// `error` field), extract a human-readable message; `None` otherwise. +/// `error` field that is not `null` or `false`), extract a human-readable +/// message; `None` otherwise. pub fn tool_error_message(result: &Value) -> Option { - if let Some(error) = result.get("error") { + if let Some(error) = result + .get("error") + .filter(|e| !matches!(e, Value::Null | Value::Bool(false))) + { let message = match error { Value::String(s) => s.clone(), Value::Object(o) => o @@ -382,7 +318,7 @@ pub fn tool_error_message(result: &Value) -> Option { pub fn format_tool_error(name: &str, message: &str) -> Message { let content = to_json_compact_py(&serde_json::json!({ "error": message })) .unwrap_or_else(|_| format!("{{\"error\": {message:?}}}")); - Message::tool_result(name, content) + tool_message(name, content) } // --------------------------------------------------------------------- @@ -393,75 +329,19 @@ const DEFAULT_MAX_STEPS: usize = 8; const DEFAULT_MAX_RETRIES: usize = 2; const MAX_CONSECUTIVE_TOOL_ERRORS: usize = 3; -/// A generic, non-committal nudge appended as a `user` message when a -/// model output is malformed/empty/unknown-tool twice in a row. +/// The first sentence of the `user` note added to a retry's prompt after +/// an invalid reply; the note then says what was wrong. const RETRY_NOTE: &str = "Respond with a tool call from the list or a final answer."; +/// The `user` note added to the prompt when the retries are spent on +/// repeated or empty tool calls. +const FINAL_ANSWER_NOTE: &str = "Give your final answer as plain text now, with no tool call."; + /// Executes one tool call and returns its JSON result. pub trait ToolCaller { fn call(&mut self, name: &str, args: &Value) -> Result; } -/// Runs the model forward from `prompt_ids`, sampling up to -/// `max_new_tokens` tokens (fewer if an eos id is hit). Greedy only. -/// -/// Implementations live in `crate::backend`, built on lean's public CPU -/// and GPU forward-pass APIs; this trait itself has no engine dependency. -pub trait Generator { - fn generate(&mut self, prompt_ids: &[u32], max_new_tokens: usize) -> Result>; -} - -/// A `Generator` that replays a scripted sequence of outputs for tests. -pub struct FixtureGenerator { - script: Vec>, - next: usize, -} - -impl FixtureGenerator { - pub fn new(script: Vec>) -> Self { - Self { script, next: 0 } - } -} - -impl Generator for FixtureGenerator { - fn generate(&mut self, _prompt_ids: &[u32], _max_new_tokens: usize) -> Result> { - let out = self.script.get(self.next).cloned().ok_or_else(|| { - anyhow::anyhow!( - "FixtureGenerator script exhausted after {} call(s)", - self.next - ) - })?; - self.next += 1; - Ok(out) - } -} - -/// A `ToolCaller` that serves canned JSON results from -/// `/.json` and records every call it saw, in order. -pub struct FixtureCaller { - results_dir: std::path::PathBuf, - pub calls_seen: Vec<(String, Value)>, -} - -impl FixtureCaller { - pub fn new(results_dir: impl Into) -> Self { - Self { - results_dir: results_dir.into(), - calls_seen: Vec::new(), - } - } -} - -impl ToolCaller for FixtureCaller { - fn call(&mut self, name: &str, args: &Value) -> Result { - self.calls_seen.push((name.to_string(), args.clone())); - let path = self.results_dir.join(format!("{name}.json")); - let text = std::fs::read_to_string(&path) - .with_context(|| format!("reading fixture result {path:?}"))?; - Ok(serde_json::from_str(&text)?) - } -} - /// One tool call the model asked for, awaiting a result via /// [`Agent::provide_tool_results`]. #[derive(Debug, Clone, PartialEq)] @@ -474,15 +354,24 @@ pub struct PendingToolCall { /// One render -> generate -> parse round of the agent loop. #[derive(Debug)] pub struct Step { - pub prompt_tokens: Vec, + /// The prompt of the accepted attempt. + pub prompt: String, + /// The model's raw reply to that prompt. pub generated_text: String, pub parsed: ParsedOutput, + /// Attempts regenerated because the reply was invalid or repeated. pub retries: usize, + /// Tool error results fed back in this step's prompt. pub tool_errors: usize, + /// Tool results fed back as "too large" errors in this step's prompt. + pub oversized_results: usize, + /// A reply repeated a call already made this turn. pub repeat_guard: bool, + /// The step ended with the model asked for a plain-text answer. pub forced_text_answer: bool, } +#[derive(Debug)] pub struct Transcript { pub steps: Vec, pub final_text: String, @@ -498,6 +387,21 @@ pub enum StepOutcome { text: String, step: Step, }, + /// The prompt no longer fits the engine's context (`needed` tokens, + /// `available` in the context). The agent is left as it was before the + /// call that returned this: after `start`, call `start` again with + /// fewer tools or a shorter utterance; after `provide_tool_results`, + /// call it again with smaller results (or errors in their place), or + /// lower `max_new_tokens` first. + ContextOverflow { + needed: usize, + available: usize, + }, + /// The turn was stopped through [`Agent::abort_flag`]. + Aborted, + /// The turn cannot go on, except after a `provide_tool_results` whose + /// call ids do not match the pending calls: the agent is left as it + /// was, and a call with the right ids continues the turn. Error { message: String, step: Option, @@ -505,58 +409,88 @@ pub enum StepOutcome { } struct AttemptOutput { - prompt_tokens: Vec, + prompt: String, generated_text: String, parsed: std::result::Result, } -pub struct Agent { - template: AgentChatTemplate, - tokenizer: tokenizers::Tokenizer, +/// Why a reply was not accepted. +enum Rejection { + /// Unparsable, or a call to a tool that is not in the list. + Invalid, + /// A `` block with no calls in it. + Empty, + /// Every call in it was already made this turn. + Repeat, +} + +/// Everything a turn changes, saved before a call so a context overflow +/// can put it back. +struct TurnState { + messages_len: usize, + step_index: usize, + pending_calls: Vec, + consecutive_tool_errors: usize, + step_tool_errors: usize, + step_oversized_results: usize, + calls_made_this_turn: Vec, +} + +/// The step-wise agent loop over a [`Generator`]. One turn answers one +/// utterance: [`Agent::start`] begins it (and a new conversation), each +/// [`StepOutcome::NeedTools`] is answered with +/// [`Agent::provide_tool_results`], until a final text answer. +pub struct Agent { + template: ChatTemplate, generator: G, - caller: C, system_prompt: String, max_new_tokens: usize, max_steps: usize, max_retries: usize, + max_result_chars: usize, + abort: AbortFlag, + on_text: Box, messages: Vec, tools: Vec, step_index: usize, pending_calls: Vec, consecutive_tool_errors: usize, + step_tool_errors: usize, + step_oversized_results: usize, calls_made_this_turn: Vec, } -impl Agent { +impl Agent { pub fn new( - template: AgentChatTemplate, - tokenizer: tokenizers::Tokenizer, + template: ChatTemplate, generator: G, - caller: C, system_prompt: impl Into, max_new_tokens: usize, ) -> Self { Self { template, - tokenizer, generator, - caller, system_prompt: system_prompt.into(), max_new_tokens, max_steps: DEFAULT_MAX_STEPS, max_retries: DEFAULT_MAX_RETRIES, + max_result_chars: MAX_RESULT_CHARS, + abort: AbortFlag::new(), + on_text: Box::new(|_| {}), messages: Vec::new(), tools: Vec::new(), step_index: 0, pending_calls: Vec::new(), consecutive_tool_errors: 0, + step_tool_errors: 0, + step_oversized_results: 0, calls_made_this_turn: Vec::new(), } } - pub fn caller(&self) -> &C { - &self.caller + pub fn generator(&self) -> &G { + &self.generator } pub fn set_max_steps(&mut self, max_steps: usize) { @@ -567,108 +501,288 @@ impl Agent { self.max_retries = max_retries; } + pub fn set_max_new_tokens(&mut self, max_new_tokens: usize) { + self.max_new_tokens = max_new_tokens; + } + + /// A tool result whose formatted content is longer than this (in + /// bytes) is fed back to the model as a tool error saying it was too + /// large. Default [`MAX_RESULT_CHARS`]. + pub fn set_max_result_chars(&mut self, max_result_chars: usize) { + self.max_result_chars = max_result_chars; + } + + /// Receives the model's raw output as it is generated, every step + /// (tool calls included, before they are parsed). + pub fn set_on_text(&mut self, on_text: impl FnMut(&str) + 'static) { + self.on_text = Box::new(on_text); + } + + /// A handle that stops the current turn: after `abort()`, the running + /// step returns [`StepOutcome::Aborted`]. [`Agent::start`] clears it. + pub fn abort_flag(&self) -> AbortFlag { + self.abort.clone() + } + pub fn reset(&mut self) { self.messages.clear(); self.tools.clear(); self.step_index = 0; self.pending_calls.clear(); self.consecutive_tool_errors = 0; + self.step_tool_errors = 0; + self.step_oversized_results = 0; self.calls_made_this_turn.clear(); } - pub fn start(&mut self, utterance: &str, tools: &[Tool]) -> StepOutcome { + /// Begins a turn for `utterance` with `tools`. Each turn is a new + /// conversation: the system prompt and this one user message. + pub async fn start(&mut self, utterance: &str, tools: &[Tool]) -> StepOutcome { + self.abort.reset(); + self.reset(); self.tools = tools.to_vec(); self.messages = vec![ Message::system(&self.system_prompt), Message::user(utterance), ]; - self.step_index = 0; - self.pending_calls.clear(); - self.consecutive_tool_errors = 0; - self.calls_made_this_turn.clear(); + let outcome = self.step_inner().await; + if matches!(outcome, StepOutcome::ContextOverflow { .. }) { + self.reset(); + } + outcome + } - self.step_inner() + /// [`Agent::start`] with the tools of an MCP `tools/list` array, read + /// with [`tools_from_mcp`]; an invalid entry is an `Error` outcome. + pub async fn start_from_mcp(&mut self, utterance: &str, raw_tools: &[Value]) -> StepOutcome { + match tools_from_mcp(raw_tools) { + Ok(tools) => self.start(utterance, &tools).await, + Err(e) => StepOutcome::Error { + message: format!("reading tools: {e:#}"), + step: None, + }, + } } - pub fn start_from_mcp(&mut self, utterance: &str, raw_tools: &[Value]) -> StepOutcome { - let tools: Vec = raw_tools - .iter() - .filter_map(|t| { - Some(Tool::from_mcp( - t.get("name")?.as_str()?, - t.get("description")?.as_str()?, - t.get("inputSchema")?.clone(), - )) - }) - .collect(); - self.start(utterance, &tools) - } - - /// Continue the current turn with the results of the tool calls from - /// the most recent `NeedTools` outcome, keyed by `call_id`. - pub fn provide_tool_results(&mut self, results: Vec<(String, Value)>) -> StepOutcome { - for (call_id, result) in results { - if let Some(pending) = self.pending_calls.iter().find(|c| c.call_id == call_id) { - if let Some(error_message) = tool_error_message(&result) { - self.consecutive_tool_errors += 1; - if self.consecutive_tool_errors > MAX_CONSECUTIVE_TOOL_ERRORS { - self.pending_calls.clear(); - return StepOutcome::Error { - message: format!( - "{} consecutive tool errors (giving up after {}): {error_message}", - self.consecutive_tool_errors, MAX_CONSECUTIVE_TOOL_ERRORS - ), - step: None, - }; - } - self.messages - .push(format_tool_error(&pending.name, &error_message)); + /// Continues the turn with the results of the calls from the most + /// recent `NeedTools` outcome, one per call, keyed by `call_id`. An + /// unknown, duplicate or missing id is an `Error` outcome that leaves + /// the agent as it was. + pub async fn provide_tool_results(&mut self, results: Vec<(String, Value)>) -> StepOutcome { + let ordered = match self.results_in_call_order(results) { + Ok(ordered) => ordered, + Err(message) => { + return StepOutcome::Error { + message, + step: None, + } + } + }; + let saved = self.save_turn(); + + let pending = std::mem::take(&mut self.pending_calls); + self.step_tool_errors = 0; + self.step_oversized_results = 0; + for (call, result) in pending.iter().zip(ordered) { + if let Some(error_message) = tool_error_message(&result) { + self.consecutive_tool_errors += 1; + self.step_tool_errors += 1; + if self.consecutive_tool_errors > MAX_CONSECUTIVE_TOOL_ERRORS { + return StepOutcome::Error { + message: format!( + "{} consecutive tool errors (giving up after {}): {error_message}", + self.consecutive_tool_errors, MAX_CONSECUTIVE_TOOL_ERRORS + ), + step: None, + }; + } + // A failed call may be retried as is. + let failed = ToolCall { + name: call.name.clone(), + arguments: call.arguments.clone(), + }; + self.calls_made_this_turn.retain(|c| *c != failed); + self.messages + .push(format_tool_error(&call.name, &error_message)); + } else { + self.consecutive_tool_errors = 0; + let message = format_tool_result(&call.name, &result); + if message.content.len() > self.max_result_chars { + self.step_oversized_results += 1; + self.messages.push(format_tool_error( + &call.name, + &format!( + "result too large: {} characters, the limit is {}", + message.content.len(), + self.max_result_chars + ), + )); } else { - self.consecutive_tool_errors = 0; - self.messages - .push(format_tool_result(&pending.name, &result)); + self.messages.push(message); } } } - self.pending_calls.clear(); - self.step_inner() - } - - fn generate_attempt(&mut self) -> std::result::Result { - let rendered = self - .template - .render(&self.messages, &self.tools, true) - .map_err(|e| format!("rendering prompt: {e:#}"))?; - let encoding = self - .tokenizer - .encode(rendered, false) - .map_err(|e| format!("tokenizer encode failed: {e}"))?; - let prompt_tokens = encoding.get_ids().to_vec(); - let out_ids = self - .generator - .generate(&prompt_tokens, self.max_new_tokens) - .map_err(|e| format!("generation failed: {e:#}"))?; + + let outcome = self.step_inner().await; + if matches!(outcome, StepOutcome::ContextOverflow { .. }) { + self.restore_turn(saved); + } + outcome + } + + /// `results` in the order of the pending calls, or why they do not + /// match them one to one. + fn results_in_call_order( + &self, + mut results: Vec<(String, Value)>, + ) -> std::result::Result, String> { + let mut ordered = Vec::with_capacity(self.pending_calls.len()); + for call in &self.pending_calls { + let index = results + .iter() + .position(|(id, _)| *id == call.call_id) + .ok_or_else(|| { + format!("no result for tool call `{}` ({})", call.call_id, call.name) + })?; + ordered.push(results.swap_remove(index).1); + } + if let Some((id, _)) = results.first() { + return Err(format!( + "no pending tool call has id `{id}`, or it has more than one result" + )); + } + Ok(ordered) + } + + fn save_turn(&self) -> TurnState { + TurnState { + messages_len: self.messages.len(), + step_index: self.step_index, + pending_calls: self.pending_calls.clone(), + consecutive_tool_errors: self.consecutive_tool_errors, + step_tool_errors: self.step_tool_errors, + step_oversized_results: self.step_oversized_results, + calls_made_this_turn: self.calls_made_this_turn.clone(), + } + } + + fn restore_turn(&mut self, saved: TurnState) { + self.messages.truncate(saved.messages_len); + self.step_index = saved.step_index; + self.pending_calls = saved.pending_calls; + self.consecutive_tool_errors = saved.consecutive_tool_errors; + self.step_tool_errors = saved.step_tool_errors; + self.step_oversized_results = saved.step_oversized_results; + self.calls_made_this_turn = saved.calls_made_this_turn; + } + + /// Renders the conversation plus `notes` (as `user` messages that are + /// not kept in the history), generates and parses one reply. `Err` is + /// the outcome to return as is (render or engine failure, overflow, + /// abort). + async fn generate_attempt( + &mut self, + notes: &[String], + ) -> std::result::Result> { + let rendered = if notes.is_empty() { + render_prompt(&self.template, &self.messages, &self.tools) + } else { + let mut messages = self.messages.clone(); + messages.extend(notes.iter().map(Message::user)); + render_prompt(&self.template, &messages, &self.tools) + }; + let prompt = rendered.map_err(|e| StepOutcome::Error { + message: format!("rendering prompt: {e:#}"), + step: None, + })?; let generated_text = self - .tokenizer - .decode(&out_ids, true) - .map_err(|e| format!("tokenizer decode failed: {e}"))?; + .generator + .generate( + &prompt, + self.max_new_tokens, + &mut *self.on_text, + &self.abort, + ) + .await + .map_err(|e| { + Box::new(match e { + GenerateError::ContextOverflow { needed, available } => { + StepOutcome::ContextOverflow { needed, available } + } + GenerateError::Aborted => StepOutcome::Aborted, + GenerateError::Other(e) => StepOutcome::Error { + message: format!("generation failed: {e:#}"), + step: None, + }, + }) + })?; let parsed = parse_output(&generated_text); Ok(AttemptOutput { - prompt_tokens, + prompt, generated_text, parsed, }) } + /// `parsed` if it can be accepted; otherwise why not, with a short + /// reason for the model and the full one for the caller. + fn check( + &self, + parsed: std::result::Result, + ) -> std::result::Result { + let calls = match parsed { + Err(e) => { + return Err(( + Rejection::Invalid, + "Your last reply was not a valid tool call.".to_string(), + e, + )) + } + Ok(ParsedOutput::ToolCalls(calls)) => calls, + Ok(text) => return Ok(text), + }; + if calls.is_empty() { + return Err(( + Rejection::Empty, + "Your last reply had no tool call in it.".to_string(), + "empty tool call list".to_string(), + )); + } + if let Some(unknown) = calls + .iter() + .find(|c| !self.tools.iter().any(|t| t.function.name == c.name)) + { + return Err(( + Rejection::Invalid, + format!("`{}` is not one of the tools.", unknown.name), + format!("call to unknown tool `{}`", unknown.name), + )); + } + if let Some(repeat) = calls.iter().find(|c| self.calls_made_this_turn.contains(c)) { + return Err(( + Rejection::Repeat, + format!( + "`{}` was already called with these arguments; use its result.", + repeat.name + ), + format!("repeated call {}({})", repeat.name, repeat.arguments), + )); + } + Ok(ParsedOutput::ToolCalls(calls)) + } + /// One render -> generate -> parse round using `self.messages`/`self.tools`. /// - /// Malformed-output retry policy: an empty tool-call array, unparsable - /// output, or a call naming a tool outside `self.tools` does not get - /// appended to history; the agent regenerates, nudging with - /// [`RETRY_NOTE`] once `self.max_retries` is exhausted. A call - /// identical (same name+args) to one already made earlier this turn is - /// never re-executed (the repeated-call loop guard). - fn step_inner(&mut self) -> StepOutcome { + /// Retry policy: an unparsable reply, an empty tool-call block, a call + /// to a tool outside `self.tools`, or a call identical (same name and + /// arguments) to one already made this turn whose result was not an + /// error is not accepted. The agent regenerates up to + /// `self.max_retries` times, each time with a `user` note saying what + /// was wrong; the notes are part of that step's prompts only, never of + /// the history. When the retries are spent on repeated or empty calls, + /// the model is asked once more for a plain-text answer; otherwise the + /// step is an error carrying the last reason. + async fn step_inner(&mut self) -> StepOutcome { if self.step_index >= self.max_steps { return StepOutcome::Error { message: format!( @@ -680,145 +794,139 @@ impl Agent { } self.step_index += 1; - let mut retries = 0usize; + let mut notes: Vec = Vec::new(); let mut repeat_guard = false; loop { - let attempt = match self.generate_attempt() { + let attempt = match self.generate_attempt(¬es).await { Ok(a) => a, - Err(message) => { - return StepOutcome::Error { - message, - step: None, - } - } - }; - - let empty_calls = - matches!(&attempt.parsed, Ok(ParsedOutput::ToolCalls(calls)) if calls.is_empty()); - let invalid = match &attempt.parsed { - Err(_) => true, - Ok(ParsedOutput::ToolCalls(calls)) => { - calls.is_empty() - || calls - .iter() - .any(|c| !self.tools.iter().any(|t| t.function.name == c.name)) - } - Ok(ParsedOutput::Text(_)) => false, + Err(outcome) => return *outcome, }; - let is_repeat = !invalid - && matches!(&attempt.parsed, Ok(ParsedOutput::ToolCalls(calls)) if calls.iter().any(|c| self.calls_made_this_turn.contains(c))); - if is_repeat { - repeat_guard = true; - } - - if invalid || is_repeat { - if retries < self.max_retries { - retries += 1; - self.messages.push(Message::user(RETRY_NOTE)); - continue; - } - if is_repeat || empty_calls { - return self.force_final_answer(repeat_guard); - } - return StepOutcome::Error { - message: "model produced no valid call".to_string(), - step: None, - }; - } - let AttemptOutput { - prompt_tokens, + prompt, generated_text, parsed, } = attempt; - let parsed = parsed.expect("checked valid above"); - let tool_errors = self.consecutive_tool_errors; - - return match parsed { - ParsedOutput::ToolCalls(calls) => { - let pending = pending_calls_from(&calls); - let entries = tool_call_entries(&pending); - self.messages.push(Message::assistant_tool_calls(entries)); - self.pending_calls = pending.clone(); - self.calls_made_this_turn.extend(calls.iter().cloned()); - + match self.check(parsed) { + Ok(parsed) => { let step = Step { - prompt_tokens, + prompt, generated_text, - parsed: ParsedOutput::ToolCalls(calls), - retries, - tool_errors, + parsed, + retries: notes.len(), + tool_errors: self.step_tool_errors, + oversized_results: self.step_oversized_results, repeat_guard, forced_text_answer: false, }; - StepOutcome::NeedTools { - calls: pending, - step, - } + return self.accept(step); } - ParsedOutput::Text(text) => { - self.messages.push(Message::assistant_text(&text)); - let step = Step { - prompt_tokens, - generated_text, - parsed: ParsedOutput::Text(text.clone()), - retries, - tool_errors, - repeat_guard, - forced_text_answer: false, + Err((rejection, note, detail)) => { + if matches!(rejection, Rejection::Repeat) { + repeat_guard = true; + } + if notes.len() < self.max_retries { + notes.push(format!("{note} {RETRY_NOTE}")); + continue; + } + if matches!(rejection, Rejection::Repeat | Rejection::Empty) { + return self.force_final_answer(notes, repeat_guard).await; + } + return StepOutcome::Error { + message: format!( + "model produced no valid call after {} retries: {detail}", + notes.len() + ), + step: None, }; - StepOutcome::Final { text, step } } - }; + } } } - /// Last resort when the retry budget is exhausted on a repeated or - /// empty tool call: nudge once more and accept whatever text comes - /// back (falling back to the raw output if the model still emits a - /// tool call, rather than looping forever). - fn force_final_answer(&mut self, repeat_guard: bool) -> StepOutcome { - self.messages.push(Message::user( - "Give your final answer as plain text now, with no tool call.", - )); - let attempt = match self.generate_attempt() { - Ok(a) => a, - Err(message) => { - return StepOutcome::Error { - message, - step: None, + fn accept(&mut self, step: Step) -> StepOutcome { + match &step.parsed { + ParsedOutput::ToolCalls(calls) => { + let pending = pending_calls_from(calls); + self.messages + .push(Message::assistant_tool_calls("", chat_tool_calls(&pending))); + self.pending_calls = pending.clone(); + self.calls_made_this_turn.extend(calls.iter().cloned()); + StepOutcome::NeedTools { + calls: pending, + step, } } + ParsedOutput::Text(text) => { + let text = text.clone(); + self.messages.push(Message::assistant(&text)); + StepOutcome::Final { text, step } + } + } + } + + /// Last resort when the retries are spent on repeated or empty tool + /// calls: one more generation, in the same step, with + /// [`FINAL_ANSWER_NOTE`] added. A text reply is the final answer; a + /// reply that still calls a tool is an error (its markup is never + /// returned as the answer). + async fn force_final_answer( + &mut self, + mut notes: Vec, + repeat_guard: bool, + ) -> StepOutcome { + let retries = notes.len(); + notes.push(FINAL_ANSWER_NOTE.to_string()); + let attempt = match self.generate_attempt(¬es).await { + Ok(a) => a, + Err(outcome) => return *outcome, }; let AttemptOutput { - prompt_tokens, + prompt, generated_text, parsed, } = attempt; - let text = match parsed { - Ok(ParsedOutput::Text(t)) => t, - _ => generated_text.clone(), + let parsed = match parsed { + Ok(parsed) => parsed, + Err(e) => { + return StepOutcome::Error { + message: format!("asked for a plain-text answer, the model replied: {e}"), + step: None, + } + } }; - self.messages.push(Message::assistant_text(&text)); + let is_text = matches!(parsed, ParsedOutput::Text(_)); let step = Step { - prompt_tokens, + prompt, generated_text, - parsed: ParsedOutput::Text(text.clone()), - retries: self.max_retries, - tool_errors: self.consecutive_tool_errors, + parsed, + retries, + tool_errors: self.step_tool_errors, + oversized_results: self.step_oversized_results, repeat_guard, forced_text_answer: true, }; - StepOutcome::Final { text, step } + if is_text { + self.accept(step) + } else { + StepOutcome::Error { + message: "asked for a plain-text answer, the model called a tool again".to_string(), + step: Some(step), + } + } } - /// Synchronous convenience for native callers (tests, eval, a CLI) - /// with a synchronous `ToolCaller`: runs the whole turn to completion. - pub fn run(&mut self, utterance: &str, tools: &[Tool], max_steps: usize) -> Result { - self.max_steps = max_steps; - let mut steps = Vec::with_capacity(max_steps); - let mut outcome = self.start(utterance, tools); + /// Runs a whole turn with a synchronous `caller` (tests, eval, a CLI), + /// up to the agent's `max_steps`. A caller error is fed to the model + /// as a tool error, like an MCP error result. + pub async fn run( + &mut self, + utterance: &str, + tools: &[Tool], + caller: &mut C, + ) -> Result { + let mut steps = Vec::new(); + let mut outcome = self.start(utterance, tools).await; loop { match outcome { @@ -831,24 +939,30 @@ impl Agent { } StepOutcome::NeedTools { calls, step } => { steps.push(step); - let mut results = Vec::with_capacity(calls.len()); - for call in &calls { - let result = self.caller.call(&call.name, &call.arguments)?; - results.push((call.call_id.clone(), result)); - } - outcome = self.provide_tool_results(results); + let results = calls + .iter() + .map(|call| (call.call_id.clone(), call_or_error(caller, call))) + .collect(); + outcome = self.provide_tool_results(results).await; } - StepOutcome::Error { message, step } => { - if let Some(step) = step { - steps.push(step); - } - bail!("{message}"); + StepOutcome::ContextOverflow { needed, available } => { + bail!("context overflow: needed {needed} tokens, {available} available"); } + StepOutcome::Aborted => bail!("aborted"), + StepOutcome::Error { message, .. } => bail!("{message}"), } } } } +/// `caller`'s result for `call`, or its error as an `{"error": ...}` +/// result. +pub fn call_or_error(caller: &mut C, call: &PendingToolCall) -> Value { + caller + .call(&call.name, &call.arguments) + .unwrap_or_else(|e| serde_json::json!({ "error": format!("{e:#}") })) +} + fn pending_calls_from(calls: &[ToolCall]) -> Vec { calls .iter() @@ -861,22 +975,17 @@ fn pending_calls_from(calls: &[ToolCall]) -> Vec { .collect() } -fn tool_call_entries(pending: &[PendingToolCall]) -> Vec { +fn chat_tool_calls(pending: &[PendingToolCall]) -> Vec { pending .iter() - .map(|c| ToolCallEntry { - kind: "function".to_string(), - function: ToolCallFunction { - name: c.name.clone(), - arguments: c.arguments.clone(), - }, - }) + .map(|c| lean_chat::ToolCall::new(c.name.clone(), c.arguments.clone())) .collect() } #[cfg(test)] mod tests { use super::*; + use crate::test_support::FakeGenerator; /// Qwen2.5-Instruct's tools-branch template, copied from the model's /// own `tokenizer_config.json`. @@ -893,9 +1002,9 @@ mod tests { #[test] fn renders_tools_preamble_and_single_user_turn() { - let tmpl = AgentChatTemplate::new(QWEN25_TOOLS_TEMPLATE).unwrap(); + let tmpl = ChatTemplate::new(QWEN25_TOOLS_TEMPLATE); let messages = vec![Message::user("What's the weather in Paris?")]; - let out = tmpl.render(&messages, &[weather_tool()], true).unwrap(); + let out = render_prompt(&tmpl, &messages, &[weather_tool()]).unwrap(); assert!(out.starts_with("<|im_start|>system\n")); assert!(out.contains( "\n{\"type\": \"function\", \"function\": {\"name\": \"get_weather\"" @@ -906,22 +1015,22 @@ mod tests { #[test] fn renders_multiturn_with_tool_call_and_result() { - let tmpl = AgentChatTemplate::new(QWEN25_TOOLS_TEMPLATE).unwrap(); + let tmpl = ChatTemplate::new(QWEN25_TOOLS_TEMPLATE); let messages = vec![ Message::user("What's the weather in Paris?"), - Message::assistant_tool_calls(vec![ToolCallEntry { - kind: "function".into(), - function: ToolCallFunction { - name: "get_weather".into(), - arguments: serde_json::json!({"city": "Paris"}), - }, - }]), + Message::assistant_tool_calls( + "", + vec![lean_chat::ToolCall::new( + "get_weather", + serde_json::json!({"city": "Paris"}), + )], + ), format_tool_result( "get_weather", &serde_json::json!({"temp_c": 14, "condition": "cloudy"}), ), ]; - let out = tmpl.render(&messages, &[weather_tool()], true).unwrap(); + let out = render_prompt(&tmpl, &messages, &[weather_tool()]).unwrap(); assert!(out.contains("\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n")); assert!(out.contains( "\n{\"temp_c\": 14, \"condition\": \"cloudy\"}\n" @@ -973,18 +1082,71 @@ mod tests { assert!(parse_output("\nnot json\n").is_err()); } + #[test] + fn call_without_arguments_has_empty_arguments() { + let empty = serde_json::json!({}); + for raw in [ + "\n{\"name\": \"a\"}\n", + "\n{\"name\": \"a\", \"arguments\": null}\n", + "\n{\"name\": \"a\", \"arguments\": \"{}\"}\n", + ] { + assert_eq!( + parse_output(raw).unwrap(), + ParsedOutput::ToolCalls(vec![ToolCall { + name: "a".into(), + arguments: empty.clone() + }]), + "{raw}" + ); + } + assert!( + parse_output("\n{\"name\": \"a\", \"arguments\": 3}\n").is_err() + ); + } + + #[test] + fn mcp_tools_without_description_or_schema_are_accepted() { + let tools = tools_from_mcp(&[ + serde_json::json!({"name": "ping"}), + serde_json::json!({"name": "echo", "description": null, "inputSchema": {"type": "object"}}), + ]) + .unwrap(); + assert_eq!(tools[0].function.name, "ping"); + assert_eq!(tools[0].function.description, ""); + assert_eq!( + tools[0].function.parameters, + serde_json::json!({"type": "object", "properties": {}}) + ); + assert_eq!(tools[1].function.description, ""); + assert!(tools_from_mcp(&[serde_json::json!({"description": "no name"})]).is_err()); + assert!( + tools_from_mcp(&[serde_json::json!({"name": "x", "inputSchema": "nope"})]).is_err() + ); + } + + #[test] + fn null_or_false_error_field_is_not_an_error() { + assert_eq!( + tool_error_message(&serde_json::json!({"error": null, "ok": 1})), + None + ); + assert_eq!( + tool_error_message(&serde_json::json!({"error": false})), + None + ); + assert_eq!( + tool_error_message(&serde_json::json!({"error": "no such group"})), + Some("no such group".to_string()) + ); + } + #[test] fn format_tool_result_compacts_embedded_json_string() { let msg = format_tool_result( "get_weather", &serde_json::json!({"body": "{\n \"temp_c\": 14\n}"}), ); - assert_eq!( - msg.content, - Some(Value::String( - "{\"body\": \"{\\\"temp_c\\\": 14}\"}".to_string() - )) - ); + assert_eq!(msg.content, "{\"body\": \"{\\\"temp_c\\\": 14}\"}"); } #[test] @@ -1003,48 +1165,345 @@ mod tests { } } - fn id(tokenizer: &tokenizers::Tokenizer, s: &str) -> Vec { - tokenizer.encode(s, false).unwrap().get_ids().to_vec() + struct FailingCaller; + impl ToolCaller for FailingCaller { + fn call(&mut self, name: &str, _args: &Value) -> Result { + bail!("{name} is offline") + } } - /// End-to-end step-wise loop against a scripted `FixtureGenerator`: the - /// model calls a tool, gets a result, then answers in text. - #[test] - fn agent_loop_calls_tool_then_answers() { - // No pre-tokenizer: each whole turn is one atomic vocab entry, so - // encode/decode round-trips it exactly. This test only exercises - // the loop's control flow, not real subword tokenization. - let tool_call_text = "\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n"; - let answer_text = "It is cloudy."; - let vocab = [(tool_call_text, 0u32), (answer_text, 1), ("", 2)]; - let model = tokenizers::models::wordlevel::WordLevel::builder() - .vocab(vocab.iter().map(|(s, i)| (s.to_string(), *i)).collect()) - .unk_token("".into()) - .build() - .unwrap(); - let tokenizer = tokenizers::Tokenizer::new(model); - - let tool_call_ids = id(&tokenizer, tool_call_text); - let answer_ids = id(&tokenizer, answer_text); - - let generator = FixtureGenerator::new(vec![tool_call_ids, answer_ids]); - let template = AgentChatTemplate::new(QWEN25_TOOLS_TEMPLATE).unwrap(); - let mut agent = Agent::new( - template, - tokenizer, + const PARIS_CALL: &str = "\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n"; + const LYON_CALL: &str = "\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Lyon\"}}\n"; + + fn agent_with(generator: FakeGenerator) -> Agent { + Agent::new( + ChatTemplate::new(QWEN25_TOOLS_TEMPLATE), generator, - EchoCaller, "You are a helpful assistant.", 64, - ); + ) + } + + fn start(agent: &mut Agent) -> StepOutcome { + pollster::block_on(agent.start("What's the weather in Paris?", &[weather_tool()])) + } + + fn provide(agent: &mut Agent, result: Value) -> StepOutcome { + pollster::block_on(agent.provide_tool_results(vec![("call_0".to_string(), result)])) + } + + fn sunny() -> Value { + serde_json::json!({"temp_c": 21, "condition": "sunny"}) + } - let transcript = agent - .run("What's the weather in Paris?", &[weather_tool()], 4) - .unwrap(); + /// End-to-end step-wise loop: the model calls a tool, gets a result, + /// then answers in text. + #[test] + fn agent_loop_calls_tool_then_answers() { + let answer_text = "It is cloudy."; + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL, answer_text])); + agent.set_max_steps(4); + + let transcript = pollster::block_on(agent.run( + "What's the weather in Paris?", + &[weather_tool()], + &mut EchoCaller, + )) + .unwrap(); assert_eq!(transcript.steps.len(), 2); assert!( matches!(&transcript.steps[0].parsed, ParsedOutput::ToolCalls(calls) if calls.len() == 1 && calls[0].name == "get_weather") ); assert!(matches!(&transcript.steps[1].parsed, ParsedOutput::Text(_))); + assert!(transcript.steps[1].prompt.contains( + "\n{\"called\": \"get_weather\", \"with\": {\"city\": \"Paris\"}}\n" + )); + assert_eq!(transcript.final_text, answer_text); + } + + #[test] + fn run_feeds_a_caller_error_to_the_model() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL, "It is offline."])); + let transcript = pollster::block_on(agent.run( + "What's the weather in Paris?", + &[weather_tool()], + &mut FailingCaller, + )) + .unwrap(); + assert_eq!(transcript.steps[1].tool_errors, 1); + assert!(transcript.steps[1].prompt.contains( + "\n{\"error\": \"get_weather is offline\"}\n" + )); + } + + #[test] + fn streams_text_and_abort_stops_the_turn() { + let mut agent = agent_with(FakeGenerator::script(&["Hello."])); + let streamed = std::rc::Rc::new(std::cell::RefCell::new(String::new())); + let sink = streamed.clone(); + agent.set_on_text(move |t| sink.borrow_mut().push_str(t)); + let first = pollster::block_on(agent.start("hi", &[weather_tool()])); + assert!(matches!(first, StepOutcome::Final { .. }), "{first:?}"); + assert_eq!(*streamed.borrow(), "Hello."); + + agent.abort_flag().abort(); + let outcome = pollster::block_on(agent.provide_tool_results(Vec::new())); + assert!(matches!(outcome, StepOutcome::Aborted), "{outcome:?}"); + } + + #[test] + fn abort_during_a_generation_stops_the_step() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL])); + let flag = agent.abort_flag(); + agent.set_on_text(move |_| flag.abort()); + assert!(matches!(start(&mut agent), StepOutcome::Aborted)); + } + + #[test] + fn retry_notes_are_not_kept_in_history() { + let generator = FakeGenerator::script(&[ + "\nnot json\n", + PARIS_CALL, + "It is sunny.", + ]); + let prompts = generator.prompts(); + let mut agent = agent_with(generator); + + let outcome = start(&mut agent); + let StepOutcome::NeedTools { step, .. } = outcome else { + panic!("{outcome:?}"); + }; + assert_eq!(step.retries, 1); + assert!(prompts.borrow()[1].contains("Your last reply was not a valid tool call.")); + + let outcome = provide(&mut agent, sunny()); + assert!(matches!(outcome, StepOutcome::Final { .. }), "{outcome:?}"); + let prompts = prompts.borrow(); + assert!(!prompts[2].contains(RETRY_NOTE)); + // The next step's prompt extends the first one, so an engine can + // reuse its KV cache for all of it. + assert!(prompts[2].starts_with(&prompts[0])); + } + + #[test] + fn spent_retries_report_the_last_reason() { + let mut agent = agent_with(FakeGenerator::script(&[ + "\nnot json\n", + "\n{\"name\": \"get_weather\"\n", + "\n{\"name\": \"nope\", \"arguments\": {}}\n", + ])); + let outcome = start(&mut agent); + let StepOutcome::Error { message, .. } = outcome else { + panic!("{outcome:?}"); + }; + assert!(message.contains("after 2 retries"), "{message}"); + assert!(message.contains("unknown tool `nope`"), "{message}"); + } + + #[test] + fn repeated_call_ends_in_a_forced_text_answer() { + let generator = FakeGenerator::script(&[ + PARIS_CALL, + PARIS_CALL, + PARIS_CALL, + PARIS_CALL, + "It is sunny.", + ]); + let prompts = generator.prompts(); + let mut agent = agent_with(generator); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + + let outcome = provide(&mut agent, sunny()); + let StepOutcome::Final { text, step } = outcome else { + panic!("{outcome:?}"); + }; + assert_eq!(text, "It is sunny."); + assert!(step.forced_text_answer && step.repeat_guard); + assert_eq!(step.retries, 2); + assert!(prompts.borrow()[4].contains(FINAL_ANSWER_NOTE)); + + // The note is not kept: the history ends with the answer. + let next = pollster::block_on(agent.provide_tool_results(Vec::new())); + assert!(matches!(next, StepOutcome::Error { .. })); + let last = prompts.borrow().last().cloned().unwrap(); + assert!(!last.contains(FINAL_ANSWER_NOTE)); + } + + #[test] + fn forced_answer_that_calls_a_tool_is_an_error_not_markup() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL; 5])); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + let outcome = provide(&mut agent, sunny()); + let StepOutcome::Error { message, step } = outcome else { + panic!("{outcome:?}"); + }; + assert!(message.contains("called a tool again"), "{message}"); + assert!(step.is_some_and(|s| s.forced_text_answer)); + } + + #[test] + fn a_failed_call_may_be_retried_as_is() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL, PARIS_CALL])); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + let outcome = provide(&mut agent, serde_json::json!({"error": "timeout"})); + let StepOutcome::NeedTools { step, .. } = outcome else { + panic!("{outcome:?}"); + }; + assert!(!step.repeat_guard); + assert_eq!(step.retries, 0); + assert_eq!(step.tool_errors, 1); + } + + #[test] + fn tool_errors_count_per_step() { + let mut agent = agent_with(FakeGenerator::script(&[ + PARIS_CALL, + LYON_CALL, + "\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Nice\"}}\n", + ])); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + let error = serde_json::json!({"error": "timeout"}); + let StepOutcome::NeedTools { step, .. } = provide(&mut agent, error.clone()) else { + panic!() + }; + assert_eq!(step.tool_errors, 1); + let StepOutcome::NeedTools { step, .. } = provide(&mut agent, error) else { + panic!() + }; + assert_eq!(step.tool_errors, 1); + } + + #[test] + fn too_many_consecutive_tool_errors_end_the_turn() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL; 4])); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + let error = + serde_json::json!({"isError": true, "content": [{"type": "text", "text": "bad city"}]}); + for _ in 0..MAX_CONSECUTIVE_TOOL_ERRORS { + let outcome = provide(&mut agent, error.clone()); + assert!( + matches!(outcome, StepOutcome::NeedTools { .. }), + "{outcome:?}" + ); + } + let outcome = provide(&mut agent, error); + let StepOutcome::Error { message, .. } = outcome else { + panic!("{outcome:?}"); + }; + assert!(message.contains("4 consecutive tool errors"), "{message}"); + } + + #[test] + fn max_steps_ends_the_turn() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL])); + agent.set_max_steps(1); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + let outcome = provide(&mut agent, sunny()); + assert!( + matches!(&outcome, StepOutcome::Error { message, .. } if message.contains("max_steps")), + "{outcome:?}" + ); + } + + #[test] + fn generator_failure_is_an_error_outcome() { + let mut agent = agent_with(FakeGenerator::script(&[])); + let outcome = start(&mut agent); + assert!( + matches!(&outcome, StepOutcome::Error { message, .. } if message.contains("generation failed")), + "{outcome:?}" + ); + } + + #[test] + fn context_overflow_mid_turn_is_recoverable() { + const LIMIT: usize = 4000; + let mut replies = vec![PARIS_CALL, "It is sunny."].into_iter(); + let generator = FakeGenerator::new(move |prompt| { + if prompt.len() > LIMIT { + return Err(GenerateError::ContextOverflow { + needed: prompt.len(), + available: LIMIT, + }); + } + Ok(replies.next().unwrap().to_string()) + }); + let prompts = generator.prompts(); + let mut agent = agent_with(generator); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + + let big = serde_json::json!({"report": "x".repeat(LIMIT)}); + let outcome = provide(&mut agent, big); + assert!( + matches!( + outcome, + StepOutcome::ContextOverflow { + available: LIMIT, + .. + } + ), + "{outcome:?}" + ); + + // The same call takes a smaller result. + let outcome = provide(&mut agent, sunny()); + let StepOutcome::Final { text, step } = outcome else { + panic!("{outcome:?}"); + }; + assert_eq!(text, "It is sunny."); + assert!(!step.prompt.contains("xxxx")); + assert_eq!(prompts.borrow().len(), 3); + } + + #[test] + fn oversized_result_is_fed_back_as_an_error() { + let generator = FakeGenerator::script(&[PARIS_CALL, "It is too much."]); + let mut agent = agent_with(generator); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + let big = serde_json::json!({"report": "x".repeat(MAX_RESULT_CHARS)}); + let StepOutcome::Final { step, .. } = provide(&mut agent, big) else { + panic!() + }; + assert_eq!(step.oversized_results, 1); + assert!(step.prompt.contains("result too large")); + assert!(!step.prompt.contains("xxxx")); + } + + #[test] + fn tool_results_must_match_the_pending_calls() { + let mut agent = agent_with(FakeGenerator::script(&[PARIS_CALL, "It is sunny."])); + assert!(matches!(start(&mut agent), StepOutcome::NeedTools { .. })); + + for results in [ + vec![], + vec![("call_7".to_string(), sunny())], + vec![ + ("call_0".to_string(), sunny()), + ("call_0".to_string(), sunny()), + ], + ] { + let outcome = pollster::block_on(agent.provide_tool_results(results)); + assert!(matches!(outcome, StepOutcome::Error { .. }), "{outcome:?}"); + } + // The agent was left as it was: the right id continues the turn. + assert!(matches!( + provide(&mut agent, sunny()), + StepOutcome::Final { .. } + )); + } + + #[test] + fn start_from_mcp_takes_tools_without_a_description() { + let mut agent = agent_with(FakeGenerator::script(&[ + "\n{\"name\": \"ping\"}\n", + ])); + let outcome = pollster::block_on( + agent.start_from_mcp("ping it", &[serde_json::json!({"name": "ping"})]), + ); + let StepOutcome::NeedTools { calls, .. } = outcome else { + panic!("{outcome:?}"); + }; + assert_eq!(calls[0].name, "ping"); + assert_eq!(calls[0].arguments, serde_json::json!({})); } } diff --git a/crates/agent/src/backend.rs b/crates/agent/src/backend.rs deleted file mode 100644 index 8e469cb..0000000 --- a/crates/agent/src/backend.rs +++ /dev/null @@ -1,143 +0,0 @@ -//! `Generator` implementations built on lean's public native CPU and GPU -//! forward-pass APIs (`lean::cpu`, `lean::model`/`lean::engine`), and a -//! capability-based pick between them, mirroring lean's own CLI (which -//! takes `--engine gpu|cpu` and drives `lean::model::forward_prefill` or -//! `lean::cpu::forward_prefill` directly, picking the GPU path by default -//! and falling back to CPU when no adapter is available). -//! -//! Native only: wasm bindings for the agent live in the `agent-wasm` crate -//! and build their own generators against the same lean APIs, inside the -//! wasm-bindgen surface (which cannot share this module's `pollster` -//! blocking call). - -use anyhow::Result; - -use crate::agent::Generator; - -/// Greedy decode on lean's native CPU engine: a fresh `CpuKvCache` sized -/// to `prompt_ids.len() + max_new_tokens` every call, matching the -/// agent's one-call-per-step usage (no KV-cache prefix reuse across -/// steps; agent turns are short enough that this is the simple, obviously -/// correct choice). -pub struct CpuGenerator { - model: lean::cpu::CpuModel, -} - -impl CpuGenerator { - pub fn load(gguf_path: &str) -> Result { - Ok(Self { - model: lean::cpu::CpuModel::load(gguf_path)?, - }) - } -} - -impl Generator for CpuGenerator { - fn generate(&mut self, prompt_ids: &[u32], max_new_tokens: usize) -> Result> { - let max_ctx = prompt_ids.len() + max_new_tokens; - let mut cache = lean::cpu::CpuKvCache::new(&self.model.config, max_ctx); - let logits = lean::cpu::forward_prefill(&self.model, &mut cache, prompt_ids); - let mut next_id = lean::model::argmax(&logits); - let mut generated = Vec::with_capacity(max_new_tokens); - for _ in 0..max_new_tokens { - if self.model.config.eos_token_ids.contains(&next_id) { - break; - } - generated.push(next_id); - next_id = lean::cpu::forward_decode_step_argmax(&self.model, &mut cache, next_id); - } - Ok(generated) - } -} - -/// Greedy decode on lean's wgpu engine: builds a fresh rope table and -/// `KvCache` per call, blocking on lean's async forward functions with -/// `pollster` (native only, same pattern `lean-cli` uses). -pub struct GpuGenerator { - engine: lean::engine::Engine, - model: lean::model::GpuModel, -} - -impl GpuGenerator { - /// Fails if no wgpu adapter is available; callers should fall back to - /// [`CpuGenerator`] in that case (see [`best_generator`]). - pub fn load(gguf_path: &str) -> Result { - let engine = pollster::block_on(lean::engine::Engine::new_async())?; - let model = lean::model::GpuModel::load(&engine, gguf_path, true)?; - Ok(Self { engine, model }) - } -} - -impl Generator for GpuGenerator { - fn generate(&mut self, prompt_ids: &[u32], max_new_tokens: usize) -> Result> { - // A fresh generation must reset the model's buffer pool first, or - // decode can rebind to a stale KvCache from a previous call. - self.model.pool.reset(); - let max_ctx = prompt_ids.len() as u32 + max_new_tokens as u32 + 4; - let mut cache = lean::model::KvCache::new(&self.engine, &self.model.config, max_ctx); - let (cos, sin) = lean::model::build_rope_tables( - self.model.config.head_dim, - self.model.config.rope_theta, - max_ctx as usize, - ); - let cos_buf = self.engine.buf_f32(&cos, "rope_cos"); - let sin_buf = self.engine.buf_f32(&sin, "rope_sin"); - - let logits = pollster::block_on(lean::model::forward_prefill( - &self.engine, - &self.model, - &mut cache, - prompt_ids, - &cos_buf, - &sin_buf, - None, - )); - let mut next_id = lean::model::argmax(&logits); - let mut generated = Vec::with_capacity(max_new_tokens); - for _ in 0..max_new_tokens { - if self.model.config.eos_token_ids.contains(&next_id) { - break; - } - generated.push(next_id); - next_id = pollster::block_on(lean::model::forward_decode_step_argmax( - &self.engine, - &self.model, - &mut cache, - next_id, - &cos_buf, - &sin_buf, - None, - )); - } - Ok(generated) - } -} - -/// Either engine, picked once at load time by capability (GPU when an -/// adapter is available, CPU otherwise) rather than by benchmarking. -pub enum AnyGenerator { - Gpu(Box), - Cpu(Box), -} - -impl Generator for AnyGenerator { - fn generate(&mut self, prompt_ids: &[u32], max_new_tokens: usize) -> Result> { - match self { - AnyGenerator::Gpu(g) => g.generate(prompt_ids, max_new_tokens), - AnyGenerator::Cpu(c) => c.generate(prompt_ids, max_new_tokens), - } - } -} - -/// Tries the GPU engine first (an adapter request failing, or the model's -/// `head_dim` being unsupported by the GPU attention kernels, are both -/// capability facts, not performance ones); falls back to the CPU engine -/// on either failure. Never benchmarks the two to choose. -pub fn best_generator(gguf_path: &str) -> Result { - match GpuGenerator::load(gguf_path) { - Ok(g) => Ok(AnyGenerator::Gpu(Box::new(g))), - Err(gpu_err) => { - eprintln!("GPU engine unavailable ({gpu_err:#}), falling back to the CPU engine"); - Ok(AnyGenerator::Cpu(Box::new(CpuGenerator::load(gguf_path)?))) - } - } -} diff --git a/crates/agent/src/eval.rs b/crates/agent/src/eval.rs index 5219752..8acdd3c 100644 --- a/crates/agent/src/eval.rs +++ b/crates/agent/src/eval.rs @@ -7,7 +7,9 @@ //! Native only (`std::fs`, `std::time::Instant`); not part of the wasm //! build. -use crate::agent::{Agent, FixtureCaller, Generator, StepOutcome, Tool, ToolCall, ToolCaller}; +use crate::agent::{ + call_or_error, tools_from_mcp, Agent, Generator, StepOutcome, Tool, ToolCall, ToolCaller, +}; use anyhow::{Context, Result}; use serde::Deserialize; use serde_json::Value; @@ -37,13 +39,39 @@ pub struct EvalCase { pub kind: String, pub expected: Vec, pub accept: Accept, - pub tools12_ok: bool, + pub subset_ok: bool, #[serde(default)] pub notes: String, #[serde(default)] pub lang: Option, } +/// A [`ToolCaller`] that serves canned JSON results from +/// `/.json` and records every call it saw, in order. +pub struct FixtureCaller { + results_dir: std::path::PathBuf, + pub calls_seen: Vec<(String, Value)>, +} + +impl FixtureCaller { + pub fn new(results_dir: impl Into) -> Self { + Self { + results_dir: results_dir.into(), + calls_seen: Vec::new(), + } + } +} + +impl ToolCaller for FixtureCaller { + fn call(&mut self, name: &str, args: &Value) -> Result { + self.calls_seen.push((name.to_string(), args.clone())); + let path = self.results_dir.join(format!("{name}.json")); + let text = std::fs::read_to_string(&path) + .with_context(|| format!("reading fixture result {path:?}"))?; + Ok(serde_json::from_str(&text)?) + } +} + pub fn load_cases(path: impl AsRef) -> Result> { let path = path.as_ref(); let raw = std::fs::read_to_string(path) @@ -51,28 +79,12 @@ pub fn load_cases(path: impl AsRef) -> Result> { serde_json::from_str(&raw).with_context(|| format!("parsing eval cases from {path:?}")) } +/// Which tools the agent is given: every tool in `tools.json`, or the +/// eval subset in `tools-subset.json`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ToolSet { All, - Twelve, -} - -impl ToolSet { - pub fn tool_count(&self) -> usize { - match self { - ToolSet::All => 34, - ToolSet::Twelve => 12, - } - } -} - -impl std::fmt::Display for ToolSet { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - ToolSet::All => write!(f, "34 tools"), - ToolSet::Twelve => write!(f, "12 tools"), - } - } + Subset, } fn load_mcp_tools(path: impl AsRef) -> Result> { @@ -80,25 +92,15 @@ fn load_mcp_tools(path: impl AsRef) -> Result> { let raw: Vec = serde_json::from_str( &std::fs::read_to_string(path).with_context(|| format!("reading tools from {path:?}"))?, )?; - raw.into_iter() - .map(|t| { - let name = t["name"] - .as_str() - .ok_or_else(|| anyhow::anyhow!("tool entry missing string `name`: {t}"))?; - let description = t["description"] - .as_str() - .ok_or_else(|| anyhow::anyhow!("tool entry missing string `description`: {t}"))?; - Ok(Tool::from_mcp(name, description, t["inputSchema"].clone())) - }) - .collect() + tools_from_mcp(&raw).with_context(|| format!("reading tools from {path:?}")) } pub fn load_all_tools(fixtures_dir: impl AsRef) -> Result> { load_mcp_tools(fixtures_dir.as_ref().join("tools.json")) } -/// Select the tools for a `ToolSet`, in `tools-12.json`'s own order for -/// `Twelve` (the listing tool comes first; the model's greedy output +/// Select the tools for a `ToolSet`, in `tools-subset.json`'s own order for +/// `Subset` (the listing tool comes first; the model's greedy output /// depends on tool order, so this order is not arbitrary). pub fn select_tools( all: &[Tool], @@ -107,9 +109,9 @@ pub fn select_tools( ) -> Result> { match subset { ToolSet::All => Ok(all.to_vec()), - ToolSet::Twelve => { - let twelve = load_mcp_tools(fixtures_dir.as_ref().join("tools-12.json"))?; - twelve + ToolSet::Subset => { + let subset = load_mcp_tools(fixtures_dir.as_ref().join("tools-subset.json"))?; + subset .iter() .map(|t| { all.iter() @@ -117,7 +119,7 @@ pub fn select_tools( .cloned() .ok_or_else(|| { anyhow::anyhow!( - "tools-12.json tool `{}` not found in the full tool set", + "tools-subset.json tool `{}` not found in the full tool set", t.function.name ) }) @@ -252,10 +254,9 @@ pub fn score( "all required (mutating) calls made in order with correct args".to_string(), ) } else { - let last = case - .expected - .last() - .expect("EvalCase.expected is never empty"); + let Some(last) = case.expected.last() else { + return (false, "subset mode: the case expects no call".to_string()); + }; if calls_made .iter() .any(|c| c.name == last.name && args_match(&last.args, &c.arguments)) @@ -319,8 +320,8 @@ pub fn score( } } -pub fn run_case( - agent: &mut Agent, +pub async fn run_case( + agent: &mut Agent, tools: &[Tool], case: &EvalCase, caller: &mut FixtureCaller, @@ -336,7 +337,7 @@ pub fn run_case( tool_errors += step.tool_errors; }; - let mut outcome = agent.start(&case.utterance, tools); + let mut outcome = agent.start(&case.utterance, tools).await; let (correct, reason) = loop { steps += 1; match outcome { @@ -349,9 +350,7 @@ pub fn run_case( accumulate(&step); let mut results = Vec::with_capacity(calls.len()); for call in &calls { - let result = caller - .call(&call.name, &call.arguments) - .unwrap_or(Value::Null); + let result = call_or_error(caller, call); if call.name == "get_households_and_groups_and_players" { households = Some(result.clone()); } @@ -361,8 +360,15 @@ pub fn run_case( let calls_made = to_tool_calls(&caller.calls_seen); break score(case, &calls_made, households.as_ref(), false); } - outcome = agent.provide_tool_results(results); + outcome = agent.provide_tool_results(results).await; } + StepOutcome::ContextOverflow { needed, available } => { + break ( + false, + format!("context overflow: needed {needed} tokens, {available} available"), + ); + } + StepOutcome::Aborted => break (false, "aborted".to_string()), StepOutcome::Error { message, step } => { if let Some(step) = step { accumulate(&step); @@ -401,7 +407,7 @@ fn skipped_result(case: &EvalCase) -> CaseResult { id: case.id.clone(), correct: false, skipped: true, - reason: "tools12_ok=false: not solvable under the 12-tool subset".to_string(), + reason: "subset_ok=false: not solvable under the 13-tool subset".to_string(), steps: 0, calls_made: Vec::new(), total_ms: 0.0, @@ -414,6 +420,7 @@ fn skipped_result(case: &EvalCase) -> CaseResult { pub struct EvalReport { pub model: String, pub tool_set: ToolSet, + pub tool_count: usize, pub backend: String, pub date: String, pub results: Vec, @@ -428,8 +435,8 @@ pub struct EvalReport { pub max_steps: usize, } -pub fn run_all( - agent: &mut Agent, +pub async fn run_all( + agent: &mut Agent, tools: &[Tool], cases: &[EvalCase], tool_set: ToolSet, @@ -440,12 +447,12 @@ pub fn run_all( let results_dir = results_dir.as_ref(); let mut results = Vec::with_capacity(cases.len()); for case in cases { - if tool_set == ToolSet::Twelve && !case.tools12_ok { + if tool_set == ToolSet::Subset && !case.subset_ok { results.push(skipped_result(case)); continue; } let mut caller = FixtureCaller::new(results_dir); - results.push(run_case(agent, tools, case, &mut caller)); + results.push(run_case(agent, tools, case, &mut caller).await); } let scored: Vec<&CaseResult> = results.iter().filter(|r| !r.skipped).collect(); @@ -459,6 +466,7 @@ pub fn run_all( EvalReport { model: model.into(), tool_set, + tool_count: tools.len(), backend: "native".to_string(), date: date.into(), results, @@ -480,18 +488,18 @@ pub fn render_markdown(report: &EvalReport) -> String { let _ = writeln!( out, - "# Sonos MCP agent eval (lean/Qwen2.5-3B-Instruct, unconstrained), {}", - report.date + "# Sonos MCP agent eval, {}, unconstrained, {}", + report.model, report.date ); let _ = writeln!(out); let _ = writeln!( out, - "## Run: {}, {}, {}, label={}", - report.model, report.tool_set, report.backend, report.label + "## Run: {}, {} tools, {}, label={}", + report.model, report.tool_count, report.backend, report.label ); let _ = writeln!(out); let _ = writeln!(out, "- system: {:?}", report.system_prompt); - let _ = writeln!(out, "- tool_count: {}", report.tool_set.tool_count()); + let _ = writeln!(out, "- tool_count: {}", report.tool_count); let _ = writeln!(out, "- max_new_tokens: {}", report.max_new_tokens); let _ = writeln!(out, "- max_steps: {}", report.max_steps); let _ = writeln!(out); @@ -527,3 +535,179 @@ pub fn render_markdown(report: &EvalReport) -> String { out } + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::ChatTemplate; + use crate::test_support::FakeGenerator; + use serde_json::json; + + const QWEN25_TOOLS_TEMPLATE: &str = + include_str!("../tests/fixtures/qwen25_chat_template.jinja"); + const REPO_ROOT: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../.."); + + fn fixtures_dir() -> String { + format!("{REPO_ROOT}/crates/agent/fixtures/sonos") + } + + fn cases() -> Vec { + load_cases(format!("{REPO_ROOT}/eval/utterances.json")).unwrap() + } + + fn case(id: &str) -> EvalCase { + cases().into_iter().find(|c| c.id == id).unwrap() + } + + fn call(name: &str, arguments: Value) -> ToolCall { + ToolCall { + name: name.to_string(), + arguments, + } + } + + fn households() -> Value { + let path = format!( + "{}/results/get_households_and_groups_and_players.json", + fixtures_dir() + ); + serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap() + } + + #[test] + fn eval_set_has_46_cases_33_solvable_with_the_subset() { + let cases = cases(); + assert_eq!(cases.len(), 46); + assert_eq!(cases.iter().filter(|c| c.subset_ok).count(), 33); + let mut ids: Vec<&str> = cases.iter().map(|c| c.id.as_str()).collect(); + ids.sort(); + ids.dedup(); + assert_eq!(ids.len(), 46, "case ids are unique"); + + let all = load_all_tools(fixtures_dir()).unwrap(); + let subset = select_tools(&all, ToolSet::Subset, fixtures_dir()).unwrap(); + assert_eq!(all.len(), 34); + assert_eq!(subset.len(), 13); + let has = |tools: &[Tool], name: &str| tools.iter().any(|t| t.function.name == name); + for c in &cases { + for e in &c.expected { + assert!(has(&all, &e.name), "{}: unknown tool {}", c.id, e.name); + } + let solvable = c.expected.iter().all(|e| has(&subset, &e.name)); + assert_eq!(solvable, c.subset_ok, "{}: subset_ok", c.id); + } + } + + #[test] + fn exact_mode_needs_the_same_calls() { + let c = case("s02"); + let listing = call("get_households_and_groups_and_players", json!({})); + assert!(score(&c, std::slice::from_ref(&listing), None, true).0); + assert!(!score(&c, &[listing.clone(), listing], None, true).0); + assert!(!score(&c, &[call("get_now_playing", json!({}))], None, true).0); + } + + #[test] + fn subset_mode_needs_the_mutating_calls_in_order() { + let c = case("m01"); + let pause = call( + "pause", + json!({"group_id": "RINCON_KITCHEN01:1", "extra": 1}), + ); + assert!(score(&c, std::slice::from_ref(&pause), None, true).0); + assert!( + !score( + &c, + &[call("pause", json!({"group_id": "RINCON_LIVING01:2"}))], + None, + true + ) + .0 + ); + assert!(!score(&c, &[], None, true).0); + } + + #[test] + fn subset_mode_with_only_reads_needs_the_final_read() { + let c = case("m07"); + let read = call( + "get_group_volume", + json!({"group_id": "RINCON_BEDROOM01:3"}), + ); + assert!(score(&c, &[read], None, true).0); + assert!( + !score( + &c, + &[call("get_group_volume", json!({"group_id": "x"}))], + None, + true + ) + .0 + ); + } + + #[test] + fn no_mutation_mode_needs_a_text_answer_and_no_mutating_call() { + let c = case("m21"); + let listing = call("get_households_and_groups_and_players", json!({})); + assert!(score(&c, std::slice::from_ref(&listing), None, true).0); + assert!(!score(&c, std::slice::from_ref(&listing), None, false).0); + assert!( + !score( + &c, + &[call("pause", json!({"group_id": "garage"}))], + None, + true + ) + .0 + ); + } + + #[test] + fn any_play_mode_needs_a_known_group_id() { + let c = case("m22"); + let households = households(); + let play = |group: &str| { + call( + "play_sonos_favorite", + json!({"group_id": group, "favorite_id": "FV:1", "shuffle": false}), + ) + }; + assert!(score(&c, &[play("RINCON_KITCHEN01:1")], Some(&households), true).0); + assert!(!score(&c, &[play("garage")], Some(&households), true).0); + assert!(!score(&c, &[play("RINCON_KITCHEN01:1")], None, true).0); + let mcp_shaped = json!({"content": [{"type": "text", "text": households.to_string()}]}); + assert!(score(&c, &[play("RINCON_KITCHEN01:1")], Some(&mcp_shaped), true).0); + } + + #[test] + fn too_many_calls_fail_any_mode() { + let c = case("m01"); + let pause = call("pause", json!({"group_id": "RINCON_KITCHEN01:1"})); + let calls = vec![pause; MAX_TOOL_CALLS + 1]; + assert!(!score(&c, &calls, None, true).0); + } + + #[test] + fn run_case_scores_a_scripted_run() { + let generator = FakeGenerator::script(&[ + "\n{\"name\": \"get_households_and_groups_and_players\"}\n", + "\n{\"name\": \"pause\", \"arguments\": {\"group_id\": \"RINCON_KITCHEN01:1\"}}\n", + "The kitchen is paused.", + ]); + let mut agent = Agent::new( + ChatTemplate::new(QWEN25_TOOLS_TEMPLATE), + generator, + "You are a helpful home assistant.", + 64, + ); + let all = load_all_tools(fixtures_dir()).unwrap(); + let tools = select_tools(&all, ToolSet::Subset, fixtures_dir()).unwrap(); + let mut caller = FixtureCaller::new(format!("{}/results", fixtures_dir())); + let result = pollster::block_on(run_case(&mut agent, &tools, &case("m01"), &mut caller)); + assert!(result.correct, "{}", result.reason); + assert_eq!(result.steps, 3); + assert_eq!(result.calls_made.len(), 2); + assert_eq!(result.tool_errors, 0); + } +} diff --git a/crates/agent/src/generator.rs b/crates/agent/src/generator.rs new file mode 100644 index 0000000..ac15a60 --- /dev/null +++ b/crates/agent/src/generator.rs @@ -0,0 +1,94 @@ +//! What the agent loop needs from an inference engine: prompt text in, +//! reply text out, streamed as it is decoded, with a way to stop it and a +//! typed error when the prompt does not fit the engine's context. +//! +//! The loop renders every prompt itself (the model's chat template, with +//! the tools) and parses the reply itself, so any engine that can complete +//! a piece of text can drive it. `crate::lean_generator` implements it on +//! lean's engine. + +use std::fmt; +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +/// Why a [`Generator::generate`] call produced no reply. +#[derive(Debug)] +pub enum GenerateError { + /// The prompt plus `max_new_tokens` needs `needed` tokens and the + /// engine's context holds `available`. Nothing was generated and the + /// engine is still usable. + ContextOverflow { needed: usize, available: usize }, + /// The call was stopped through its [`AbortFlag`]. + Aborted, + /// Anything else (the engine failed). + Other(anyhow::Error), +} + +impl fmt::Display for GenerateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + GenerateError::ContextOverflow { needed, available } => { + write!( + f, + "context overflow: needed {needed} tokens, {available} available" + ) + } + GenerateError::Aborted => write!(f, "aborted"), + GenerateError::Other(e) => write!(f, "{e:#}"), + } + } +} + +impl std::error::Error for GenerateError {} + +impl From for GenerateError { + fn from(e: anyhow::Error) -> Self { + GenerateError::Other(e) + } +} + +/// Stops a running generation: keep a clone, call [`AbortFlag::abort`] +/// from anywhere (another thread, a message handler), and the call in +/// flight returns [`GenerateError::Aborted`] at its next token. +#[derive(Debug, Clone, Default)] +pub struct AbortFlag(Arc); + +impl AbortFlag { + pub fn new() -> Self { + Self::default() + } + + pub fn abort(&self) { + self.0.store(true, Ordering::Relaxed); + } + + pub fn is_aborted(&self) -> bool { + self.0.load(Ordering::Relaxed) + } + + /// Clears the flag, for the next turn. + pub fn reset(&self) { + self.0.store(false, Ordering::Relaxed); + } +} + +/// An engine that completes text. The agent loop drives it from one +/// thread (a browser worker, or a blocking executor natively), so the +/// returned future does not need to be `Send`. +pub trait Generator { + /// Generates the continuation of `prompt`, exactly as given (the + /// prompt is already rendered with the chat template), up to + /// `max_new_tokens` tokens, stopping early at the model's + /// end-of-sequence token. Each piece of the reply goes to `on_text` as + /// soon as it is decoded; the whole reply is returned. Returns + /// [`GenerateError::Aborted`] once `abort` is set, at the engine's + /// next check. + fn generate( + &mut self, + prompt: &str, + max_new_tokens: usize, + on_text: &mut dyn FnMut(&str), + abort: &AbortFlag, + ) -> impl Future>; +} diff --git a/crates/agent/src/lean_generator.rs b/crates/agent/src/lean_generator.rs new file mode 100644 index 0000000..a9134f9 --- /dev/null +++ b/crates/agent/src/lean_generator.rs @@ -0,0 +1,79 @@ +//! [`Generator`] on lean's engine: [`lean::Lean`] picks its backend by +//! capability at load (WebGPU, then the CPU forward pass on a thread pool, +//! then on one thread) and reuses the KV cache of the previous call for the +//! longest prefix the next prompt shares with it, so each agent step only +//! prefills what the previous one did not already hold (the system prompt +//! and tools are prefilled once per turn). One agent per engine. +//! +//! lean checks for an abort once per decoded token, not during prefill; +//! the agent's flag reaches lean's own flag each time a piece of text is +//! streamed. + +use lean::{GenParams, Lean, LeanError, LoadOptions}; + +use crate::generator::{AbortFlag, GenerateError, Generator}; + +/// Greedy decoding on a loaded [`Lean`]. +pub struct LeanGenerator { + lean: Lean, +} + +impl LeanGenerator { + pub async fn load(opts: LoadOptions) -> Result { + Ok(Self::new(Lean::load(opts).await?)) + } + + pub fn new(lean: Lean) -> Self { + Self { lean } + } + + /// The backend lean picked: `"webgpu"`, `"threads"` or `"single"`. + pub fn backend(&self) -> &'static str { + self.lean.backend() + } +} + +impl Generator for LeanGenerator { + async fn generate( + &mut self, + prompt: &str, + max_new_tokens: usize, + on_text: &mut dyn FnMut(&str), + abort: &AbortFlag, + ) -> Result { + if abort.is_aborted() { + return Err(GenerateError::Aborted); + } + let params = GenParams { + max_new_tokens: u32::try_from(max_new_tokens).unwrap_or(u32::MAX), + temperature: 0.0, + ..GenParams::default() + }; + // lean checks its own flag once per token; this call's caller + // holds the agent's flag, so it is forwarded after each piece + // of text. + let stop = lean::AbortFlag::new(); + let reply = self + .lean + .generate( + prompt, + ¶ms, + |text| { + on_text(text); + if abort.is_aborted() { + stop.abort(); + } + }, + Some(&stop), + ) + .await; + match reply { + Ok(_) if abort.is_aborted() => Err(GenerateError::Aborted), + Ok(text) => Ok(text), + Err(LeanError::ContextOverflow { needed, available }) => { + Err(GenerateError::ContextOverflow { needed, available }) + } + Err(LeanError::Other(e)) => Err(GenerateError::Other(e)), + } + } +} diff --git a/crates/agent/src/lib.rs b/crates/agent/src/lib.rs index f88d1f7..ce7697d 100644 --- a/crates/agent/src/lib.rs +++ b/crates/agent/src/lib.rs @@ -1,14 +1,18 @@ -//! A tool-calling agent on top of the lean inference engine. +//! A tool-calling agent loop that works with any inference engine. //! -//! This crate holds everything the engine itself does not: the -//! Qwen2.5-Instruct tool-calling chat template and output parser, the -//! step-wise agent loop, and an offline eval harness. It depends on -//! lean's public Rust API for the model forward pass, using lean's normal -//! engine selection (GPU when available, CPU otherwise); it never pins a -//! backend. +//! This crate holds everything the engine itself does not: the prompt for +//! each step (the model's own chat template with the tools, rendered by +//! lean-chat), the Qwen2.5-Instruct tool-call parser, the step-wise agent +//! loop, and an offline eval harness. The loop reaches the model only +//! through the [`generator::Generator`] trait (prompt text in, reply text +//! out). [`lean_generator::LeanGenerator`] implements it on lean's engine, +//! which picks its backend by capability. pub mod agent; #[cfg(not(target_arch = "wasm32"))] -pub mod backend; -#[cfg(not(target_arch = "wasm32"))] pub mod eval; +pub mod generator; +#[cfg(feature = "lean")] +pub mod lean_generator; +#[cfg(test)] +mod test_support; diff --git a/crates/agent/src/test_support.rs b/crates/agent/src/test_support.rs new file mode 100644 index 0000000..1f17544 --- /dev/null +++ b/crates/agent/src/test_support.rs @@ -0,0 +1,70 @@ +//! A scripted [`Generator`] for the unit tests. + +use std::cell::RefCell; +use std::collections::VecDeque; +use std::rc::Rc; + +use crate::generator::{AbortFlag, GenerateError, Generator}; + +type Reply = Box Result>; + +/// Answers each prompt with `reply(prompt)`, streams the reply to +/// `on_text`, and records every prompt it was given. Like an engine, it +/// returns [`GenerateError::Aborted`] when the flag is set before the call +/// or while the reply streams. +pub(crate) struct FakeGenerator { + reply: Reply, + prompts: Rc>>, +} + +impl FakeGenerator { + pub(crate) fn new(reply: impl FnMut(&str) -> Result + 'static) -> Self { + Self { + reply: Box::new(reply), + prompts: Rc::default(), + } + } + + /// Replies with `replies` in order, then fails. + pub(crate) fn script(replies: &[&str]) -> Self { + let mut replies: VecDeque = replies.iter().map(|r| r.to_string()).collect(); + Self::new(move |_| { + replies.pop_front().ok_or_else(|| { + GenerateError::Other(anyhow::anyhow!("FakeGenerator script exhausted")) + }) + }) + } + + /// Every prompt generated from so far, shared with the generator. + pub(crate) fn prompts(&self) -> Rc>> { + self.prompts.clone() + } +} + +impl Generator for FakeGenerator { + fn generate( + &mut self, + prompt: &str, + _max_new_tokens: usize, + on_text: &mut dyn FnMut(&str), + abort: &AbortFlag, + ) -> impl std::future::Future> { + let reply = if abort.is_aborted() { + Err(GenerateError::Aborted) + } else { + self.prompts.borrow_mut().push(prompt.to_string()); + match (self.reply)(prompt) { + Ok(text) => { + on_text(&text); + if abort.is_aborted() { + Err(GenerateError::Aborted) + } else { + Ok(text) + } + } + Err(e) => Err(e), + } + }; + std::future::ready(reply) + } +} diff --git a/crates/agent/tests/agent_live.rs b/crates/agent/tests/agent_live.rs index 2b6c301..fe9646d 100644 --- a/crates/agent/tests/agent_live.rs +++ b/crates/agent/tests/agent_live.rs @@ -1,9 +1,10 @@ //! End-to-end smoke test of the agent loop against the real -//! Qwen2.5-3B-Instruct model on lean's native CPU engine: the model must +//! Qwen2.5-3B-Instruct model on lean's engine (backend picked by +//! capability: GPU when a wgpu adapter is available, else CPU): the model must //! pick the one tool it's given and call it, then answer in text once it //! has the result. Doesn't assert on the wording of the final answer //! (that's a model-quality question, not a correctness one); see -//! `docs/runs/2026-10-07-mcp-agent-lean.md` for scored eval runs. +//! `docs/runs/` for scored eval runs. //! //! Ignored by default (needs the GGUF + tokenizer files on disk): //! @@ -13,8 +14,9 @@ //! cargo test -p agent --release -- --ignored agent_live //! ``` -use agent::agent::{Agent, AgentChatTemplate, ParsedOutput, ToolCaller}; -use agent::backend::CpuGenerator; +use agent::agent::{Agent, ParsedOutput, ToolCaller}; +use agent::lean_generator::LeanGenerator; +use lean::{Backend, GgufSource, LoadOptions, TokenizerSource}; use serde_json::{json, Value}; struct WeatherCaller; @@ -29,18 +31,14 @@ impl ToolCaller for WeatherCaller { #[test] #[ignore = "needs LEAN_GGUF_QWEN25_3B (GGUF path) and LEAN_TOKENIZER_DIR_QWEN25_3B (dir with tokenizer.json + tokenizer_config.json) on disk; never committed to this repo"] -fn agent_cpu_live_calls_tool_then_answers() { +fn agent_live_calls_tool_then_answers() { let gguf_path = std::env::var("LEAN_GGUF_QWEN25_3B").expect("set LEAN_GGUF_QWEN25_3B to run this test"); let tokenizer_dir = std::env::var("LEAN_TOKENIZER_DIR_QWEN25_3B") .expect("set LEAN_TOKENIZER_DIR_QWEN25_3B to run this test"); - let tokenizer = tokenizers::Tokenizer::from_file(format!("{tokenizer_dir}/tokenizer.json")) - .expect("loading tokenizer.json"); - let cfg_json = std::fs::read_to_string(format!("{tokenizer_dir}/tokenizer_config.json")) - .expect("reading tokenizer_config.json"); - let template = - AgentChatTemplate::from_tokenizer_config_json(&cfg_json).expect("compiling chat_template"); + let template = lean_chat::read_chat_template(&format!("{tokenizer_dir}/tokenizer_config.json")) + .expect("reading the chat template"); let tool = agent::agent::Tool::from_mcp( "get_weather", @@ -48,19 +46,28 @@ fn agent_cpu_live_calls_tool_then_answers() { json!({"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}), ); - let generator = CpuGenerator::load(&gguf_path).expect("loading GGUF"); + let generator = pollster::block_on(LeanGenerator::load(LoadOptions { + gguf: GgufSource::Path(gguf_path), + tokenizer: TokenizerSource::Path(format!("{tokenizer_dir}/tokenizer.json")), + max_ctx: 2048, + backend: Backend::Auto, + })) + .expect("loading the model"); + println!("backend: {}", generator.backend()); let mut agent = Agent::new( template, - tokenizer, generator, - WeatherCaller, "You are a helpful assistant with access to tools. Use them when needed.", 96, ); - let transcript = agent - .run("What's the weather like in Paris right now?", &[tool], 4) - .expect("agent run failed"); + agent.set_max_steps(4); + let transcript = pollster::block_on(agent.run( + "What's the weather like in Paris right now?", + &[tool], + &mut WeatherCaller, + )) + .expect("agent run failed"); assert!(!transcript.steps.is_empty()); let first = &transcript.steps[0]; diff --git a/crates/agent/tests/fixture_parity_agent_tools.rs b/crates/agent/tests/fixture_parity_agent_tools.rs index 6757458..d686543 100644 --- a/crates/agent/tests/fixture_parity_agent_tools.rs +++ b/crates/agent/tests/fixture_parity_agent_tools.rs @@ -1,6 +1,8 @@ -//! Token-exact parity for the Qwen2.5-Instruct tools template against -//! `transformers.apply_chat_template(..., tools=...)` itself (no model -//! forward pass needed, this is a tokenizer/template check). +//! Byte-exact parity of the agent's prompts (`render_prompt`, lean-chat +//! with the model's own chat template) against +//! `transformers.apply_chat_template(..., tools=...)` itself, and +//! token-exact parity of their tokenization (no model forward pass needed, +//! this is a template/tokenizer check). //! //! See `reference/gen_fixture_agent_tools.py` for how the fixture was //! generated. @@ -12,7 +14,7 @@ //! cargo test -p agent --release -- --ignored fixture_parity_agent_tools //! ``` -use agent::agent::{AgentChatTemplate, Message, Tool}; +use agent::agent::{render_prompt, Message, Tool}; use serde::Deserialize; use serde_json::Value; @@ -38,10 +40,8 @@ fn fixture_parity_agent_tools() { let tokenizer = tokenizers::Tokenizer::from_file(format!("{tokenizer_dir}/tokenizer.json")) .expect("loading tokenizer.json"); - let cfg_json = std::fs::read_to_string(format!("{tokenizer_dir}/tokenizer_config.json")) - .expect("reading tokenizer_config.json"); - let template = - AgentChatTemplate::from_tokenizer_config_json(&cfg_json).expect("compiling chat_template"); + let template = lean_chat::read_chat_template(&format!("{tokenizer_dir}/tokenizer_config.json")) + .expect("reading the chat template"); let fixture_json = std::fs::read_to_string(concat!( env!("CARGO_MANIFEST_DIR"), @@ -52,8 +52,7 @@ fn fixture_parity_agent_tools() { serde_json::from_str(&fixture_json).expect("parsing fixture_agent_tools.json"); for case in &fixture.cases { - let rendered = template - .render(&case.messages, &case.tools, true) + let rendered = render_prompt(&template, &case.messages, &case.tools) .unwrap_or_else(|e| panic!("case {}: render failed: {e:#}", case.name)); assert_eq!( rendered, case.rendered, @@ -74,15 +73,16 @@ fn fixture_parity_agent_tools() { } } -/// `Message`'s `content: Option` round-trips either JSON shape the -/// fixtures carry: a plain string (every `system`/`user`/`tool` message) -/// or `null` (an `assistant` message that only has `tool_calls`). +/// `Message` reads either `content` shape the fixtures carry: a plain +/// string (every `system`/`user`/`tool` message) or `null` (an `assistant` +/// message that only has `tool_calls`, read as `""`, which the template +/// treats the same way). #[test] fn message_deserializes_null_content_and_string_tool_calls_arguments() { let json = r#"{"role": "assistant", "content": null, "tool_calls": [{"id": "call_0", "type": "function", "function": {"name": "get_now_playing", "arguments": "{}"}}]}"#; let msg: Message = serde_json::from_str(json).unwrap(); - assert_eq!(msg.content, None); - let tool_calls = msg.tool_calls.unwrap(); + assert_eq!(msg.content, ""); + let tool_calls = msg.tool_calls; assert_eq!( tool_calls[0].function.arguments, Value::String("{}".to_string()) diff --git a/docs/runs/2026-10-07-sonos-eval.md b/docs/runs/2026-10-07-sonos-eval.md new file mode 100644 index 0000000..c572667 --- /dev/null +++ b/docs/runs/2026-10-07-sonos-eval.md @@ -0,0 +1,61 @@ +# Sonos MCP agent eval, Qwen2.5-3B-Instruct Q4_0, CPU, unconstrained, 2026-10-07 + +46 cases from `eval/utterances.json`: 33 scored with the 13-tool set, 13 skipped because they need a tool outside it. + +## Run: Qwen2.5-3B-Instruct-q4_0.gguf (lean, CPU, unconstrained), 13 tools, native, label=sonos-subset + +- system: "You are a helpful home assistant with access to Sonos speaker controls." +- tool_count: 13 +- max_new_tokens: 256 +- max_steps: 6 + +| id | correct | steps | total_s | retries | tool_errors | reason | +|----|---------|-------|---------|---------|-------------|--------| +| s01 | false | 2 | 463.213 | 0 | 0 | final read get_now_playing({"group_id":"RINCON_LIVING01:2"}) never made with correct args | +| s02 | true | 2 | 540.176 | 0 | 0 | matched expected sequence exactly | +| s03 | false | 2 | 493.883 | 0 | 0 | final read get_sonos_favorites({"household_id":"Sonos_ABC123XYZ"}) never made with correct args | +| s04 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| s05 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| s06 | false | 2 | 445.383 | 0 | 0 | missing required call skip_to_next_track({"group_id":"RINCON_LIVING01:2"}) in order after step 0 | +| s07 | false | 1 | 217.598 | 0 | 0 | missing required call skip_to_previous_track({"group_id":"RINCON_LIVING01:2"}) in order after step 0 | +| s08 | false | 2 | 443.362 | 0 | 0 | missing required call pause({"group_id":"RINCON_LIVING01:2"}) in order after step 0 | +| s09 | false | 2 | 439.755 | 0 | 0 | missing required call resume({"group_id":"RINCON_LIVING01:2"}) in order after step 0 | +| s10 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m01 | false | 2 | 443.202 | 0 | 0 | missing required call pause({"group_id":"RINCON_KITCHEN01:1"}) in order after step 0 | +| m02 | false | 2 | 458.214 | 0 | 0 | missing required call set_group_volume({"group_id":"RINCON_LIVING01:2","volume":20}) in order after step 0 | +| m03 | false | 2 | 446.979 | 0 | 0 | missing required call set_group_mute({"group_id":"RINCON_BEDROOM01:3","muted":true}) in order after step 0 | +| m04 | false | 3 | 840.498 | 0 | 0 | missing required call skip_to_next_track({"group_id":"RINCON_KITCHEN01:1"}) in order after step 0 | +| m05 | false | 2 | 566.183 | 0 | 0 | missing required call play_sonos_favorite({"group_id":"RINCON_LIVING01:2","favorite_id":"FV:1","shuffle":false}) in order after step 0 | +| m06 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m07 | false | 2 | 447.899 | 0 | 0 | final read get_group_volume({"group_id":"RINCON_BEDROOM01:3"}) never made with correct args | +| m08 | false | 2 | 455.013 | 0 | 0 | missing required call adjust_group_volume({"group_id":"RINCON_KITCHEN01:1","volume_delta":10}) in order after step 0 | +| m09 | false | 2 | 542.481 | 0 | 0 | missing required call set_group_mute({"group_id":"RINCON_LIVING01:2","muted":false}) in order after step 0 | +| m10 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m11 | false | 3 | 875.686 | 0 | 0 | missing required call play_artist({"group_id":"RINCON_KITCHEN01:1","artist":"Nirvana","shuffle":false}) in order after step 0 | +| m12 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| s11 | false | 2 | 449.862 | 0 | 0 | missing required call set_group_volume({"group_id":"RINCON_LIVING01:2","volume":35}) in order after step 0 | +| s12 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| s13 | false | 3 | 1181.724 | 1 | 0 | missing required call set_group_mute({"group_id":"RINCON_KITCHEN01:1","muted":false}) in order after step 0 | +| s14 | false | 2 | 455.865 | 0 | 0 | final read get_now_playing({"group_id":"RINCON_BEDROOM01:3"}) never made with correct args | +| s15 | false | 2 | 440.646 | 0 | 0 | final read get_now_playing({"group_id":"RINCON_LIVING01:2"}) never made with correct args | +| s16 | false | 2 | 435.274 | 0 | 0 | missing required call adjust_group_volume({"group_id":"RINCON_LIVING01:2","volume_delta":-10}) in order after step 0 | +| s17 | false | 2 | 518.446 | 0 | 0 | missing required call skip_to_next_track({"group_id":"RINCON_KITCHEN01:1"}) in order after step 0 | +| m13 | true | 3 | 915.273 | 0 | 0 | all required (mutating) calls made in order with correct args | +| m14 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m15 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m16 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m17 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m18 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m19 | skipped | - | - | - | - | subset_ok=false: not solvable under the 13-tool subset | +| m20 | false | 2 | 441.131 | 0 | 0 | missing required call pause({"group_id":"RINCON_LIVING01:2"}) in order after step 0 | +| m21 | false | 2 | 421.905 | 0 | 0 | no_mutation mode: mutating call pause({"group_id":"garage"}) was made | +| m22 | false | 1 | 231.110 | 0 | 0 | any_play mode: no play_* call with a valid group_id was made | +| m23 | false | 2 | 512.525 | 0 | 0 | missing required call pause({"group_id":"RINCON_KITCHEN01:1"}) in order after step 0 | +| m24 | false | 2 | 422.575 | 0 | 0 | missing required call pause({"group_id":"RINCON_KITCHEN01:1"}) in order after step 0 | +| m25 | false | 2 | 511.158 | 0 | 0 | missing required call pause({"group_id":"RINCON_KITCHEN01:1"}) in order after step 0 | +| m26 | false | 2 | 531.815 | 0 | 0 | missing required call adjust_group_volume({"group_id":"RINCON_LIVING01:2","volume_delta":10}) in order after step 0 | +| m27 | false | 2 | 505.882 | 0 | 0 | missing required call play_artist({"group_id":"RINCON_KITCHEN01:1","artist":"Nirvana","shuffle":false}) in order after step 0 | +| m28 | false | 2 | 441.762 | 0 | 0 | final read get_now_playing({"group_id":"RINCON_KITCHEN01:1"}) never made with correct args | +| s18 | true | 1 | 214.497 | 0 | 0 | no mutating call made; final text answer given | + +**Summary**: correct 9.1% (33 scored, 13 skipped), mean_steps 2.03, mean_total_s 507.605, total_retries 1, total_tool_errors 0 diff --git a/eval/utterances.json b/eval/utterances.json index 888bb64..7393ce4 100644 --- a/eval/utterances.json +++ b/eval/utterances.json @@ -12,8 +12,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "No room named; the only currently PLAYING group in the canned household is Living Room, so that's the canonical target. A model may legitimately call get_households_and_groups_and_players first to find it — that read is free." + "subset_ok": true, + "notes": "No room named; the only currently PLAYING group in the canned household is Living Room, so that's the canonical target. A model may legitimately call get_households_and_groups_and_players first to find it; that read is free." }, { "id": "s02", @@ -26,7 +26,7 @@ } ], "accept": "exact", - "tools12_ok": true, + "subset_ok": true, "notes": "Single no-arg discovery call; nothing else is needed or sensible before it." }, { @@ -42,7 +42,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Only one household in the fixture, so household_id is effectively fixed context." }, { @@ -58,8 +58,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "get_sonos_playlists is not in the 12-tool subset." + "subset_ok": false, + "notes": "get_sonos_playlists is not in the 13-tool subset." }, { "id": "s05", @@ -74,8 +74,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "get_registered_music_services is not in the 12-tool subset." + "subset_ok": false, + "notes": "get_registered_music_services is not in the 13-tool subset." }, { "id": "s06", @@ -90,7 +90,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "No room named; defaults to the only actively playing group (Living Room)." }, { @@ -106,7 +106,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Same default-group reasoning as s06." }, { @@ -122,7 +122,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "pause is stateless per its description ('always execute when asked') so no need to check playback state first." }, { @@ -138,7 +138,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Mirrors s08 for resume." }, { @@ -154,8 +154,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "get_shuffle_repeat_crossfade is not in the 12-tool subset." + "subset_ok": false, + "notes": "get_shuffle_repeat_crossfade is not in the 13-tool subset." }, { "id": "m01", @@ -174,7 +174,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Named room requires resolving 'kitchen' to its group id first." }, { @@ -195,7 +195,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "\"down to\" is an absolute target, not a delta, so set_group_volume (not adjust_group_volume) is correct." }, { @@ -216,7 +216,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "muted is required by the schema even though the user only said 'mute' (true is implied but must still be passed explicitly)." }, { @@ -236,7 +236,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Named room requires the lookup." }, { @@ -264,8 +264,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "shuffle is required by play_sonos_favorite's schema even though the user never mentioned shuffling — the model must default it to false, which is easy to get wrong (some models omit required-but-unmentioned booleans)." + "subset_ok": true, + "notes": "shuffle is required by play_sonos_favorite's schema even though the user never mentioned shuffling: the model must default it to false, which is easy to get wrong (some models omit required-but-unmentioned booleans)." }, { "id": "m06", @@ -287,8 +287,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "add_players_to_group is not in the 12-tool subset. Also tests whether the model picks the kitchen's existing group_id as the destination (adding the bedroom player into it) rather than inventing a new group or reversing which side is the 'group_id' vs the 'player_ids'." + "subset_ok": false, + "notes": "add_players_to_group is not in the 13-tool subset. Also tests whether the model picks the kitchen's existing group_id as the destination (adding the bedroom player into it) rather than inventing a new group or reversing which side is the 'group_id' vs the 'player_ids'." }, { "id": "m07", @@ -307,8 +307,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "A read-only multi-step case: both calls are reads, so there is no mutating call to check — correctness here means the second call happened with the right group_id." + "subset_ok": true, + "notes": "A read-only multi-step case: both calls are reads, so there is no mutating call to check; correctness here means the second call happened with the right group_id." }, { "id": "m08", @@ -328,8 +328,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "A relative delta ('by 10') should map to adjust_group_volume, not set_group_volume — tests the set vs. adjust distinction in the opposite direction from m02." + "subset_ok": true, + "notes": "A relative delta ('by 10') should map to adjust_group_volume, not set_group_volume. It tests the set vs. adjust distinction in the opposite direction from m02." }, { "id": "m09", @@ -349,7 +349,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Mirrors m03 with muted: false." }, { @@ -372,8 +372,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "move_audio_to_players is not in the 12-tool subset. Its first argument is named source_group_id rather than group_id like every other group-scoped tool — an easy field-name slip." + "subset_ok": false, + "notes": "move_audio_to_players is not in the 13-tool subset. Its first argument is named source_group_id rather than group_id like every other group-scoped tool, an easy field-name slip." }, { "id": "m11", @@ -394,7 +394,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "music_service is required by play_artist's schema but not determined by the utterance; the scorer's subset-args matching (expected keys must match, extra keys in the actual call are ignored) means the model may include any music_service value without being penalized. shuffle defaults to false per m05's pattern for a required-but-unmentioned boolean." }, { @@ -417,8 +417,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "remove_players_from_group is not in the 12-tool subset. The fixture Kitchen group has two players (RINCON_KITCHEN01, RINCON_KITCHEN02) so 'ungroup' is unambiguous: remove the non-coordinator player RINCON_KITCHEN02, leaving RINCON_KITCHEN01 as the group's sole (coordinator) player." + "subset_ok": false, + "notes": "remove_players_from_group is not in the 13-tool subset. The fixture Kitchen group has two players (RINCON_KITCHEN01, RINCON_KITCHEN02) so 'ungroup' is unambiguous: remove the non-coordinator player RINCON_KITCHEN02, leaving RINCON_KITCHEN01 as the group's sole (coordinator) player." }, { "id": "s11", @@ -438,7 +438,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Named-room volume set by explicit number; the lookup resolves 'living room' to its group id." }, { @@ -459,8 +459,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "Player-level mute (contrast with s13's group-level mute) — naming the individual speaker ('Kitchen Surround') rather than the room should route to set_player_mute/player_id, not set_group_mute/group_id. set_player_mute is not in the 12-tool subset." + "subset_ok": false, + "notes": "Player-level mute (contrast with s13's group-level mute): naming the individual speaker ('Kitchen Surround') rather than the room should route to set_player_mute/player_id, not set_group_mute/group_id. set_player_mute is not in the 13-tool subset." }, { "id": "s13", @@ -480,8 +480,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "Room name (not an individual speaker) routes to set_group_mute/group_id — contrast with s12's player-level mute." + "subset_ok": true, + "notes": "Room name (not an individual speaker) routes to set_group_mute/group_id, in contrast with s12's player-level mute." }, { "id": "s14", @@ -500,7 +500,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Bedroom is IDLE in the canned household; the correct call is still get_now_playing on the bedroom's group id, regardless of playback state." }, { @@ -520,7 +520,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "No room named; Living Room is the only PLAYING group in the canned household, so it's the canonical target, mirroring s01." }, { @@ -541,8 +541,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "No room named — the playing group (Living Room) is the canonical target, same as s01/s15. Delta convention: an unquantified 'a bit'/'a little' nudge is ±10; this case is the 'down' direction." + "subset_ok": true, + "notes": "No room named: the playing group (Living Room) is the canonical target, same as s01/s15. Delta convention: an unquantified 'a bit'/'a little' nudge is ±10; this case is the 'down' direction." }, { "id": "s17", @@ -561,7 +561,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Named-room skip, mirrors m04's pattern but phrased as a single terse command." }, { @@ -589,7 +589,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Same pattern as m05 (Jazz Vibes/living room) but with the 'Morning News' favourite (FV:2) in the kitchen, to check the model doesn't hardcode m05's ids." }, { @@ -617,8 +617,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "Sonos playlist by name, not a music-service playlist (play_playlist). get_sonos_playlists and play_sonos_playlist are both outside the 12-tool subset." + "subset_ok": false, + "notes": "Sonos playlist by name, not a music-service playlist (play_playlist). get_sonos_playlists and play_sonos_playlist are both outside the 13-tool subset." }, { "id": "m15", @@ -638,8 +638,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "A named live broadcast station routes to play_radio (radio_name), not play_station, which is for algorithmic/seeded stations (see m-cases below). play_radio is not in the 12-tool subset." + "subset_ok": false, + "notes": "A named live broadcast station routes to play_radio (radio_name), not play_station, which is for algorithmic/seeded stations (see m-cases below). play_radio is not in the 13-tool subset." }, { "id": "m16", @@ -660,8 +660,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "A specific track named by title and artist routes to play_track, not play_station's track-seeded radio (contrast with a 'Bohemian Rhapsody radio' phrasing, which would seed play_station instead). play_track is not in the 12-tool subset." + "subset_ok": false, + "notes": "A specific track named by title and artist routes to play_track, not play_station's track-seeded radio (contrast with a 'Bohemian Rhapsody radio' phrasing, which would seed play_station instead). play_track is not in the 13-tool subset." }, { "id": "m17", @@ -683,8 +683,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "shuffle:true must come from play_album's own shuffle argument (per its schema note), not a separate set_shuffle_repeat_crossfade call. play_album is not in the 12-tool subset." + "subset_ok": false, + "notes": "shuffle:true must come from play_album's own shuffle argument (per its schema note), not a separate set_shuffle_repeat_crossfade call. play_album is not in the 13-tool subset." }, { "id": "m18", @@ -713,8 +713,8 @@ } ], "accept": "subset", - "tools12_ok": false, - "notes": "Accepted sequence: group the two rooms first (kitchen's existing group id absorbs the bedroom player, per add_players_to_group's own guidance to use the larger/coordinator group as the target), then play a genre-seeded station on the combined group. A fresh get_households_and_groups_and_players call between add_players_to_group and play_station (per that tool's 'always call it again' note) is a free read under accept:subset and not required, since the fixture's group id doesn't change. Neither add_players_to_group nor play_station is in the 12-tool subset." + "subset_ok": false, + "notes": "Accepted sequence: group the two rooms first (kitchen's existing group id absorbs the bedroom player, per add_players_to_group's own guidance to use the larger/coordinator group as the target), then play a genre-seeded station on the combined group. A fresh get_households_and_groups_and_players call between add_players_to_group and play_station (per that tool's 'always call it again' note) is a free read under accept:subset and not required, since the fixture's group id doesn't change. Neither add_players_to_group nor play_station is in the 13-tool subset." }, { "id": "m19", @@ -737,7 +737,7 @@ } ], "accept": "subset", - "tools12_ok": false, + "subset_ok": false, "notes": "Same tool as m10 (kitchen/living room instead of living room/bedroom) with a different destination player set (both Living Room players). Checks the model doesn't hardcode m10's args." }, { @@ -757,7 +757,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "'Everything' means every group actually playing, not every group that exists: only Living Room is PLAYING in the canned household (Kitchen is PAUSED, Bedroom is IDLE), so this resolves to a single pause call, not one per group." }, { @@ -766,7 +766,7 @@ "kind": "multi", "expected": [], "accept": "no_mutation", - "tools12_ok": true, + "subset_ok": true, "notes": "No 'Garage' room/group/player exists in the canned household. Expect no mutating tool call (the model should say it can't find that room, not guess an id), scored with accept 'no_mutation'." }, { @@ -775,7 +775,7 @@ "kind": "multi", "expected": [], "accept": "any_play", - "tools12_ok": true, + "subset_ok": true, "notes": "Fully ambiguous request: no room, artist, track, station, or playlist named. Any single play_* call with a valid group_id is accepted (accept 'any_play'); there is no one canonical expected sequence." }, { @@ -802,8 +802,8 @@ } ], "accept": "subset", - "tools12_ok": true, - "notes": "Two rooms, two different actions in one utterance — tests that the model doesn't apply one action to both rooms. Volume delta uses the same ±10 'a bit' convention as s16 (this is the 'up' direction, unquantified)." + "subset_ok": true, + "notes": "Two rooms, two different actions in one utterance: it tests that the model doesn't apply one action to both rooms. Volume delta uses the same ±10 'a bit' convention as s16 (this is the 'up' direction, unquantified)." }, { "id": "m24", @@ -822,7 +822,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Typo'd room name ('kitchn') should still fuzzy-match to Kitchen rather than being treated as an unknown room." }, { @@ -843,7 +843,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "French for 'Pause the kitchen.' Tests multilingual room-name and intent resolution; expected calls are identical to an equivalent English case (cf. m01)." }, { @@ -865,7 +865,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "French for 'Turn up the living room's volume', unquantified. Uses the same ±10 'a bit' delta convention as s16/m23 ('salon' = living room)." }, { @@ -887,7 +887,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "Same pattern as m11 but with a speaker name ('BLOUPBLIP') that has no match in the fixture household; the expected call uses the fixture's kitchen group id (RINCON_KITCHEN01:1) as a stand-in target the way m11/m13 do, since the fixture has no room named 'bloupblip'. With no id known yet, the model should call get_households_and_groups_and_players first rather than answer with a prose refusal on step 0." }, { @@ -907,7 +907,7 @@ } ], "accept": "subset", - "tools12_ok": true, + "subset_ok": true, "notes": "'Kitchen Surround' is a *player* name (playerId RINCON_KITCHEN02) belonging to the Kitchen *group* (groupId RINCON_KITCHEN01:1) in the fixture household. get_now_playing takes a group_id, not a player_id, so the model must resolve the player's group and call it with RINCON_KITCHEN01:1; a player id passed as group_id is rejected by the server as an invalid group_id." }, { @@ -916,7 +916,7 @@ "kind": "single", "expected": [], "accept": "no_mutation", - "tools12_ok": true, + "subset_ok": true, "notes": "Chat-shaped question with tools present: the right answer is plain text with no mutating call (accept 'no_mutation'); a read-only tool call first is allowed." } ] diff --git a/scripts/build.sh b/scripts/build.sh index 7b24bd3..8af6851 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -1,11 +1,6 @@ #!/usr/bin/env bash -# Builds the agent-wasm package and assembles the deployed site into -# _site/, following the pattern of llm-web's scripts/build.sh (build tag -# rewrite, local-path leak check on the built wasm). -# -# The web/ page (phase 2 of this repo) is not built yet; this currently -# just builds and checks the wasm package at _site/pkg. Once web/ exists, -# add its copy and ENGINE_BUILD rewrite here, following llm-web's script. +# Builds the agent-wasm package, checks the built wasm for local paths, and +# copies the package to _site/pkg. # # Requires wasm-pack and wasm-bindgen-cli (version matching Cargo.lock's # wasm-bindgen exactly: `cargo install wasm-bindgen-cli --version @@ -20,9 +15,13 @@ ENGINE_BUILD="${ENGINE_BUILD:-dev}" REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$REPO_ROOT" +# SIMD128 comes from .cargo/config.toml. The remap flag (keeps the home +# directory out of compiled-in source paths) is passed with --config, which +# cargo merges with that file's rustflags; RUSTFLAGS would replace them and +# drop +simd128. +REMAP_CONFIG="target.wasm32-unknown-unknown.rustflags=[\"--remap-path-prefix=$HOME=/home\"]" echo "==> Building agent-wasm (web feature) for ENGINE_BUILD=$ENGINE_BUILD" -RUSTFLAGS="-C target-feature=+simd128 --remap-path-prefix=$HOME=/home" \ - wasm-pack build crates/agent-wasm --target web --release +wasm-pack build crates/agent-wasm --target web --release --config "$REMAP_CONFIG" WASM="crates/agent-wasm/pkg/agent_wasm_bg.wasm" @@ -40,4 +39,4 @@ rm -rf _site mkdir -p _site/pkg cp -R crates/agent-wasm/pkg/. _site/pkg/ -echo "==> Wrote _site (pkg only; no page yet, see this script's header)" +echo "==> Wrote _site/pkg"