diff --git a/Cargo.lock b/Cargo.lock index efb55396..d4110580 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -675,6 +675,18 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libsy" +version = "0.1.0" +dependencies = [ + "async-trait", + "rand 0.8.6", + "reqwest", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "libyaml-rs" version = "0.3.0" diff --git a/Cargo.toml b/Cargo.toml index 56e23bbf..e11c2c32 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ [workspace] resolver = "2" members = [ + "crates/libsy", "crates/switchyard-components", "crates/switchyard-components-v2", "crates/switchyard-components-v2-macros", diff --git a/crates/libsy/Cargo.toml b/crates/libsy/Cargo.toml new file mode 100644 index 00000000..28182acc --- /dev/null +++ b/crates/libsy/Cargo.toml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libsy" +version = "0.1.0" +description = "Switchyard library crate" +authors.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true + +[dependencies] +async-trait = "0.1" +rand = "0.8" +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" + +[dev-dependencies] +tokio = { version = "1", features = ["macros", "rt"] } diff --git a/crates/libsy/DESIGN.md b/crates/libsy/DESIGN.md new file mode 100644 index 00000000..7b6d5781 --- /dev/null +++ b/crates/libsy/DESIGN.md @@ -0,0 +1,220 @@ +# Switchyard-Lib Design + +Status: design draft. The Rust crate in `src/` (`lib.rs` traits, `rand.rs`, `llm_class.rs`, and the +`client.rs` wrapper) is the proof of concept and the design of record; the type names below match the +code (the `AgentApi*` prefix). This doc describes the intended shape and marks where the POC still +diverges (see *POC Gaps*). The model is language-independent — "trait" means interface/contract — +but it is worked out against the Rust code, not any binding or port. + +## Summary + +`libsy` is a lightweight **library** for multi-LLM agent optimization, with **routing** as the first +case. It sits where a request is about to hit a model and decides — statefully, using more than the +request — which model(s) to call, how to rewrite the call, or whether to skip it. + +It is not a service: no server, and the optimization core owns no transport — model calls are the +host's job (see *The Loop*). It is **provider- and transport-agnostic** so it embeds in inference +platforms, agentic systems, and routing proxies; the bundled `HttpCaller` (used by the client +wrapper) is a convenience default, not a coupling. + +## The Loop + +The whole library is one pair of traits and one loop: + +```text +AgentApiOptAlgorithm --optimizer()--> AgentApiOptimizer --feed()/optimize()--> Decision + (config/factory) (per-session, stateful) +``` + +The host feeds the optimizer inputs and calls `optimize()`, which returns either "make these model +calls and feed me the results" (`ModelInference`) or "done, return to the agent" (`Return`): + +```text +feed(request) +loop: + match optimize(): + ModelInference(reqs) -> for r in reqs { feed(host_call(r)) } // host owns the transport + Return -> break +``` + +A one-round router and a multi-round LLM classifier are the *same* loop; the host's code does not +change when the algorithm's round count does. Most integrators do not write this loop by hand — +`RoutedClient` drives it for them (see *Client Wrapper*). + +```rust +#[async_trait] +pub trait AgentApiOptimizer: Send + Sync { + async fn feed(&mut self, input: AgentApiOptInput, enrichment: EnrichmentData) + -> Result<(), Box>; + async fn optimize(&mut self) -> Result, Box>; +} + +pub trait AgentApiOptAlgorithm: Send + Sync { + fn optimizer(&self) -> Box>; // fresh instance per session +} +``` + +`AgentApiOptAlgorithm` is immutable, shareable config; the `AgentApiOptimizer` holds the per-session +mutable state. `D` is the algorithm's typed **decision metadata** (see *Construction* for the config +path). + +## Client Wrapper + +`RoutedClient` turns the loop into a **blackbox LLM client**, so most integrators never touch +`feed`/`optimize`. It is a drop-in for an existing client — one `complete` call in, one response out: + +```rust +impl RoutedClient { + pub fn new(algorithm: Box>, caller: Box) -> Self; + pub fn with_http(algorithm: Box>, base_url, api_key) -> Self; + + async fn complete(&self, request: AgentApiRequest) -> Result>; + async fn complete_with(&self, request, enrichment) -> Result>; +} +``` + +Each `complete` mints a fresh optimizer, feeds the request, and drives the loop to `Return` — +performing every model call the optimizer asks for and feeding the results back. Single- and +multi-round algorithms are indistinguishable to the caller: the classifier's internal classifier +call is just another call the client makes and consumes, and the routed response is what comes back. + +The one piece of I/O `libsy` does not own is the model call itself, expressed as a trait: + +```rust +#[async_trait] +pub trait ModelCaller: Send + Sync { + async fn call(&self, request: AgentApiRequest) -> Result>; +} +``` + +Integrators implement `ModelCaller` over their own transport. `HttpCaller` is the built-in default — +it POSTs to `{base_url}/chat/completions` on any OpenAI-compatible endpoint, and `with_http` wires it +in one line. This keeps the transport-agnostic core intact: the default is a convenience, not +something the algorithms depend on. + +## Data Model + +Neutral request/response — the host converts to/from provider wire types at the edge: + +```rust +pub struct AgentApiRequest { pub prompt: String, pub model: String } // +params/tools later +pub struct AgentApiResponse { pub completion: String } // +token usage, latency later +``` + +Inputs are a tagged union, so optimization is driven by more than the request — responses **and** +events from the agentic stack: + +```rust +pub enum AgentApiOptInput { + Request(AgentApiRequest), + Response(AgentApiResponse), // design intent; see POC Gaps for the current shape + Signal(Signal), // tool/task/plan/budget/telemetry events + Metadata(BTreeMap), +} + +pub enum Signal { // open, versioned; algorithms opt in, ignore the rest + ToolCallCompleted { tool: String, ok: bool, latency_ms: u64 }, + TaskStarted { task: String }, TaskCompleted { task: String, ok: bool }, + Budget { spent_tokens: u64, remaining_tokens: Option }, + Telemetry { key: String, value: f64 }, // e.g. replica health, queue depth + // …additive; a new signal never breaks an existing algorithm +} +``` + +A decision names the calls to make plus its own explanation: + +```rust +pub enum Decision { ModelInference(AgentApiOptimizerResponse), Return } + +pub struct AgentApiOptimizerResponse { + pub requests: Vec, + pub enrichment_data: Vec, + pub decision_reasoning: Option, // human-readable "why" (traces) + pub decision_info: Option, // structured "why" (metrics, eval) +} +``` + +Correlation is carried on every feed and every emitted request — the spine of both policy and +observability: + +```rust +pub struct EnrichmentData { + pub session_id: Option, pub agent_id: Option, + pub task_id: Option, pub tool_id: Option, // session/agent/task/tool + pub correlation_id: Option, // external trace join + pub extra_metadata: Option>, +} +``` + +## Algorithms Are Traits + +Implement `AgentApiOptAlgorithm` + `AgentApiOptimizer`; no base class, no framework. Reference impls: + +- **`RandomRouterAlgorithm`** (`rand.rs`) — weighted random over N targets. One round: draw a target, + rewrite `model`, return; after the response is fed, `Return`. +- **`LlmClassifierAlgorithm`** (`llm_class.rs`) — multi-round: emit a classifier call, parse its + score, route strong/weak by threshold, return; after that response, `Return`. + +Future algorithms on the same surface: latency/health-aware routing (`Telemetry`), cost-budget +routing (`Budget`), cascade/escalation (`ToolCallCompleted` + response quality), semantic caching. + +## Statefulness and Concurrency + +- One optimizer per session; discarded at session end. Immutable algorithm config, mutable instance. +- Feeds arrive asynchronously but are applied serially — the host owns a per-session queue/task; the + instance is not internally synchronized. Separate sessions are separate instances, so cross-session + use is safe with no shared state. +- Cross-session policy (fleet load, global budgets) enters as fed `Signal`s, not hidden state — so + every decision is a function of that session's fed history, and reproducible for evaluation. + +## Observability + +Production and research/benchmarking need the same data, so it is a core feature: + +- `optimize()` and each host model call are **spans**, tagged with the correlation set + algorithm + name + round. +- Metrics on those spans: **token counts, timings, failures**, decisions by model/tier. Token usage + is a fed input (on `AgentApiResponse`/`Budget`), not a guess. +- `decision_reasoning`/`decision_info` are recorded per decision: an explanation in production, the + dataset in research (win-rate / cost / latency deltas). Shadow and A/B modes fall out of the loop + directly; a benchmark is a batch of recorded sessions. +- Emitted through a thin sink abstraction, so `libsy` mandates no telemetry backend and a no-op sink + compiles it away. + +## Construction + +Two entry points to select an algorithm: + +- **Direct** (built) — instantiate a concrete algorithm; you get typed decision metadata. + `RandomRouterAlgorithm { models, rng_seed }.optimizer()`. Wrap it in a `RoutedClient` (see *Client + Wrapper*) for a client-style API. +- **Config** (future) — `SwitchyardOptimizer::from_config(cfg, registry)` (aliased `SwitchyardRouter`) + selects a named algorithm from a builder registry. Because the registry is heterogeneous, the + config path standardizes decision metadata to a **serializable structured value** instead of the + concrete `D`; the observability contract is preserved either way. That erasure is the one cost of + the config path. + +Errors are returned as `Box` today (a typed `OptError` is a future cleanup), so illegal +sequences (optimize-before-feed, out-of-phase response) surface as errors, not panics. Routing fails +**open**: when a decision can't be made safely (e.g. an unparseable classifier score) the algorithm +keeps traffic on a safe tier and records why. + +## POC Gaps + +The code in `src/` diverges from this design in known ways: + +- **`AgentApiOptInput` has no `Signal` variant** (only Request/Response/Metadata), so event/ + signal-driven optimization is designed but not yet wired. Its `Response` variant also currently + carries a request-shaped struct (completion in `prompt`) rather than `AgentApiResponse` — a rough + edge the client wrapper works around when feeding responses back. +- **`EnrichmentData` lacks `tool_id`** — the tool-correlation dimension shown above is proposed, not + built. +- **No observability layer** (spans, the metrics sink, token/timing/failure tagging) — + `decision_reasoning`/`decision_info` are the hook it will build on, and they already exist. +- **No config-driven construction** — `RoutedClient` covers the direct/blackbox path, but + `SwitchyardOptimizer`/`SwitchyardRouter`, the builder registry, and decision-metadata erasure do + not exist yet. +- **Errors are `Box`**, not the typed `OptError` the design calls for. +- **Requests are single-prompt** — a message model, params, and tool schema are open questions. + +Keep this section accurate. Update the sketches above if the traits in `src/` change. diff --git a/crates/libsy/README.md b/crates/libsy/README.md new file mode 100644 index 00000000..cd8645bf --- /dev/null +++ b/crates/libsy/README.md @@ -0,0 +1,192 @@ +# libsy — Switchyard-Lib + +A lightweight library for multi-LLM agent optimization, with **routing** as the first case. It +decides — statefully, using more than the request — which model(s) to call and how, and never +performs the call itself ("ask, don't call"), which keeps it provider- and transport-agnostic. +This README shows the scope and shape through code; the narrative design is in +[`DESIGN.md`](DESIGN.md). + +## Core: two traits + +The whole core is a stateful, per-session state machine plus a factory that mints one per session. + +```rust +// Fed inputs (request, model responses, metadata); asked to decide what to do next. +#[async_trait] +pub trait AgentApiOptimizer: Send + Sync { + async fn feed(&mut self, input: AgentApiOptInput, enrichment: EnrichmentData) + -> Result<(), Box>; + async fn optimize(&mut self) -> Result, Box>; +} + +// Factory: a fresh optimizer per session. `D` is the algorithm's typed decision metadata. +pub trait AgentApiOptAlgorithm: Send + Sync { + fn optimizer(&self) -> Box>; +} +``` + +Supporting types (in `lib.rs`): + +```rust +pub enum Decision { + ModelInference(AgentApiOptimizerResponse), // host: make these calls, feed results back + Return(), // done: hand control back to the agent +} + +pub enum AgentApiOptInput { + Request(AgentApiRequest), + Response(AgentApiRequest), // MISSING: reuses the request struct (completion in `prompt`); + // should carry AgentApiResponse + token usage/latency + Metadata(BTreeMap), + // MISSING: no `Signal(..)` variant yet for agentic-stack events (tool/task/budget/telemetry) +} + +pub struct AgentApiRequest { pub prompt: String, pub model: String } // MISSING: single-prompt only +pub struct AgentApiResponse { pub completion: String } + +pub struct EnrichmentData { // correlation carried on every feed + pub session_id: Option, pub agent_id: Option, + pub task_id: Option, pub correlation_id: Option, + pub extra_metadata: Option>, + // MISSING: no `tool_id` yet +} +``` + +## Wrapper: a blackbox LLM client + +`RoutedClient` runs the `feed`/`optimize` loop for you, so routing becomes a drop-in for an ordinary +client: request in, response out. The one bit of I/O the core does not own is the model call, which +you supply as a `ModelCaller` (or use the built-in HTTP one). + +```rust +impl RoutedClient { + pub fn new(algorithm: Box>, caller: Box) -> Self; + // Convenience: default OpenAI-compatible HTTP caller. + pub fn with_http(algorithm: Box>, base_url, api_key) -> Self; + + async fn complete(&self, request: AgentApiRequest) -> Result>; + async fn complete_with(&self, request, enrichment) -> Result>; +} + +#[async_trait] +pub trait ModelCaller: Send + Sync { + async fn call(&self, request: AgentApiRequest) -> Result>; +} +``` + +## Implementing a router (LLM classifier, minimized) + +A multi-round router: run a classifier model, then route to a strong/weak model by its score. This +is why `optimize` returns `ModelInference` more than once — the host calls back in with each +response. (Full, error-checked version in [`src/llm_class.rs`](src/llm_class.rs); a one-round +weighted-random router is in [`src/rand.rs`](src/rand.rs).) + +```rust +#[derive(Clone)] +pub struct LlmClassifier { pub classifier: String, pub strong: String, pub weak: String, pub threshold: f64 } + +// `D = ()`: this router attaches no typed decision metadata (kept minimal here). +impl AgentApiOptAlgorithm<()> for LlmClassifier { + fn optimizer(&self) -> Box> { + Box::new(Session { cfg: self.clone(), phase: Phase::Classify, prompt: None, score: None }) + } +} + +enum Phase { Classify, Route, Done } +struct Session { cfg: LlmClassifier, phase: Phase, prompt: Option, score: Option } + +#[async_trait] +impl AgentApiOptimizer<()> for Session { + async fn feed(&mut self, input: AgentApiOptInput, _e: EnrichmentData) -> Result<(), Box> { + match input { + AgentApiOptInput::Request(r) => self.prompt = Some(r.prompt), + // Responses arrive in phase order: first the classifier score, then the final answer. + AgentApiOptInput::Response(r) => match self.phase { + Phase::Classify => { self.score = r.prompt.trim().parse().ok(); self.phase = Phase::Route; } + _ => self.phase = Phase::Done, + }, + _ => {} + } + Ok(()) + } + + async fn optimize(&mut self) -> Result, Box> { + let prompt = self.prompt.clone().ok_or("optimize before request was fed")?; + Ok(match self.phase { + Phase::Classify => call(&self.cfg.classifier, &format!("score 0..1: {prompt}")), + Phase::Route => { + // Fail open: an unparseable score (None) routes to the strong model. + let strong = self.score.map_or(true, |s| s >= self.cfg.threshold); + call(if strong { &self.cfg.strong } else { &self.cfg.weak }, &prompt) + } + Phase::Done => Decision::Return(), + }) + } +} + +// Emit a single model call for the host to perform. +fn call(model: &str, prompt: &str) -> Decision<()> { + Decision::ModelInference(AgentApiOptimizerResponse { + requests: vec![AgentApiRequest { prompt: prompt.into(), model: model.into() }], + enrichment_data: vec![], decision_reasoning: None, decision_info: None, + }) +} +``` + +## Using the wrapper (research agent) + +The agent owns a `RoutedClient` and nothing routing-specific; routing is one `complete` call. Here +it uses the multi-round LLM classifier, so a single `complete` triggers two model calls (classify, +then route) — invisible to the agent. Runnable: [`examples/research_agent.rs`](examples/research_agent.rs). + +```rust +// Configure routing once, in one place. The classifier scores each request with a +// model call, then routes strong/weak — two rounds the agent never sees. +let algorithm = LlmClassifierAlgorithm { + classifier_model: "classifier/model".into(), strong_model: "strong/model".into(), + weak_model: "weak/model".into(), threshold: 0.5 }; +let client = RoutedClient::new(Box::new(algorithm), Box::new(StubCaller)); // or ::with_http(..) + +// ...inside the agent, the ONLY integration point — one call, one response +// (the classifier + routed calls happen under the hood). "auto" is a placeholder +// the router replaces; the agent never learns which model served it. +let answer = client.complete(AgentApiRequest { prompt: step, model: "auto".into() }).await?; +``` + +Swapping in a one-round `RandomRouterAlgorithm` or `with_http` needs **no** change to the agent. + +## Using the core API directly (research agent) + +Same trivial agent, but driving `feed`/`optimize` by hand — exactly what the wrapper hides. Reach +for this when you need the raw loop (custom transport per round, inspecting `decision_reasoning`, +interleaving your own signals). Runnable: [`examples/research_agent_core.rs`](examples/research_agent_core.rs). + +```rust +let mut optimizer = algorithm.optimizer(); // fresh, isolated state per request +optimizer.feed(AgentApiOptInput::Request(req), EnrichmentData::default()).await?; + +let mut last = None; +// One round for weighted-random routing, N for the classifier — the loop is identical. +while let Decision::ModelInference(decision) = optimizer.optimize().await? { + // decision.decision_reasoning / decision_info explain the route (logging/eval); optional. + for req in decision.requests { + let model = req.model.clone(); + let answer = call_model(&req).await; // the host owns the transport ("ask, don't call") + // MISSING: Response reuses AgentApiRequest (completion goes in `prompt`). + optimizer.feed(AgentApiOptInput::Response(AgentApiRequest { + prompt: answer.completion.clone(), model }), EnrichmentData::default()).await?; + last = Some(answer); // the routed answer; the classifier call is consumed internally + } +} +``` + +## Not yet built + +Annotated inline above; collected here: + +- **`Signal` inputs** for agentic-stack events (tool/task/budget/telemetry) — routing on more than the request. +- **`Response` should carry `AgentApiResponse`** (with token usage/latency), not reuse the request struct. +- **`tool_id`** on `EnrichmentData` to complete session/agent/task/tool correlation. +- **Observability** — spans + a metrics sink for token counts/timings/failures (`decision_reasoning`/`decision_info` are the hook). +- **Config-driven construction** — a `SwitchyardOptimizer`/`SwitchyardRouter` built from config over a builder registry. +- **Typed errors** (`OptError`) instead of `Box`, and a **message model** (system/history, params, tools) beyond a single prompt. diff --git a/crates/libsy/examples/research_agent.rs b/crates/libsy/examples/research_agent.rs new file mode 100644 index 00000000..e4ce3c28 --- /dev/null +++ b/crates/libsy/examples/research_agent.rs @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Minimal example: dropping the `RoutedClient` wrapper into a research agent, +//! routed by a multi-round LLM classifier. +//! +//! Almost everything is a stub — the point is the *integration*. The agent makes +//! one `client.complete(..)` call per step; under the hood the wrapper runs the +//! classifier call *and then* the routed call (see the two `model call:` lines +//! per step). That multi-step routing is completely invisible to the agent. Run: +//! cargo run -p libsy --example research_agent + +use async_trait::async_trait; +use libsy::client::{ModelCaller, RoutedClient}; +use libsy::llm_class::{ClassifierRoutingDecision, LlmClassifierAlgorithm}; +use libsy::{AgentApiRequest, AgentApiResponse}; +use std::error::Error; + +const CLASSIFIER: &str = "classifier/model"; +const STRONG: &str = "strong/model"; +const WEAK: &str = "weak/model"; + +/// Stub transport. Real integrators supply their own `ModelCaller`, or use +/// `RoutedClient::with_http(algorithm, base_url, api_key)`. +struct StubCaller; + +#[async_trait] +impl ModelCaller for StubCaller { + async fn call(&self, req: AgentApiRequest) -> Result> { + // Printing each call makes the multi-step routing visible: the classifier + // call first, then the routed model call. + println!(" -> model call: {}", req.model); + let completion = if req.model == CLASSIFIER { + // The classifier is asked for a strong-win-rate score; 0.8 >= threshold + // routes the follow-up call to the strong model. + "0.8".to_string() + } else { + format!("answer from {}", req.model) + }; + Ok(AgentApiResponse { completion }) + } +} + +/// A research agent that owns a routed client and nothing routing-specific. +struct ResearchAgent { + client: RoutedClient, +} + +impl ResearchAgent { + /// Trivial plan: one lookup per sub-question (stub). + fn plan(&self, question: &str) -> Vec { + vec![format!("look up: {question}")] + } + + /// Trivial synthesis (stub). + fn synthesize(&self, notes: Vec) -> String { + notes.join("\n") + } + + async fn run(&self, question: &str) -> Result> { + let mut notes = Vec::new(); + for step in self.plan(question) { + // The only integration point: one call from the agent's view. The + // classifier + routed calls happen inside the wrapper. `model: "auto"` + // is a placeholder the router replaces. + let answer = self + .client + .complete(AgentApiRequest { + prompt: step, + model: "auto".to_string(), + }) + .await?; + notes.push(answer.completion); + } + Ok(self.synthesize(notes)) + } +} + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Box> { + // Configure routing once. The classifier scores each request with a model call, + // then routes to the strong or weak model — a two-round algorithm the agent is + // oblivious to. Swapping in a one-round `RandomRouterAlgorithm` (or + // `RoutedClient::with_http(..)`) needs no change to `ResearchAgent`. + let algorithm = LlmClassifierAlgorithm { + classifier_model: CLASSIFIER.to_string(), + strong_model: STRONG.to_string(), + weak_model: WEAK.to_string(), + threshold: 0.5, + }; + let client = RoutedClient::new(Box::new(algorithm), Box::new(StubCaller)); + + let agent = ResearchAgent { client }; + println!("{}", agent.run("what is switchyard?").await?); + Ok(()) +} diff --git a/crates/libsy/examples/research_agent_core.rs b/crates/libsy/examples/research_agent_core.rs new file mode 100644 index 00000000..6787dfcf --- /dev/null +++ b/crates/libsy/examples/research_agent_core.rs @@ -0,0 +1,130 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Second research-agent example — driving the CORE `feed`/`optimize` API by +//! hand (no `RoutedClient`), routed by a multi-round LLM classifier. Watch +//! `routed_call()` loop twice per request: once for the classifier call, once for +//! the routed call (the two `model call:` lines). `RoutedClient` hides exactly this. +//! Run with: +//! cargo run -p libsy --example research_agent_core + +use libsy::llm_class::{ClassifierRoutingDecision, LlmClassifierAlgorithm}; +use libsy::{ + AgentApiOptAlgorithm, AgentApiOptInput, AgentApiRequest, AgentApiResponse, Decision, + EnrichmentData, +}; +use std::error::Error; + +const CLASSIFIER: &str = "classifier/model"; +const STRONG: &str = "strong/model"; +const WEAK: &str = "weak/model"; + +/// Stand-in for the host's real model call. The core API never makes the call +/// itself: `optimize()` hands back an `AgentApiRequest` and the host performs +/// the call and feeds the result back ("ask, don't call"). That is what keeps +/// libsy transport- and provider-agnostic — no HTTP client in the core. +async fn call_model(req: &AgentApiRequest) -> AgentApiResponse { + // Printed so the two rounds (classifier, then routed target) are visible. + println!(" -> model call: {}", req.model); + let completion = if req.model == CLASSIFIER { + "0.8".to_string() // strong-win-rate score; >= threshold routes to strong + } else { + format!("answer from {}", req.model) + }; + AgentApiResponse { completion } +} + +struct ResearchAgent { + // The agent owns the algorithm *factory*, not an optimizer: a fresh, isolated + // optimizer is minted per request (routing state never leaks between requests). + algorithm: Box>, +} + +impl ResearchAgent { + /// Trivial plan: one lookup per question (stub). + fn plan(&self, question: &str) -> Vec { + vec![format!("look up: {question}")] + } + + /// The core-API equivalent of `RoutedClient::complete`: mint a per-request + /// optimizer, feed the request, and drive `feed`/`optimize` to `Return`, + /// making each model call the optimizer asks for. This is the reusable loop + /// the wrapper hides — the agent's `run` just builds requests and calls it. + async fn routed_call( + &self, + request: AgentApiRequest, + ) -> Result> { + // 1. Mint a per-request optimizer and feed the inbound request. `feed` + // takes `&mut self`, so a session is single-owner: the host serializes + // inputs (here, trivially, one request). + let mut optimizer = self.algorithm.optimizer(); + optimizer + .feed( + AgentApiOptInput::Request(request), + EnrichmentData::default(), + ) + .await?; + + // 2. Loop: the optimizer decides which model calls to make; the host makes + // them and feeds the responses back. `Return` ends the session. The + // classifier runs this loop TWICE (classify, then route) — but the loop + // code is identical to a single-round router's (the whole point). + let mut last = None; + while let Decision::ModelInference(decision) = optimizer.optimize().await? { + // `decision.decision_reasoning` / `decision.decision_info` explain the + // route (great for logging/eval); nothing forces you to consume them. + for req in decision.requests { + let model = req.model.clone(); + let answer = call_model(&req).await; + // MISSING: the `Response` input reuses `AgentApiRequest` (completion + // stuffed into `prompt`) instead of carrying `AgentApiResponse` — a + // known rough edge. Token usage / latency also have nowhere to go yet. + optimizer + .feed( + AgentApiOptInput::Response(AgentApiRequest { + prompt: answer.completion.clone(), + model, + }), + EnrichmentData::default(), + ) + .await?; + last = Some(answer); + } + } + + // The last response is the routed answer; the intermediate classifier call + // is consumed inside the loop and never surfaces here. + last.ok_or_else(|| "optimizer returned before requesting any model call".into()) + } + + async fn run(&self, question: &str) -> Result> { + let mut notes = Vec::new(); + for step in self.plan(question) { + // The agent builds a request (`model: "auto"` is a placeholder the + // router replaces) and hands it off — it never drives the loop itself. + let request = AgentApiRequest { + prompt: step, + model: "auto".to_string(), + }; + notes.push(self.routed_call(request).await?.completion); + } + Ok(notes.join("\n")) + } +} + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Box> { + // Configure routing once. The classifier is a two-round algorithm; + // routed_call's loop does not care how many rounds the algorithm uses. + let algorithm = LlmClassifierAlgorithm { + classifier_model: CLASSIFIER.to_string(), + strong_model: STRONG.to_string(), + weak_model: WEAK.to_string(), + threshold: 0.5, + }; + let agent = ResearchAgent { + algorithm: Box::new(algorithm), + }; + println!("{}", agent.run("what is switchyard?").await?); + Ok(()) +} diff --git a/crates/libsy/src/client.rs b/crates/libsy/src/client.rs new file mode 100644 index 00000000..70f6da72 --- /dev/null +++ b/crates/libsy/src/client.rs @@ -0,0 +1,485 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! An LLM-client-style wrapper over the optimizer interfaces. +//! +//! [`RoutedClient`] hides the `feed` / `optimize` loop behind a single +//! `complete` call, turning routing (and any other optimization) into a black +//! box that can stand in for an ordinary LLM client. The actual model call is +//! delegated to a [`ModelCaller`]: integrators plug in their own transport, or +//! use the built-in [`HttpCaller`] against any OpenAI-compatible endpoint. + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use std::error::Error; + +use crate::{ + AgentApiOptAlgorithm, AgentApiOptInput, AgentApiRequest, AgentApiResponse, Decision, + EnrichmentData, +}; + +/// Observes the generic routed-client loop without owning routing policy. +/// +/// Callbacks are registered on [`RoutedClientBuilder`] and run for every +/// algorithm the client hosts. Algorithm-specific decision metadata remains in +/// `D`; the callback layer only observes the common lifecycle. +#[async_trait] +pub trait RouterCallback: Send + Sync { + /// Called after each optimizer decision, before any requested model calls. + async fn on_decision( + &self, + _decision: &Decision, + _enrichment: &EnrichmentData, + ) -> Result<(), Box> { + Ok(()) + } + + /// Called immediately before the host performs one model call. + async fn before_model_call( + &self, + _request: &AgentApiRequest, + _enrichment: &EnrichmentData, + ) -> Result<(), Box> { + Ok(()) + } + + /// Called after one model call succeeds and before its response is fed back. + async fn after_model_call( + &self, + _request: &AgentApiRequest, + _response: &AgentApiResponse, + _enrichment: &EnrichmentData, + ) -> Result<(), Box> { + Ok(()) + } +} + +/// Performs a single model call. This is the one piece of I/O the client does +/// not own — integrators implement it with their own HTTP/transport stack, or +/// use [`HttpCaller`]. +#[async_trait] +pub trait ModelCaller: Send + Sync { + /// Call the model named by `request.model` with `request.prompt`, returning + /// its completion. + async fn call(&self, request: AgentApiRequest) -> Result>; +} + +/// A blackbox, LLM-client-style front end over an optimization algorithm. +/// +/// Each [`complete`](RoutedClient::complete) mints a fresh optimizer for the +/// request, drives the `feed` -> `optimize` loop to completion — performing +/// every model call the optimizer asks for via the [`ModelCaller`] — and +/// returns the final model response. Single-round routers (weighted random) and +/// multi-round ones (LLM classifier) are handled by the same loop; the caller +/// sees only "request in, response out". +pub struct RoutedClient { + algorithm: Box>, + caller: Box, + callbacks: Vec>>, +} + +impl RoutedClient { + /// Start a builder that can register callbacks before constructing the + /// client. + pub fn builder( + algorithm: Box>, + caller: Box, + ) -> RoutedClientBuilder { + RoutedClientBuilder::new(algorithm, caller) + } + + /// Build a client from an algorithm and a model caller. + pub fn new(algorithm: Box>, caller: Box) -> Self { + RoutedClient { + algorithm, + caller, + callbacks: Vec::new(), + } + } + + /// Build a client that makes model calls with the built-in [`HttpCaller`] + /// against an OpenAI-compatible endpoint at `base_url`. + pub fn with_http( + algorithm: Box>, + base_url: impl Into, + api_key: impl Into, + ) -> Self { + RoutedClient::new(algorithm, Box::new(HttpCaller::new(base_url, api_key))) + } + + async fn on_decision( + &self, + decision: &Decision, + enrichment: &EnrichmentData, + ) -> Result<(), Box> { + for callback in &self.callbacks { + callback.on_decision(decision, enrichment).await?; + } + Ok(()) + } + + async fn before_model_call( + &self, + request: &AgentApiRequest, + enrichment: &EnrichmentData, + ) -> Result<(), Box> { + for callback in &self.callbacks { + callback.before_model_call(request, enrichment).await?; + } + Ok(()) + } + + async fn after_model_call( + &self, + request: &AgentApiRequest, + response: &AgentApiResponse, + enrichment: &EnrichmentData, + ) -> Result<(), Box> { + for callback in &self.callbacks { + callback + .after_model_call(request, response, enrichment) + .await?; + } + Ok(()) + } + + /// Optimize and run `request` to completion, returning the final response. + pub async fn complete( + &self, + request: AgentApiRequest, + ) -> Result> { + self.complete_with(request, EnrichmentData::default()).await + } + + /// Like [`complete`](RoutedClient::complete), but attaches `enrichment` + /// (session/agent/task correlation) to every input fed to the optimizer. + pub async fn complete_with( + &self, + request: AgentApiRequest, + enrichment: EnrichmentData, + ) -> Result> { + let mut optimizer = self.algorithm.optimizer(); + optimizer + .feed(AgentApiOptInput::Request(request), enrichment.clone()) + .await?; + + // Round-agnostic drive loop: keep performing the model calls the + // optimizer asks for and feeding their responses back until it returns. + let mut last: Option = None; + loop { + let decision = optimizer.optimize().await?; + self.on_decision(&decision, &enrichment).await?; + match decision { + Decision::ModelInference(response) => { + for req in response.requests { + let model = req.model.clone(); + self.before_model_call(&req, &enrichment).await?; + let result = self.caller.call(req.clone()).await?; + self.after_model_call(&req, &result, &enrichment).await?; + // The Response input variant carries a request-shaped + // struct whose `prompt` holds the completion text. + optimizer + .feed( + AgentApiOptInput::Response(AgentApiRequest { + prompt: result.completion.clone(), + model, + }), + enrichment.clone(), + ) + .await?; + last = Some(result); + } + } + Decision::Return() => break, + } + } + + last.ok_or_else(|| "optimizer returned before requesting any model call".into()) + } +} + +/// Builder for registering routed-client callbacks before construction. +pub struct RoutedClientBuilder { + algorithm: Box>, + caller: Box, + callbacks: Vec>>, +} + +impl RoutedClientBuilder { + /// Create a builder from the algorithm and host-owned model caller. + pub fn new(algorithm: Box>, caller: Box) -> Self { + RoutedClientBuilder { + algorithm, + caller, + callbacks: Vec::new(), + } + } + + /// Register a callback that observes the generic optimize/call/feed loop. + pub fn callback(mut self, callback: impl RouterCallback + 'static) -> Self { + self.callbacks.push(Box::new(callback)); + self + } + + /// Build the routed client. + pub fn build(self) -> RoutedClient { + RoutedClient { + algorithm: self.algorithm, + caller: self.caller, + callbacks: self.callbacks, + } + } +} + +/// Default [`ModelCaller`]: POSTs to `{base_url}/chat/completions` on any +/// OpenAI-compatible endpoint, sending the prompt as a single user message. +pub struct HttpCaller { + client: reqwest::Client, + base_url: String, + api_key: String, +} + +impl HttpCaller { + /// Create a caller targeting `base_url` (e.g. `https://api.openai.com/v1`) + /// with the given bearer `api_key`. + pub fn new(base_url: impl Into, api_key: impl Into) -> Self { + HttpCaller { + client: reqwest::Client::new(), + base_url: base_url.into(), + api_key: api_key.into(), + } + } +} + +#[derive(Serialize)] +struct ChatMessage<'a> { + role: &'a str, + content: &'a str, +} + +#[derive(Serialize)] +struct ChatCompletionRequest<'a> { + model: &'a str, + messages: Vec>, +} + +#[derive(Deserialize)] +struct ChatCompletionResponse { + choices: Vec, +} + +#[derive(Deserialize)] +struct Choice { + message: ResponseMessage, +} + +#[derive(Deserialize)] +struct ResponseMessage { + content: String, +} + +#[async_trait] +impl ModelCaller for HttpCaller { + async fn call(&self, request: AgentApiRequest) -> Result> { + let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/')); + let payload = ChatCompletionRequest { + model: &request.model, + messages: vec![ChatMessage { + role: "user", + content: &request.prompt, + }], + }; + let response = self + .client + .post(url) + .bearer_auth(&self.api_key) + .json(&payload) + .send() + .await? + .error_for_status()? + .json::() + .await?; + let completion = response + .choices + .into_iter() + .next() + .map(|choice| choice.message.content) + .ok_or("model response contained no choices")?; + Ok(AgentApiResponse { completion }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::llm_class::LlmClassifierAlgorithm; + use crate::rand::{RandomRouterAlgorithm, WeightedModel}; + use std::sync::{Arc, Mutex}; + + /// A no-network caller that records requests and returns scripted answers: a + /// parseable score for the classifier model, an answer tagged with the model + /// otherwise. + struct RecordingCaller { + classifier_model: String, + calls: Arc>>, + } + + #[async_trait] + impl ModelCaller for RecordingCaller { + async fn call(&self, request: AgentApiRequest) -> Result> { + let completion = if request.model == self.classifier_model { + "0.9".to_string() + } else { + format!("answer from {}", request.model) + }; + self.calls + .lock() + .map_err(|_| "recording lock poisoned")? + .push(request); + Ok(AgentApiResponse { completion }) + } + } + + struct RecordingCallback { + events: Arc>>, + } + + impl RecordingCallback { + fn push(&self, event: impl Into) -> Result<(), Box> { + self.events + .lock() + .map_err(|_| "callback lock poisoned")? + .push(event.into()); + Ok(()) + } + } + + #[async_trait] + impl RouterCallback for RecordingCallback { + async fn on_decision( + &self, + decision: &Decision, + _enrichment: &EnrichmentData, + ) -> Result<(), Box> { + match decision { + Decision::ModelInference(_) => self.push("decision:model_inference"), + Decision::Return() => self.push("decision:return"), + } + } + + async fn before_model_call( + &self, + request: &AgentApiRequest, + _enrichment: &EnrichmentData, + ) -> Result<(), Box> { + self.push(format!("before:{}", request.model)) + } + + async fn after_model_call( + &self, + request: &AgentApiRequest, + response: &AgentApiResponse, + _enrichment: &EnrichmentData, + ) -> Result<(), Box> { + self.push(format!("after:{}:{}", request.model, response.completion)) + } + } + + #[tokio::test] + async fn rand_client_routes_once_and_returns_the_response() -> Result<(), Box> { + let algorithm = RandomRouterAlgorithm { + models: vec![WeightedModel::new("frontier/model", 1.0)], + rng_seed: Some(7), + }; + let calls = Arc::new(Mutex::new(Vec::new())); + let caller = RecordingCaller { + classifier_model: "unused".to_string(), + calls: Arc::clone(&calls), + }; + let events = Arc::new(Mutex::new(Vec::new())); + let client = RoutedClient::builder(Box::new(algorithm), Box::new(caller)) + .callback(RecordingCallback { + events: Arc::clone(&events), + }) + .build(); + + let response = client + .complete(AgentApiRequest { + prompt: "hi".to_string(), + model: "auto".to_string(), + }) + .await?; + + let recorded = calls.lock().map_err(|_| "lock poisoned")?; + assert_eq!(recorded.len(), 1); + assert_eq!(recorded[0].model, "frontier/model"); + assert_eq!(response.completion, "answer from frontier/model"); + drop(recorded); + + let events = events.lock().map_err(|_| "lock poisoned")?.clone(); + assert_eq!( + events, + vec![ + "decision:model_inference", + "before:frontier/model", + "after:frontier/model:answer from frontier/model", + "decision:return", + ] + ); + Ok(()) + } + + #[tokio::test] + async fn classifier_client_drives_multi_round_and_returns_routed_response( + ) -> Result<(), Box> { + let algorithm = LlmClassifierAlgorithm { + classifier_model: "router/classifier".to_string(), + strong_model: "frontier/model".to_string(), + weak_model: "cheap/model".to_string(), + threshold: 0.5, + }; + let calls = Arc::new(Mutex::new(Vec::new())); + let caller = RecordingCaller { + classifier_model: "router/classifier".to_string(), + calls: Arc::clone(&calls), + }; + let events = Arc::new(Mutex::new(Vec::new())); + let client = RoutedClient::builder(Box::new(algorithm), Box::new(caller)) + .callback(RecordingCallback { + events: Arc::clone(&events), + }) + .build(); + + let response = client + .complete(AgentApiRequest { + prompt: "prove it".to_string(), + model: "auto".to_string(), + }) + .await?; + + let recorded = calls.lock().map_err(|_| "lock poisoned")?; + // Two model calls: the classifier, then the routed target. + assert_eq!(recorded.len(), 2); + assert_eq!(recorded[0].model, "router/classifier"); + assert!(recorded[0].prompt.contains("prove it")); + assert_eq!(recorded[1].model, "frontier/model"); // score 0.9 >= 0.5 -> strong + // The returned response is the routed call's, not the classifier's. + assert_eq!(response.completion, "answer from frontier/model"); + drop(recorded); + + let events = events.lock().map_err(|_| "lock poisoned")?.clone(); + assert_eq!( + events, + vec![ + "decision:model_inference", + "before:router/classifier", + "after:router/classifier:0.9", + "decision:model_inference", + "before:frontier/model", + "after:frontier/model:answer from frontier/model", + "decision:return", + ] + ); + Ok(()) + } +} diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs new file mode 100644 index 00000000..060e066d --- /dev/null +++ b/crates/libsy/src/lib.rs @@ -0,0 +1,100 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Switchyard library crate. + +pub mod client; +pub mod llm_class; +pub mod rand; + +use async_trait::async_trait; +use std::collections::BTreeMap; +use std::error::Error; + +/// AgentApiRequest internal representation of an llm request. +/// This structure is designed to be converted to a from provider specifiic structs without loosing +/// information +#[derive(Clone, Debug)] +pub struct AgentApiRequest { + pub prompt: String, + pub model: String, +} + +/// AgentApiRequest internal representation of an llm response. +/// This structure is designed to be converted to a from provider specifiic structs without loosing +/// information +#[derive(Clone, Debug)] +pub struct AgentApiResponse { + pub completion: String, +} + +/// enrichment and correlation data use for routing the request +#[derive(Clone, Default)] +pub struct EnrichmentData { + pub session_id: Option, + pub agent_id: Option, + pub task_id: Option, + pub correlation_id: Option, + pub extra_metadata: Option>, +} + +/// A Routing / optimization decision made by the AgentApiOptimizer +pub enum Decision { + ModelInference(AgentApiOptimizerResponse), + Return(), +} + +/// Input to the AgentApiOptimizer, can be a request, response or metadata +pub enum AgentApiOptInput { + Request(AgentApiRequest), + Response(AgentApiRequest), + Metadata(BTreeMap), +} + +/// The response from the AgentApiOptimizer, containing the optimized requests and any additional +/// data. +pub struct AgentApiOptimizerResponse { + pub requests: Vec, + pub enrichment_data: Vec, + pub decision_reasoning: Option, + pub decision_info: Option, +} + +/// AgentApiOptimizer is a stateful optimizer that is feed requests, responses and metadata and can +/// make routing / optimization decisions based on the input and the current state. +/// A single instance of an AgentApiOptimizer is created for each session and is used to optimize +/// the requests for that session. +#[async_trait] +pub trait AgentApiOptimizer: Send + Sync { + /// Feed the optimizer with a new input and enrichment data. The optimizer can use this data to + async fn feed( + &mut self, + _input: AgentApiOptInput, + _enrichment: EnrichmentData, + ) -> Result<(), Box> { + Ok(()) + } + + /// Make a routing / optimization decision based on the current state of the optimizer. + /// Return Types: + /// ModelInference: The optimizer has decided to make a model inference request. The caller + /// should make the listed model calls and pass the responses back to the optimizer via + /// feed() for further optimization. + /// Return: The optimizer has decided to return a response to the agent (e.g. for tool + /// excution) the caller should pass control to the calling agent. + async fn optimize(&mut self) -> Result, Box> { + Ok(Decision::ModelInference(AgentApiOptimizerResponse { + requests: Vec::new(), + enrichment_data: Vec::new(), + decision_reasoning: None, + decision_info: None, + })) + } +} + +/// AgentApiOptAlgorithm is a factory for creating instances of AgentApiOptimizer. It is used to +/// create a new optimizer for each session. +pub trait AgentApiOptAlgorithm: Send + Sync { + /// Create a new instance of the AgentApiOptimizer for the given session. + fn optimizer(&self) -> Box>; +} diff --git a/crates/libsy/src/llm_class.rs b/crates/libsy/src/llm_class.rs new file mode 100644 index 00000000..df0b2114 --- /dev/null +++ b/crates/libsy/src/llm_class.rs @@ -0,0 +1,380 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! LLM-classifier router built on the AgentApi optimizer interfaces. +//! +//! Unlike a local ML classifier (which scores a prompt in-process), an LLM +//! classifier needs its own model call to classify the request. That maps +//! directly onto the optimizer's multi-round `ModelInference` -> `feed` -> +//! `optimize` loop: the router first asks the caller to run a classifier model, +//! then — once the score is fed back — asks the caller to run the routed target +//! model, and finally returns control to the agent. + +use async_trait::async_trait; +use std::error::Error; + +use crate::{ + AgentApiOptAlgorithm, AgentApiOptInput, AgentApiOptimizer, AgentApiOptimizerResponse, + AgentApiRequest, Decision, EnrichmentData, +}; + +/// Preamble prepended to the user prompt when asking the classifier model for a +/// strong-win-rate score. +const CLASSIFIER_PROMPT_PREAMBLE: &str = "Rate how strongly this request needs a frontier model. \ + Reply with a single strong-win-rate score in [0, 1]:\n"; + +/// The tier a classifier score selected. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ClassifierTier { + Strong, + Weak, +} + +impl ClassifierTier { + /// Stable string form of the tier, used in decision reasoning. + pub fn as_str(self) -> &'static str { + match self { + ClassifierTier::Strong => "strong", + ClassifierTier::Weak => "weak", + } + } +} + +/// Decision info attached to the routed `ModelInference` decision. +/// +/// This is the concrete `D` carried by [`Decision`] / [`AgentApiOptimizerResponse`] +/// for the classifier router. +#[derive(Clone, Debug, PartialEq)] +pub struct ClassifierRoutingDecision { + /// Strong-win-rate score parsed from the classifier response, or `None` when + /// the response could not be parsed (the router then defaults to strong). + pub score: Option, + /// Threshold the score was compared against. + pub threshold: f64, + /// Tier chosen for this request. + pub tier: ClassifierTier, + /// Model the request was rewritten to target. + pub selected_model: String, +} + +/// Factory that mints a fresh [`LlmClassifierRouter`] per session. +/// +/// Generalizes Switchyard's RouteLLM strong/weak selection to an LLM-driven +/// classifier expressed with the AgentApi optimizer interfaces. +pub struct LlmClassifierAlgorithm { + /// Model id used to classify the request. + pub classifier_model: String, + /// Model id chosen when the score meets the threshold. + pub strong_model: String, + /// Model id chosen when the score is below the threshold. + pub weak_model: String, + /// Score threshold at or above which the strong tier is selected. + pub threshold: f64, +} + +impl AgentApiOptAlgorithm for LlmClassifierAlgorithm { + fn optimizer(&self) -> Box> { + Box::new(LlmClassifierRouter { + classifier_model: self.classifier_model.clone(), + strong_model: self.strong_model.clone(), + weak_model: self.weak_model.clone(), + threshold: self.threshold, + phase: Phase::AwaitingRequest, + pending_request: None, + score: None, + }) + } +} + +/// Where the router is in its classify -> route -> return lifecycle. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Phase { + /// No request buffered yet. + AwaitingRequest, + /// Request buffered; next `optimize` emits the classifier call. + Classify, + /// Classifier call emitted; waiting for its response to be fed. + AwaitingScore, + /// Score received; next `optimize` emits the routed target call. + Route, + /// Routed call emitted; waiting for the target response to be fed. + AwaitingResponse, + /// Target response received; next `optimize` returns control to the agent. + Done, +} + +/// Per-session LLM-classifier router. +/// +/// Flow: `feed` the request, then `optimize` (emits the classifier call); +/// `feed` the classifier response, then `optimize` (emits the routed call); +/// `feed` the target response, then `optimize` (returns `Return`). +pub struct LlmClassifierRouter { + classifier_model: String, + strong_model: String, + weak_model: String, + threshold: f64, + phase: Phase, + pending_request: Option, + score: Option, +} + +impl LlmClassifierRouter { + /// Build the routing decision from the parsed score, defaulting to the + /// strong tier when the classifier output could not be parsed. + fn decide(&self) -> ClassifierRoutingDecision { + let (tier, selected_model) = match self.score { + Some(score) if score >= self.threshold => { + (ClassifierTier::Strong, self.strong_model.clone()) + } + Some(_) => (ClassifierTier::Weak, self.weak_model.clone()), + // Defensive default: keep traffic flowing on the strong tier when the + // classifier response was unusable. + None => (ClassifierTier::Strong, self.strong_model.clone()), + }; + ClassifierRoutingDecision { + score: self.score, + threshold: self.threshold, + tier, + selected_model, + } + } +} + +#[async_trait] +impl AgentApiOptimizer for LlmClassifierRouter { + async fn feed( + &mut self, + input: AgentApiOptInput, + _enrichment: EnrichmentData, + ) -> Result<(), Box> { + match input { + AgentApiOptInput::Request(request) => { + self.pending_request = Some(request); + self.phase = Phase::Classify; + } + AgentApiOptInput::Response(response) => match self.phase { + // First response is the classifier's output; parse a score. + Phase::AwaitingScore => { + self.score = response.prompt.trim().parse::().ok(); + self.phase = Phase::Route; + } + // Second response is the routed model's output; ready to return. + Phase::AwaitingResponse => self.phase = Phase::Done, + _ => { + return Err( + "classifier router received a response outside a pending model call".into(), + ) + } + }, + AgentApiOptInput::Metadata(_) => {} + } + Ok(()) + } + + async fn optimize(&mut self) -> Result, Box> { + match self.phase { + Phase::AwaitingRequest => Err("optimize called before a request was fed".into()), + Phase::Classify => { + let user_prompt = self + .pending_request + .as_ref() + .ok_or("classifier router has no buffered request")? + .prompt + .clone(); + self.phase = Phase::AwaitingScore; + let classifier_request = AgentApiRequest { + model: self.classifier_model.clone(), + prompt: format!("{CLASSIFIER_PROMPT_PREAMBLE}{user_prompt}"), + }; + Ok(Decision::ModelInference(AgentApiOptimizerResponse { + requests: vec![classifier_request], + enrichment_data: Vec::new(), + decision_reasoning: Some(format!( + "classifying request via {}", + self.classifier_model + )), + decision_info: None, + })) + } + Phase::AwaitingScore => { + Err("optimize called before the classifier response was fed".into()) + } + Phase::Route => { + let mut request = self + .pending_request + .take() + .ok_or("classifier router has no buffered request")?; + let decision = self.decide(); + request.model = decision.selected_model.clone(); + self.phase = Phase::AwaitingResponse; + Ok(Decision::ModelInference(AgentApiOptimizerResponse { + requests: vec![request], + enrichment_data: Vec::new(), + decision_reasoning: Some(format!( + "classifier score {:?} vs threshold {}; selected {} ({})", + decision.score, + decision.threshold, + decision.selected_model, + decision.tier.as_str() + )), + decision_info: Some(decision), + })) + } + Phase::AwaitingResponse => { + Err("optimize called before the model response was fed".into()) + } + Phase::Done => Ok(Decision::Return()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Empty enrichment payload for feeds under test. + fn enrichment() -> EnrichmentData { + EnrichmentData { + session_id: None, + agent_id: None, + task_id: None, + correlation_id: None, + extra_metadata: None, + } + } + + fn request(prompt: &str, model: &str) -> AgentApiRequest { + AgentApiRequest { + prompt: prompt.to_string(), + model: model.to_string(), + } + } + + fn algorithm(threshold: f64) -> LlmClassifierAlgorithm { + LlmClassifierAlgorithm { + classifier_model: "router/classifier".to_string(), + strong_model: "frontier/model".to_string(), + weak_model: "cheap/model".to_string(), + threshold, + } + } + + /// Drive the classify -> route -> return flow, feeding `score_text` as the + /// mocked classifier response, and return the routed decision. + async fn run_flow( + threshold: f64, + user_prompt: &str, + score_text: &str, + ) -> Result<(String, ClassifierRoutingDecision), Box> { + let mut optimizer = algorithm(threshold).optimizer(); + + // 1. Feed the user request and ask the router what to do. + optimizer + .feed( + AgentApiOptInput::Request(request(user_prompt, "client/model")), + enrichment(), + ) + .await?; + let classifier_model = match optimizer.optimize().await? { + Decision::ModelInference(response) => { + // The first inference is the classifier call itself. + let call = &response.requests[0]; + assert!( + call.prompt.contains(user_prompt), + "classifier prompt missing user text" + ); + call.model.clone() + } + Decision::Return() => return Err("expected classifier ModelInference".into()), + }; + assert_eq!(classifier_model, "router/classifier"); + + // 2. Mock the classifier LLM call by feeding its score back. + optimizer + .feed( + AgentApiOptInput::Response(request(score_text, &classifier_model)), + enrichment(), + ) + .await?; + let (routed_model, decision) = match optimizer.optimize().await? { + Decision::ModelInference(response) => { + let decision = response + .decision_info + .ok_or("expected decision info on routed ModelInference")?; + (response.requests[0].model.clone(), decision) + } + Decision::Return() => return Err("expected routed ModelInference".into()), + }; + + // 3. Mock the routed model call; the next optimize returns to the agent. + optimizer + .feed( + AgentApiOptInput::Response(request("mocked completion", &routed_model)), + enrichment(), + ) + .await?; + match optimizer.optimize().await? { + Decision::Return() => Ok((routed_model, decision)), + Decision::ModelInference(_) => Err("expected Return after routed response fed".into()), + } + } + + #[tokio::test] + async fn score_at_or_above_threshold_routes_strong() -> Result<(), Box> { + let (routed_model, decision) = run_flow(0.5, "solve this proof", "0.9").await?; + assert_eq!(routed_model, "frontier/model"); + assert_eq!(decision.tier, ClassifierTier::Strong); + assert_eq!(decision.score, Some(0.9)); + Ok(()) + } + + #[tokio::test] + async fn score_below_threshold_routes_weak() -> Result<(), Box> { + let (routed_model, decision) = run_flow(0.5, "say hello", "0.2").await?; + assert_eq!(routed_model, "cheap/model"); + assert_eq!(decision.tier, ClassifierTier::Weak); + assert_eq!(decision.score, Some(0.2)); + Ok(()) + } + + #[tokio::test] + async fn score_exactly_at_threshold_routes_strong() -> Result<(), Box> { + let (routed_model, decision) = run_flow(0.5, "borderline", "0.5").await?; + assert_eq!(routed_model, "frontier/model"); + assert_eq!(decision.tier, ClassifierTier::Strong); + Ok(()) + } + + #[tokio::test] + async fn unparseable_score_defaults_to_strong() -> Result<(), Box> { + let (routed_model, decision) = run_flow(0.5, "hi", "not-a-number").await?; + assert_eq!(routed_model, "frontier/model"); + assert_eq!(decision.tier, ClassifierTier::Strong); + assert_eq!(decision.score, None); + Ok(()) + } + + #[tokio::test] + async fn optimize_before_feed_errors() { + let mut optimizer = algorithm(0.5).optimizer(); + assert!(optimizer.optimize().await.is_err()); + } + + #[tokio::test] + async fn optimize_before_classifier_response_errors() -> Result<(), Box> { + let mut optimizer = algorithm(0.5).optimizer(); + optimizer + .feed( + AgentApiOptInput::Request(request("hi", "client/model")), + enrichment(), + ) + .await?; + // Emit the classifier call, then call optimize again without feeding the score. + assert!(matches!( + optimizer.optimize().await?, + Decision::ModelInference(_) + )); + assert!(optimizer.optimize().await.is_err()); + Ok(()) + } +} diff --git a/crates/libsy/src/rand.rs b/crates/libsy/src/rand.rs new file mode 100644 index 00000000..733d6636 --- /dev/null +++ b/crates/libsy/src/rand.rs @@ -0,0 +1,417 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Weighted random router built on the AgentApi optimizer interfaces. + +use async_trait::async_trait; +use std::error::Error; + +// `::rand` disambiguates the external crate from this module, which is also +// named `rand`. +use ::rand::rngs::StdRng; +use ::rand::{Rng, SeedableRng}; + +use crate::{ + AgentApiOptAlgorithm, AgentApiOptInput, AgentApiOptimizer, AgentApiOptimizerResponse, + AgentApiRequest, Decision, EnrichmentData, +}; + +/// A routing target and its relative selection weight. +/// +/// Weights are relative and need not sum to `1.0`; a target is chosen with +/// probability `weight / sum(weights)`. Weights must be finite and non-negative, +/// and at least one weight must be positive. +#[derive(Clone, Debug, PartialEq)] +pub struct WeightedModel { + /// Model id this target routes to. + pub model: String, + /// Relative selection weight; a zero weight is never selected. + pub weight: f64, +} + +impl WeightedModel { + /// Convenience constructor for a weighted target. + pub fn new(model: impl Into, weight: f64) -> Self { + WeightedModel { + model: model.into(), + weight, + } + } +} + +/// Decision info attached to a random-routing `ModelInference` decision. +/// +/// This is the concrete `D` carried by [`Decision`] / [`AgentApiOptimizerResponse`] +/// for the random router. +#[derive(Clone, Debug, PartialEq)] +pub struct RandomRoutingDecision { + /// Model the request was rewritten to target. + pub selected_model: String, + /// The random draw in `[0, total_weight)` that produced this decision. + pub draw: f64, + /// Sum of all target weights the draw was taken against. + pub total_weight: f64, +} + +/// Factory that mints a fresh [`RandomRouter`] per session. +/// +/// Generalizes Switchyard's strong/weak `RandomRoutingProfile` to N targets +/// selected by weighted random choice, expressed with the AgentApi optimizer +/// interfaces. +pub struct RandomRouterAlgorithm { + /// Weighted set of routing targets. + pub models: Vec, + /// Optional deterministic RNG seed for reproducible routing. + pub rng_seed: Option, +} + +impl AgentApiOptAlgorithm for RandomRouterAlgorithm { + fn optimizer(&self) -> Box> { + let rng = match self.rng_seed { + Some(seed) => StdRng::seed_from_u64(seed), + None => StdRng::from_entropy(), + }; + Box::new(RandomRouter { + models: self.models.clone(), + rng, + pending_request: None, + completed: false, + }) + } +} + +/// Per-session random router over a weighted set of N targets. +/// +/// Flow: the caller `feed`s the inbound request, then calls `optimize`, which +/// draws a weighted target, rewrites the request model, and returns +/// `ModelInference` so the caller performs the model call. The caller `feed`s the +/// response back and calls `optimize` again, which returns `Return` to hand +/// control to the agent. +pub struct RandomRouter { + models: Vec, + rng: StdRng, + pending_request: Option, + completed: bool, +} + +impl RandomRouter { + /// Draw a weighted target and build the routing decision for the current + /// request. Errors if no target carries positive, finite weight. + fn select(&mut self) -> Result> { + let total_weight: f64 = self + .models + .iter() + .map(|m| m.weight) + .filter(|w| w.is_finite() && *w > 0.0) + .sum(); + // `total_weight` sums only finite positive weights, so it is never NaN; + // `<= 0.0` therefore holds exactly when no target is selectable. + if total_weight <= 0.0 { + return Err("random router has no target with positive weight".into()); + } + // Draw in [0, total_weight) and walk cumulative weights. + let draw: f64 = self.rng.gen::() * total_weight; + let mut cumulative = 0.0; + for model in &self.models { + if !(model.weight.is_finite() && model.weight > 0.0) { + continue; + } + cumulative += model.weight; + if draw < cumulative { + return Ok(RandomRoutingDecision { + selected_model: model.model.clone(), + draw, + total_weight, + }); + } + } + // Floating-point rounding can leave the draw at the very top of the + // range; fall back to the last positively weighted target. + let last = self + .models + .iter() + .rev() + .find(|m| m.weight.is_finite() && m.weight > 0.0) + .ok_or("random router has no target with positive weight")?; + Ok(RandomRoutingDecision { + selected_model: last.model.clone(), + draw, + total_weight, + }) + } +} + +#[async_trait] +impl AgentApiOptimizer for RandomRouter { + async fn feed( + &mut self, + input: AgentApiOptInput, + _enrichment: EnrichmentData, + ) -> Result<(), Box> { + match input { + AgentApiOptInput::Request(request) => self.pending_request = Some(request), + // A fed response means the caller performed the model call; the next + // optimize should hand control back to the agent. + AgentApiOptInput::Response(_) => self.completed = true, + AgentApiOptInput::Metadata(_) => {} + } + Ok(()) + } + + async fn optimize(&mut self) -> Result, Box> { + // The model response has already come back; return control to the agent. + if self.completed { + return Ok(Decision::Return()); + } + let mut request = self + .pending_request + .take() + .ok_or("optimize called before a request was fed")?; + let decision = self.select()?; + request.model = decision.selected_model.clone(); + Ok(Decision::ModelInference(AgentApiOptimizerResponse { + requests: vec![request], + enrichment_data: Vec::new(), + decision_reasoning: Some(format!( + "weighted random draw {} of total weight {}; selected {}", + decision.draw, decision.total_weight, decision.selected_model + )), + decision_info: Some(decision), + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Empty enrichment payload for feeds under test. + fn enrichment() -> EnrichmentData { + EnrichmentData { + session_id: None, + agent_id: None, + task_id: None, + correlation_id: None, + extra_metadata: None, + } + } + + fn request(prompt: &str, model: &str) -> AgentApiRequest { + AgentApiRequest { + prompt: prompt.to_string(), + model: model.to_string(), + } + } + + fn algorithm(models: Vec, seed: u64) -> RandomRouterAlgorithm { + RandomRouterAlgorithm { + models, + rng_seed: Some(seed), + } + } + + /// Feed one request and return the model the optimizer routed it to. + async fn route_once( + optimizer: &mut Box>, + ) -> Result> { + optimizer + .feed( + AgentApiOptInput::Request(request("hi", "client/model")), + enrichment(), + ) + .await?; + match optimizer.optimize().await? { + Decision::ModelInference(response) => Ok(response.requests[0].model.clone()), + Decision::Return() => Err("expected ModelInference".into()), + } + } + + #[tokio::test] + async fn all_weight_on_one_model_always_selects_it() -> Result<(), Box> { + let mut optimizer = algorithm( + vec![ + WeightedModel::new("a/model", 0.0), + WeightedModel::new("b/model", 1.0), + WeightedModel::new("c/model", 0.0), + ], + 7, + ) + .optimizer(); + + for _ in 0..25 { + assert_eq!(route_once(&mut optimizer).await?, "b/model"); + } + Ok(()) + } + + /// Over many draws the empirical selection frequencies track the weights. + #[tokio::test] + async fn selection_frequencies_track_weights() -> Result<(), Box> { + // Weights 1:3:6 -> expected shares 0.1, 0.3, 0.6. + let mut optimizer = algorithm( + vec![ + WeightedModel::new("a/model", 1.0), + WeightedModel::new("b/model", 3.0), + WeightedModel::new("c/model", 6.0), + ], + 42, + ) + .optimizer(); + + let draws = 20_000; + let mut a = 0u32; + let mut b = 0u32; + let mut c = 0u32; + for _ in 0..draws { + match route_once(&mut optimizer).await?.as_str() { + "a/model" => a += 1, + "b/model" => b += 1, + "c/model" => c += 1, + other => return Err(format!("unexpected model {other}").into()), + } + } + + let total = f64::from(draws); + assert!( + (f64::from(a) / total - 0.1).abs() < 0.02, + "a share off: {a}" + ); + assert!( + (f64::from(b) / total - 0.3).abs() < 0.02, + "b share off: {b}" + ); + assert!( + (f64::from(c) / total - 0.6).abs() < 0.02, + "c share off: {c}" + ); + Ok(()) + } + + #[tokio::test] + async fn decision_reports_draw_within_total_weight() -> Result<(), Box> { + let mut optimizer = algorithm( + vec![ + WeightedModel::new("a/model", 2.0), + WeightedModel::new("b/model", 3.0), + ], + 7, + ) + .optimizer(); + optimizer + .feed( + AgentApiOptInput::Request(request("hi", "client/model")), + enrichment(), + ) + .await?; + + match optimizer.optimize().await? { + Decision::ModelInference(response) => { + let decision = response + .decision_info + .ok_or("expected decision info on ModelInference")?; + assert_eq!(decision.total_weight, 5.0); + assert!(decision.draw >= 0.0 && decision.draw < 5.0); + } + Decision::Return() => return Err("expected ModelInference, got Return".into()), + } + Ok(()) + } + + /// Documented flow: route -> caller performs the (mocked) model call -> + /// feed the response -> optimize returns control to the agent. + #[tokio::test] + async fn returns_to_agent_after_mocked_response_is_fed() -> Result<(), Box> { + let mut optimizer = + algorithm(vec![WeightedModel::new("frontier/model", 1.0)], 7).optimizer(); + optimizer + .feed( + AgentApiOptInput::Request(request("hi", "client/model")), + enrichment(), + ) + .await?; + + let routed_model = match optimizer.optimize().await? { + Decision::ModelInference(response) => response.requests[0].model.clone(), + Decision::Return() => return Err("expected ModelInference on first optimize".into()), + }; + assert_eq!(routed_model, "frontier/model"); + + // Mock the model call by feeding a response back into the optimizer. + optimizer + .feed( + AgentApiOptInput::Response(request("mocked completion", &routed_model)), + enrichment(), + ) + .await?; + + match optimizer.optimize().await? { + Decision::Return() => Ok(()), + Decision::ModelInference(_) => Err("expected Return after response was fed".into()), + } + } + + #[tokio::test] + async fn optimize_before_feed_errors() { + let mut optimizer = algorithm(vec![WeightedModel::new("a/model", 1.0)], 7).optimizer(); + assert!(optimizer.optimize().await.is_err()); + } + + #[tokio::test] + async fn no_positive_weight_errors() -> Result<(), Box> { + let mut optimizer = algorithm( + vec![ + WeightedModel::new("a/model", 0.0), + WeightedModel::new("b/model", 0.0), + ], + 7, + ) + .optimizer(); + optimizer + .feed( + AgentApiOptInput::Request(request("hi", "client/model")), + enrichment(), + ) + .await?; + assert!(optimizer.optimize().await.is_err()); + Ok(()) + } + + /// Feed one request and return the draw the optimizer used to route it. + async fn first_draw( + optimizer: &mut Box>, + ) -> Result> { + optimizer + .feed( + AgentApiOptInput::Request(request("hi", "client/model")), + enrichment(), + ) + .await?; + match optimizer.optimize().await? { + Decision::ModelInference(response) => response + .decision_info + .map(|d| d.draw) + .ok_or_else(|| "missing decision info".into()), + Decision::Return() => Err("expected ModelInference".into()), + } + } + + /// The seed is deterministic, so two optimizers from the same factory make + /// the same first draw. + #[tokio::test] + async fn factory_mints_independent_deterministic_optimizers() -> Result<(), Box> { + let factory = algorithm( + vec![ + WeightedModel::new("a/model", 1.0), + WeightedModel::new("b/model", 1.0), + ], + 42, + ); + + let mut first = factory.optimizer(); + let mut second = factory.optimizer(); + let draw_first = first_draw(&mut first).await?; + let draw_second = first_draw(&mut second).await?; + assert_eq!(draw_first, draw_second); + Ok(()) + } +} diff --git a/docs/superpowers/specs/2026-07-04-agentapi-route-decorator-design.md b/docs/superpowers/specs/2026-07-04-agentapi-route-decorator-design.md new file mode 100644 index 00000000..bb5b6fda --- /dev/null +++ b/docs/superpowers/specs/2026-07-04-agentapi-route-decorator-design.md @@ -0,0 +1,347 @@ +# Design: `@route` decorator — pure-Python port of the libsy optimizer interfaces + +**Date:** 2026-07-04 +**Status:** Approved (design), pending implementation plan +**Branch:** `feat/libsy-routers-poc` + +## Problem + +`crates/libsy` defines a new set of Rust interfaces for LLM traffic routing — the +`AgentApiOptimizer` pattern: a per-session stateful optimizer driven by a +`feed(input) → optimize() → Decision{ModelInference | Return}` loop, minted by an +`AgentApiOptAlgorithm` factory. Two concrete algorithms exist: + +- `rand.rs` — `RandomRouterAlgorithm`: weighted random selection over N targets. +- `llm_class.rs` — `LlmClassifierAlgorithm`: a genuinely multi-round router that + first runs a classifier model call, then routes to a strong/weak model based on + the parsed score, then returns control. + +libsy is a **pure-Rust crate with no Python bindings**. We want a Python +`@route` decorator — in the spirit of litellm's unified `completion()` entry point +— that wraps a user's async "make an API call to a model" function and drives it +through one of these routing algorithms, rewriting the model (and, for multi-round +routers, the prompt) per the optimizer's decisions. + +## Goals + +- A pure-Python port of the libsy optimizer interfaces and the two algorithms, + mirroring the Rust design 1:1 in behavior. +- A `@route(algorithm, ...)` decorator that transparently drives the wrapped + async function through the optimizer loop. +- Correct support for **multi-step** routers (the LLM classifier), where the + wrapped function is invoked more than once per user call. +- Faithful port of the Rust tests to pytest. + +## Non-Goals + +- No pyo3 / Rust bindings. This is a Python-only exploration. +- No sync-function support. Async only (matches the Rust async traits 1:1). +- No integration into the production `switchyard` chain, endpoints, or CLI. +- No streaming responses. Requests/responses are single-shot. +- No new dependencies. `litellm` is referenced only in docstrings/examples; the + decorator does not import it. + +## Decisions (from brainstorming) + +| Decision | Choice | Rationale | +|---|---|---| +| Substrate | Pure-Python port | libsy has no bindings; fastest POC, mirrors Rust 1:1. | +| Function contract | Configurable adapters with litellm-shaped defaults | Works with any fn shape; common case needs zero config. | +| Sync/async | Async only | Cleanest 1:1 mapping to the Rust async traits. | +| Optimizer lifetime | One optimizer instance per decorated call | Natural decorator mapping of the Rust "one instance per session". | + +## Architecture + +A self-contained new subpackage `switchyard/lib/agentapi/`, mirroring the +`crates/libsy` module split. It does not depend on the rest of the `switchyard` +package and nothing in the package depends on it (POC isolation). + +``` +switchyard/lib/agentapi/ +├── __init__.py # public exports: route, RandomRouter, WeightedModel, +│ # LlmClassifier, ChatRequest, ChatResponse, EnrichmentData, +│ # AgentApiOptimizer, AgentApiOptAlgorithm, Decision, OptInput +├── chat.py # ChatRequest, ChatResponse, EnrichmentData +├── optimizer.py # OptInput, Decision, OptimizerResponse, abstract +│ # AgentApiOptimizer + AgentApiOptAlgorithm +├── rand.py # WeightedModel, RandomRoutingDecision, RandomRouter +├── llm_class.py # ClassifierTier, ClassifierRoutingDecision, LlmClassifier +└── decorator.py # route(...) +``` + +### Ported interface (1:1 with `lib.rs`) + +Naming note: the working tree of `lib.rs` is mid-rename +(`ChatRequest` → `AgentApitRequest`, a typo; the router modules still reference +`ChatRequest`, so the crate does not currently compile). This port uses the +coherent pre-rename names `ChatRequest` / `ChatResponse` and does **not** +propagate the typo. + +**`chat.py`** — plain dataclasses: + +```python +@dataclass +class ChatRequest: + prompt: str + model: str + +@dataclass +class ChatResponse: + completion: str + +@dataclass +class EnrichmentData: + session_id: str | None = None + agent_id: str | None = None + task_id: str | None = None + correlation_id: str | None = None + extra_metadata: dict[str, str] | None = None +``` + +**`optimizer.py`** — the input/decision unions and the abstract roles. + +`OptInput` is a tagged union. Mirroring `AgentApiOptInput::{Request, Response, +Metadata}`. Modeled as a small closed set of dataclasses under a `OptInput` base: + +```python +class OptInput: ... +@dataclass +class RequestInput(OptInput): request: ChatRequest +@dataclass +class ResponseInput(OptInput): response: ChatResponse +@dataclass +class MetadataInput(OptInput): metadata: dict[str, str] +``` + +> Faithfulness note: the Rust `AgentApiOptInput::Response` variant currently wraps +> a `ChatRequest` (the tests stuff the completion text into its `prompt` field). +> The Python port cleans this up by using `ChatResponse.completion`, which is the +> evident intent. This is the one intentional deviation from the Rust source, and +> it is documented in the module docstring. + +`OptimizerResponse` and `Decision`: + +```python +@dataclass +class OptimizerResponse: + requests: list[ChatRequest] + enrichment_data: list[EnrichmentData] = field(default_factory=list) + decision_reasoning: str | None = None + decision_info: object | None = None # concrete `D` per algorithm; no generics + +class Decision: ... +@dataclass +class ModelInference(Decision): response: OptimizerResponse +@dataclass +class Return(Decision): pass +``` + +`AgentApiOptimizer` (abstract base) and `AgentApiOptAlgorithm` (factory): + +```python +class AgentApiOptimizer(abc.ABC): + async def feed(self, input: OptInput, enrichment: EnrichmentData) -> None: ... + @abc.abstractmethod + async def optimize(self) -> Decision: ... + +class AgentApiOptAlgorithm(abc.ABC): + @abc.abstractmethod + def optimizer(self) -> AgentApiOptimizer: ... +``` + +`feed` has a no-op default (matching the Rust default trait method). `optimize` is +made **abstract** in Python — a deliberate tightening of the Rust trait, which +ships a default `optimize` returning an empty `ModelInference`. Both concrete +algorithms override it, so requiring the override surfaces an incomplete +optimizer at construction time rather than silently no-op'ing. The `D` generic is +dropped — Python attaches `decision_info` as a plain object. + +### Algorithm: `rand.py` — `RandomRouter` + +Ports `RandomRouterAlgorithm` + `RandomRouter`. + +- `WeightedModel(model: str, weight: float)`. +- `RandomRouter(models: list[WeightedModel], rng_seed: int | None = None)` — the + factory (`AgentApiOptAlgorithm`). `optimizer()` mints a fresh per-session + optimizer. Uses `random.Random(seed)` when a seed is given, else + `random.Random()`. +- The per-session optimizer holds `models`, `rng`, `pending_request`, `completed`. + - `feed(RequestInput)` → buffer the request. + - `feed(ResponseInput)` → mark completed. + - `optimize()` when completed → `Return`. + - `optimize()` with a pending request → weighted draw: sum finite positive + weights; error if the sum ≤ 0 ("random router has no target with positive + weight"); draw in `[0, total)`, walk cumulative weights, fall back to the last + positively-weighted target on floating-point overshoot. Rewrite + `request.model`, return `ModelInference` with a `RandomRoutingDecision` + (`selected_model`, `draw`, `total_weight`) as `decision_info`. + - `optimize()` before any request was fed → error ("optimize called before a + request was fed"). + +### Algorithm: `llm_class.py` — `LlmClassifier` + +Ports `LlmClassifierAlgorithm` + `LlmClassifierRouter`, including the phase +machine. + +- `ClassifierTier` enum (`STRONG` / `WEAK`) with a stable `as_str()`. +- `ClassifierRoutingDecision(score: float | None, threshold, tier, selected_model)`. +- `LlmClassifier(classifier_model, strong_model, weak_model, threshold)` — the + factory. `optimizer()` mints a per-session optimizer starting in + `AwaitingRequest`. +- Phases: `AwaitingRequest → Classify → AwaitingScore → Route → AwaitingResponse + → Done` (mirrors the Rust `Phase` enum exactly). +- `CLASSIFIER_PROMPT_PREAMBLE` matches the Rust literal's **resolved** value (the + Rust source uses a `\`-newline continuation with leading indentation that + collapses to single spaces), i.e. the exact string: + `"Rate how strongly this request needs a frontier model. Reply with a single strong-win-rate score in [0, 1]:\n"`. +- Behavior mirrors `llm_class.rs`: + - `feed(RequestInput)` → buffer request, phase `Classify`. + - `feed(ResponseInput)` in `AwaitingScore` → parse score + (`float(text.strip())`, `None` on failure), phase `Route`. In + `AwaitingResponse` → phase `Done`. Any other phase → error ("classifier router + received a response outside a pending model call"). + - `optimize()` in `Classify` → emit `ChatRequest(classifier_model, preamble + + user_prompt)`, phase `AwaitingScore`. + - `optimize()` in `Route` → decide tier: `score >= threshold` → strong; + `score < threshold` → weak; `score is None` → **default strong** (keep traffic + flowing when the classifier output was unusable). Rewrite model, phase + `AwaitingResponse`, return `ModelInference` with `ClassifierRoutingDecision`. + - `optimize()` in `Done` → `Return`. + - `optimize()` in `AwaitingRequest` / `AwaitingScore` / `AwaitingResponse` → + the corresponding "called before … was fed" errors. + +### The decorator: `decorator.py` + +```python +def route( + algorithm: AgentApiOptAlgorithm, + *, + get_prompt: Callable[[Mapping], str] = _default_get_prompt, + apply_request: Callable[[dict, ChatRequest], None] = _default_apply_request, + get_completion: Callable[[object], str] = _default_get_completion, +) -> Callable: + ... +``` + +`route` returns a decorator that wraps an **async** function and returns an async +wrapper. The wrapper, per call: + +1. Binds the call's args to a mutable `kwargs` view (via `inspect.signature`, + normalizing positional args to keyword form so adapters see a uniform mapping). +2. Mints a fresh optimizer: `optimizer = algorithm.optimizer()`. +3. Extracts the initial prompt (`get_prompt(kwargs)`) and model (`kwargs["model"]`) + and feeds `RequestInput(ChatRequest(prompt, model))` with empty + `EnrichmentData`. +4. **Round-agnostic drive loop:** + + ```python + last_resp = None + while True: + decision = await optimizer.optimize() + if isinstance(decision, Return): + break + for req in decision.response.requests: + apply_request(kwargs, req) # sets model AND prompt for this round + last_resp = await fn(**kwargs) + completion = get_completion(last_resp) + await optimizer.feed(ResponseInput(ChatResponse(completion)), + EnrichmentData()) + return last_resp + ``` + +5. Returns `last_resp` — the wrapped fn's return value from the **last** model + call before `Return`. For the classifier this is the routed call's response; + the earlier classifier call's response is consumed internally as the score. + +**Why this supports multi-step routers:** the loop makes no assumption about how +many `ModelInference` rounds occur. Each `ModelInference` carries full +`ChatRequest`s (model + prompt); `apply_request` applies the *entire* request to +the call each round, so the classifier round correctly sends the classifier +prompt to the classifier model, and the routed round sends the original prompt to +the routed model. A 1-round router (`rand`) and an N-round router (`llm_class`) +drive through identical code. + +**Adapters — litellm-shaped defaults (zero config for the common case):** + +- `_default_get_prompt(kwargs)` → `kwargs["messages"][-1]["content"]`; falls back + to `kwargs["prompt"]` when there is no `messages` key. +- `_default_apply_request(kwargs, req)` → sets `kwargs["model"] = req.model`; if + `messages` is present, replaces the **last user message's content** with + `req.prompt` (preserving system messages / earlier history); otherwise sets + `kwargs["prompt"] = req.prompt`. +- `_default_get_completion(resp)` → `resp.choices[0].message.content`; falls back + to `resp` itself when it is a `str`. + +All three are overridable to support arbitrary function shapes. + +`functools.wraps` preserves the wrapped function's identity. If the wrapped +function is not a coroutine function, `route` raises `TypeError` at decoration +time (async-only contract). + +## Data flow + +``` +caller: await chat(model="auto", messages=[{user: "..."}]) + → wrapper binds args → kwargs + → feed(RequestInput(ChatRequest(prompt, "auto"))) + → loop: + optimize() → ModelInference(reqs) # rand: 1 round; classifier: 2 rounds + → apply_request(kwargs, req); resp = await chat(**kwargs) + → feed(ResponseInput(ChatResponse(completion(resp)))) + optimize() → Return + → return last resp +``` + +## Error handling + +- Optimizer state errors (optimize-before-feed, no-positive-weight, response out + of phase) propagate as `ValueError` (Python analog of the Rust + `Box`), with messages matching the Rust strings. +- Decorating a non-coroutine function raises `TypeError` at decoration time. +- Adapter failures (e.g. missing `model` kwarg, unparseable response shape) + surface as the natural `KeyError` / `AttributeError`; the POC does not wrap + these. + +## Testing + +Ported to pytest (`asyncio_mode = "auto"`, no explicit markers), against a mock +async wrapped function, under `tests/`: + +**`tests/test_agentapi_rand.py`** (ports `rand.rs` tests): +- all weight on one model always selects it +- selection frequencies track weights over 20k draws (with fixed seed) +- decision reports draw within total weight +- returns to agent after a mocked response is fed +- optimize-before-feed errors +- no-positive-weight errors +- factory mints independent deterministic optimizers (same seed → same first draw) + +**`tests/test_agentapi_llm_class.py`** (ports `llm_class.rs` tests): +- score ≥ threshold routes strong; score < threshold routes weak +- score exactly at threshold routes strong +- unparseable score defaults to strong (score `None`) +- optimize-before-feed errors +- optimize-before-classifier-response errors + +**`tests/test_agentapi_decorator.py`** (new, decorator end-to-end): +- rand: wrapped fn called once; returned response is the routed call's; model was + rewritten to the weighted target. +- classifier: wrapped fn called twice (classifier + routed); the classifier round + receives the classifier model and preamble prompt; the returned response is the + routed round's; the routed round receives the original prompt. +- custom adapters: a non-litellm function shape drives correctly via overridden + `get_prompt` / `apply_request` / `get_completion`. +- decorating a sync function raises `TypeError`. + +## Validation + +- `uv run ruff check .` — zero errors. +- `uv run mypy switchyard` — clean (strict). +- `uv run pytest tests/test_agentapi_*.py -v` — green. + +## Public API + +Exported from `switchyard/lib/agentapi/__init__.py`. Whether these also surface +from the top-level `switchyard/__init__.py` `__all__` is deferred to the +implementation plan (POC isolation argues for keeping them subpackage-local +initially). diff --git a/examples/react_agent.py b/examples/react_agent.py new file mode 100644 index 00000000..2a5b19f2 --- /dev/null +++ b/examples/react_agent.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimum-viable ReAct agent, with traffic routed by the agentapi `@route` decorator. + +A small Thought/Action/Observation loop over a single tool (a calculator). Every +model call goes through one function, `call_model`, which is wrapped with +`@route(...)` — so routing across models is a one-line drop-in. The decorator +rewrites the target model per the chosen algorithm before each call, and +`call_model` prints the model it was actually routed to. + +Configuration (env): + OPENAI_API_KEY required — key for the OpenAI-compatible endpoint + OPENAI_BASE_URL optional — endpoint base url (default OpenAI) + SWITCHYARD_WEAK_MODEL optional — cheap/default model (default gpt-4o-mini) + SWITCHYARD_STRONG_MODEL optional — frontier model (default gpt-4o) + SWITCHYARD_DEMO_ROUTER optional — "rand" (default) or "llm_class" + +Usage: + export OPENAI_API_KEY="sk-..." + python examples/react_agent.py "What is 17 * 24, then add 100?" +""" + +from __future__ import annotations + +import ast +import asyncio +import operator +import os +import re +import sys +from pathlib import Path +from typing import Any + +# Add package to path for development (not needed when installed via pip). +sys.path.insert(0, str(Path(__file__).parent.parent)) + +import litellm + +from switchyard.lib.agentapi import ( + AgentApiOptAlgorithm, + LlmClassifier, + RandomRouter, + WeightedModel, + route, +) + +API_KEY = os.environ.get("OPENAI_API_KEY") +BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") +WEAK_MODEL = os.environ.get("SWITCHYARD_WEAK_MODEL", "gpt-4o-mini") +STRONG_MODEL = os.environ.get("SWITCHYARD_STRONG_MODEL", "gpt-4o") + +SYSTEM_PROMPT = """\ +You are a reasoning agent. Solve the task using EXACTLY this format: + +Thought: +Action: calculator +Action Input: + +You will then receive an "Observation:" line with the tool result. Repeat +Thought/Action/Action Input as many times as needed. When you have the answer, +reply with: + +Thought: +Final Answer: + +The only tool is `calculator`, which evaluates one arithmetic expression.""" + + +def build_router() -> AgentApiOptAlgorithm: + """Pick the routing algorithm from env; defaults to a weighted random split.""" + kind = os.environ.get("SWITCHYARD_DEMO_ROUTER", "rand") + if kind == "llm_class": + # Multi-round: an LLM classifier scores each request, then routes to the + # strong or weak model. Each ReAct step therefore makes two model calls. + return LlmClassifier( + classifier_model=WEAK_MODEL, + strong_model=STRONG_MODEL, + weak_model=WEAK_MODEL, + threshold=0.5, + ) + # Single-round: send ~75% of traffic to the weak model, ~25% to the strong one. + return RandomRouter([WeightedModel(WEAK_MODEL, 3.0), WeightedModel(STRONG_MODEL, 1.0)]) + + +ROUTER = build_router() + + +@route(ROUTER) +async def call_model(model: str, messages: list[dict]) -> Any: + """Make one chat-completion call. `@route` rewrites `model` before we run.""" + print(f" -> routed to {model}") + return await litellm.acompletion( + model=model, + messages=messages, + api_key=API_KEY, + api_base=BASE_URL, + temperature=0.0, + ) + + +# Safe arithmetic evaluator: walk a parsed AST rather than calling eval(). +_OPS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Mod: operator.mod, + ast.Pow: operator.pow, + ast.USub: operator.neg, +} + + +def _eval(node: ast.AST) -> float: + if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)): + return node.value + if isinstance(node, ast.BinOp) and type(node.op) in _OPS: + return _OPS[type(node.op)](_eval(node.left), _eval(node.right)) + if isinstance(node, ast.UnaryOp) and type(node.op) in _OPS: + return _OPS[type(node.op)](_eval(node.operand)) + raise ValueError("unsupported expression") + + +def calculator(expr: str) -> str: + """Evaluate a single arithmetic expression, e.g. "17 * 24".""" + return str(_eval(ast.parse(expr, mode="eval").body)) + + +def run_tool(name: str, arg: str) -> str: + """Dispatch a tool call, returning an Observation string for the model.""" + if name == "calculator": + try: + return calculator(arg) + except (ValueError, KeyError, SyntaxError, ZeroDivisionError, TypeError) as exc: + return f"calculator error: {exc}" + return f"unknown tool: {name}" + + +_FINAL_RE = re.compile(r"Final Answer:\s*(.*)", re.DOTALL) +_ACTION_RE = re.compile(r"Action:\s*(.+)") +_INPUT_RE = re.compile(r"Action Input:\s*(.+)") + + +async def react(question: str, max_steps: int = 6) -> str: + """Run the ReAct loop until a Final Answer or the step budget is exhausted.""" + messages: list[dict] = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": question}, + ] + for step in range(1, max_steps + 1): + print(f"[step {step}]") + response = await call_model(model="auto", messages=messages) + text = response.choices[0].message.content + print("\n".join(f" {line}" for line in text.strip().splitlines())) + + final = _FINAL_RE.search(text) + if final: + return final.group(1).strip() + + messages.append({"role": "assistant", "content": text}) + action = _ACTION_RE.search(text) + action_input = _INPUT_RE.search(text) + if action and action_input: + observation = run_tool(action.group(1).strip(), action_input.group(1).strip()) + else: + observation = "no valid Action found; follow the Thought/Action/Action Input format" + print(f" Observation: {observation}") + messages.append({"role": "user", "content": f"Observation: {observation}"}) + + return "(max steps reached without a final answer)" + + +async def main() -> None: + if not API_KEY: + print("Set OPENAI_API_KEY (and optionally OPENAI_BASE_URL) to run this example.") + return + question = " ".join(sys.argv[1:]) or "What is 17 * 24, then add 100?" + print(f"Router: {type(ROUTER).__name__}") + print(f"Question: {question}\n") + answer = await react(question) + print(f"\nFinal Answer: {answer}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/switchyard/lib/agentapi/__init__.py b/switchyard/lib/agentapi/__init__.py new file mode 100644 index 00000000..612f2ff9 --- /dev/null +++ b/switchyard/lib/agentapi/__init__.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure-Python port of the libsy AgentApi optimizer interfaces and routers. + +Public entry point is the `route` decorator, which drives an async model-call +function through a routing algorithm (`RandomRouter` or `LlmClassifier`). +""" + +from switchyard.lib.agentapi.chat import ChatRequest, ChatResponse, EnrichmentData +from switchyard.lib.agentapi.decorator import route +from switchyard.lib.agentapi.llm_class import ( + ClassifierRoutingDecision, + ClassifierTier, + LlmClassifier, +) +from switchyard.lib.agentapi.optimizer import ( + AgentApiOptAlgorithm, + AgentApiOptimizer, + Decision, + MetadataInput, + ModelInference, + OptimizerResponse, + RequestInput, + ResponseInput, + Return, +) +from switchyard.lib.agentapi.rand import RandomRouter, RandomRoutingDecision, WeightedModel + +__all__ = [ + "route", + "RandomRouter", + "RandomRoutingDecision", + "WeightedModel", + "LlmClassifier", + "ClassifierRoutingDecision", + "ClassifierTier", + "ChatRequest", + "ChatResponse", + "EnrichmentData", + "AgentApiOptimizer", + "AgentApiOptAlgorithm", + "Decision", + "ModelInference", + "Return", + "OptimizerResponse", + "RequestInput", + "ResponseInput", + "MetadataInput", +] diff --git a/switchyard/lib/agentapi/chat.py b/switchyard/lib/agentapi/chat.py new file mode 100644 index 00000000..ed8876cc --- /dev/null +++ b/switchyard/lib/agentapi/chat.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal request/response/enrichment types shared by the optimizers. + +Ports libsy's ChatRequest / ChatResponse / EnrichementData. These are the +provider-neutral representations the optimizers reason about. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class ChatRequest: + """Provider-neutral request: a single prompt aimed at a model id.""" + + prompt: str + model: str + + +@dataclass +class ChatResponse: + """Provider-neutral response carrying the completion text.""" + + completion: str + + +@dataclass +class EnrichmentData: + """Correlation/enrichment metadata used for routing decisions.""" + + session_id: str | None = None + agent_id: str | None = None + task_id: str | None = None + correlation_id: str | None = None + extra_metadata: dict[str, str] | None = None diff --git a/switchyard/lib/agentapi/decorator.py b/switchyard/lib/agentapi/decorator.py new file mode 100644 index 00000000..87386b12 --- /dev/null +++ b/switchyard/lib/agentapi/decorator.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The @route decorator: drives an async model-call function through a router. + +Per call it mints one optimizer and runs the round-agnostic optimizer loop, +applying each emitted ChatRequest (model + prompt) to the wrapped function and +feeding the response back until the optimizer returns control. A single-round +router (rand) and a multi-round router (llm_class) drive through identical code. +""" + +from __future__ import annotations + +import functools +import inspect +from collections.abc import Callable, Mapping +from typing import Any + +from switchyard.lib.agentapi.chat import ChatRequest, ChatResponse, EnrichmentData +from switchyard.lib.agentapi.optimizer import ( + AgentApiOptAlgorithm, + ModelInference, + RequestInput, + ResponseInput, + Return, +) + + +def _default_get_prompt(kwargs: Mapping[str, Any]) -> str: + """litellm-shaped: last user message content, else a `prompt` kwarg.""" + messages = kwargs.get("messages") + if messages is not None: + return str(messages[-1]["content"]) + return str(kwargs["prompt"]) + + +def _default_apply_request(kwargs: dict[str, Any], req: ChatRequest) -> None: + """Apply an emitted request: set model, and set the prompt for this round. + + With `messages`, replaces the last user message's content (preserving system + messages / history); otherwise sets a `prompt` kwarg. + """ + kwargs["model"] = req.model + messages = kwargs.get("messages") + if messages is not None: + new_messages = list(messages) + for i in range(len(new_messages) - 1, -1, -1): + if new_messages[i].get("role") == "user": + new_messages[i] = {**new_messages[i], "content": req.prompt} + break + else: + new_messages.append({"role": "user", "content": req.prompt}) + kwargs["messages"] = new_messages + else: + kwargs["prompt"] = req.prompt + + +def _default_get_completion(resp: Any) -> str: + """litellm-shaped: resp.choices[0].message.content, else a str response.""" + if isinstance(resp, str): + return resp + return str(resp.choices[0].message.content) + + +def route( + algorithm: AgentApiOptAlgorithm, + *, + get_prompt: Callable[[Mapping[str, Any]], str] = _default_get_prompt, + apply_request: Callable[[dict[str, Any], ChatRequest], None] = _default_apply_request, + get_completion: Callable[[Any], str] = _default_get_completion, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Wrap an async model-call function so `algorithm` routes each call. + + The wrapped function must be a coroutine function with a keyword-compatible + signature. Returns the wrapped function's return value from the last model + call the optimizer requested (for a classifier, the routed call's response). + """ + + def decorate(fn: Callable[..., Any]) -> Callable[..., Any]: + if not inspect.iscoroutinefunction(fn): + raise TypeError(f"@route requires an async function, got {fn!r}") + sig = inspect.signature(fn) + + @functools.wraps(fn) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + call_kwargs = bound.arguments + + optimizer = algorithm.optimizer() + prompt = get_prompt(call_kwargs) + # Lenient: both routers overwrite the inbound model, and custom + # adapters may use a non-`model` parameter name (e.g. `engine`). + model = call_kwargs.get("model", "") + await optimizer.feed( + RequestInput(ChatRequest(prompt=prompt, model=model)), EnrichmentData() + ) + + last_resp: Any = None + while True: + decision = await optimizer.optimize() + if isinstance(decision, Return): + break + assert isinstance(decision, ModelInference) + for req in decision.response.requests: + apply_request(call_kwargs, req) + last_resp = await fn(**call_kwargs) + completion = get_completion(last_resp) + await optimizer.feed( + ResponseInput(ChatResponse(completion)), EnrichmentData() + ) + return last_resp + + return wrapper + + return decorate diff --git a/switchyard/lib/agentapi/llm_class.py b/switchyard/lib/agentapi/llm_class.py new file mode 100644 index 00000000..6623e173 --- /dev/null +++ b/switchyard/lib/agentapi/llm_class.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LLM-classifier router (ports llm_class.rs). + +Unlike a local ML classifier, an LLM classifier needs its own model call to +score the request. That maps onto the optimizer's multi-round loop: the router +first asks the caller to run a classifier model, then — once the score is fed +back — asks the caller to run the routed target model, and finally returns +control to the agent. +""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass + +from switchyard.lib.agentapi.chat import ChatRequest, EnrichmentData +from switchyard.lib.agentapi.optimizer import ( + AgentApiOptAlgorithm, + AgentApiOptimizer, + Decision, + ModelInference, + OptimizerResponse, + OptInput, + RequestInput, + ResponseInput, + Return, +) + +# Prepended to the user prompt when asking the classifier for a strong-win-rate +# score. Matches the resolved value of the Rust literal. +CLASSIFIER_PROMPT_PREAMBLE = ( + "Rate how strongly this request needs a frontier model. " + "Reply with a single strong-win-rate score in [0, 1]:\n" +) + + +class ClassifierTier(enum.Enum): + """The tier a classifier score selected.""" + + STRONG = "strong" + WEAK = "weak" + + def as_str(self) -> str: + return self.value + + +@dataclass +class ClassifierRoutingDecision: + """Decision info attached to the routed ModelInference decision.""" + + score: float | None + threshold: float + tier: ClassifierTier + selected_model: str + + +class _Phase(enum.Enum): + """Where the router is in its classify -> route -> return lifecycle.""" + + AWAITING_REQUEST = enum.auto() + CLASSIFY = enum.auto() + AWAITING_SCORE = enum.auto() + ROUTE = enum.auto() + AWAITING_RESPONSE = enum.auto() + DONE = enum.auto() + + +class LlmClassifier(AgentApiOptAlgorithm): + """Factory minting a fresh LLM-classifier optimizer per session.""" + + def __init__( + self, + classifier_model: str, + strong_model: str, + weak_model: str, + threshold: float, + ) -> None: + self.classifier_model = classifier_model + self.strong_model = strong_model + self.weak_model = weak_model + self.threshold = threshold + + def optimizer(self) -> AgentApiOptimizer: + return _ClassifierOptimizer( + self.classifier_model, self.strong_model, self.weak_model, self.threshold + ) + + +class _ClassifierOptimizer(AgentApiOptimizer): + """Per-session LLM-classifier router.""" + + def __init__( + self, classifier_model: str, strong_model: str, weak_model: str, threshold: float + ) -> None: + self._classifier_model = classifier_model + self._strong_model = strong_model + self._weak_model = weak_model + self._threshold = threshold + self._phase = _Phase.AWAITING_REQUEST + self._pending_request: ChatRequest | None = None + self._score: float | None = None + + async def feed(self, input: OptInput, enrichment: EnrichmentData) -> None: + if isinstance(input, RequestInput): + self._pending_request = input.request + self._phase = _Phase.CLASSIFY + elif isinstance(input, ResponseInput): + if self._phase is _Phase.AWAITING_SCORE: + self._score = _parse_score(input.response.completion) + self._phase = _Phase.ROUTE + elif self._phase is _Phase.AWAITING_RESPONSE: + self._phase = _Phase.DONE + else: + raise ValueError( + "classifier router received a response outside a pending model call" + ) + + async def optimize(self) -> Decision: + if self._phase is _Phase.AWAITING_REQUEST: + raise ValueError("optimize called before a request was fed") + if self._phase is _Phase.CLASSIFY: + assert self._pending_request is not None + user_prompt = self._pending_request.prompt + self._phase = _Phase.AWAITING_SCORE + classifier_request = ChatRequest( + prompt=f"{CLASSIFIER_PROMPT_PREAMBLE}{user_prompt}", + model=self._classifier_model, + ) + return ModelInference( + OptimizerResponse( + requests=[classifier_request], + decision_reasoning=f"classifying request via {self._classifier_model}", + ) + ) + if self._phase is _Phase.AWAITING_SCORE: + raise ValueError("optimize called before the classifier response was fed") + if self._phase is _Phase.ROUTE: + assert self._pending_request is not None + request = self._pending_request + self._pending_request = None + decision = self._decide() + request.model = decision.selected_model + self._phase = _Phase.AWAITING_RESPONSE + return ModelInference( + OptimizerResponse( + requests=[request], + decision_reasoning=( + f"classifier score {decision.score} vs threshold " + f"{decision.threshold}; selected {decision.selected_model} " + f"({decision.tier.as_str()})" + ), + decision_info=decision, + ) + ) + if self._phase is _Phase.AWAITING_RESPONSE: + raise ValueError("optimize called before the model response was fed") + # _Phase.DONE + return Return() + + def _decide(self) -> ClassifierRoutingDecision: + """Build the routing decision, defaulting to strong on an unusable score.""" + if self._score is not None and self._score >= self._threshold: + tier, model = ClassifierTier.STRONG, self._strong_model + elif self._score is not None: + tier, model = ClassifierTier.WEAK, self._weak_model + else: + # Defensive default: keep traffic flowing on the strong tier when the + # classifier response was unparseable. + tier, model = ClassifierTier.STRONG, self._strong_model + return ClassifierRoutingDecision(self._score, self._threshold, tier, model) + + +def _parse_score(text: str) -> float | None: + """Parse a strong-win-rate score, returning None when unparseable.""" + try: + return float(text.strip()) + except ValueError: + return None diff --git a/switchyard/lib/agentapi/optimizer.py b/switchyard/lib/agentapi/optimizer.py new file mode 100644 index 00000000..90700b19 --- /dev/null +++ b/switchyard/lib/agentapi/optimizer.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Optimizer role interfaces ported from libsy's lib.rs. + +An AgentApiOptimizer is a per-session state machine fed inputs (requests, +responses, metadata) that yields routing/optimization Decisions. An +AgentApiOptAlgorithm is the factory that mints a fresh optimizer per session. + +Deviation from the Rust source: the Response input variant here wraps a +ChatResponse (its completion text) rather than reusing ChatRequest, which is the +evident intent of the Rust `AgentApiOptInput::Response`. +""" + +from __future__ import annotations + +import abc +from dataclasses import dataclass, field + +from switchyard.lib.agentapi.chat import ChatRequest, ChatResponse, EnrichmentData + + +class OptInput: + """Base for optimizer inputs (request / response / metadata).""" + + +@dataclass +class RequestInput(OptInput): + """An inbound request to route.""" + + request: ChatRequest + + +@dataclass +class ResponseInput(OptInput): + """A model response fed back after the caller performed a model call.""" + + response: ChatResponse + + +@dataclass +class MetadataInput(OptInput): + """Out-of-band metadata for routing.""" + + metadata: dict[str, str] + + +@dataclass +class OptimizerResponse: + """Payload of a ModelInference decision: the model calls to perform.""" + + requests: list[ChatRequest] + enrichment_data: list[EnrichmentData] = field(default_factory=list) + decision_reasoning: str | None = None + decision_info: object | None = None + + +class Decision: + """Base for optimizer decisions (ModelInference / Return).""" + + +@dataclass +class ModelInference(Decision): + """The caller should perform the listed model calls and feed responses back.""" + + response: OptimizerResponse + + +@dataclass +class Return(Decision): + """The optimizer is done; hand control back to the calling agent.""" + + +class AgentApiOptimizer(abc.ABC): + """Per-session stateful optimizer driven by feed() -> optimize().""" + + async def feed(self, input: OptInput, enrichment: EnrichmentData) -> None: # noqa: B027 + """Feed a new input; default is a no-op (override to accumulate state).""" + + @abc.abstractmethod + async def optimize(self) -> Decision: + """Make the next routing/optimization decision from current state.""" + + +class AgentApiOptAlgorithm(abc.ABC): + """Factory that mints a fresh optimizer per session.""" + + @abc.abstractmethod + def optimizer(self) -> AgentApiOptimizer: + """Create a new optimizer instance for a session.""" diff --git a/switchyard/lib/agentapi/rand.py b/switchyard/lib/agentapi/rand.py new file mode 100644 index 00000000..23f84024 --- /dev/null +++ b/switchyard/lib/agentapi/rand.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Weighted random router (ports rand.rs) built on the optimizer interfaces.""" + +from __future__ import annotations + +import math +import random +from dataclasses import dataclass + +from switchyard.lib.agentapi.chat import ChatRequest, EnrichmentData +from switchyard.lib.agentapi.optimizer import ( + AgentApiOptAlgorithm, + AgentApiOptimizer, + Decision, + ModelInference, + OptimizerResponse, + OptInput, + RequestInput, + ResponseInput, + Return, +) + + +@dataclass +class WeightedModel: + """A routing target and its relative (non-negative, finite) selection weight.""" + + model: str + weight: float + + +@dataclass +class RandomRoutingDecision: + """Decision info attached to a random-routing ModelInference.""" + + selected_model: str + draw: float + total_weight: float + + +class RandomRouter(AgentApiOptAlgorithm): + """Factory minting a fresh weighted-random optimizer per session. + + Weights are relative; a target is chosen with probability + weight / sum(weights). A seed makes routing reproducible. + """ + + def __init__(self, models: list[WeightedModel], rng_seed: int | None = None) -> None: + self.models = models + self.rng_seed = rng_seed + + def optimizer(self) -> AgentApiOptimizer: + rng = random.Random(self.rng_seed) if self.rng_seed is not None else random.Random() + return _RandomOptimizer(list(self.models), rng) + + +class _RandomOptimizer(AgentApiOptimizer): + """Per-session random router over a weighted set of N targets.""" + + def __init__(self, models: list[WeightedModel], rng: random.Random) -> None: + self._models = models + self._rng = rng + self._pending_request: ChatRequest | None = None + self._completed = False + + async def feed(self, input: OptInput, enrichment: EnrichmentData) -> None: + if isinstance(input, RequestInput): + self._pending_request = input.request + elif isinstance(input, ResponseInput): + # A fed response means the caller performed the model call; the next + # optimize should hand control back to the agent. + self._completed = True + + async def optimize(self) -> Decision: + if self._completed: + return Return() + if self._pending_request is None: + raise ValueError("optimize called before a request was fed") + request = self._pending_request + self._pending_request = None + decision = self._select() + request.model = decision.selected_model + return ModelInference( + OptimizerResponse( + requests=[request], + decision_reasoning=( + f"weighted random draw {decision.draw} of total weight " + f"{decision.total_weight}; selected {decision.selected_model}" + ), + decision_info=decision, + ) + ) + + def _select(self) -> RandomRoutingDecision: + """Draw a weighted target; error if no target has positive finite weight.""" + selectable = [m for m in self._models if math.isfinite(m.weight) and m.weight > 0.0] + total_weight = sum(m.weight for m in selectable) + if total_weight <= 0.0: + raise ValueError("random router has no target with positive weight") + draw = self._rng.random() * total_weight + cumulative = 0.0 + for model in selectable: + cumulative += model.weight + if draw < cumulative: + return RandomRoutingDecision(model.model, draw, total_weight) + # Floating-point rounding can leave the draw at the top of the range; + # fall back to the last positively-weighted target. + return RandomRoutingDecision(selectable[-1].model, draw, total_weight) diff --git a/tests/test_agentapi_chat.py b/tests/test_agentapi_chat.py new file mode 100644 index 00000000..6ca0b014 --- /dev/null +++ b/tests/test_agentapi_chat.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the agentapi chat data types.""" + +from __future__ import annotations + +from switchyard.lib.agentapi.chat import ChatRequest, ChatResponse, EnrichmentData + + +def test_chat_request_holds_prompt_and_model(): + req = ChatRequest(prompt="hi", model="client/model") + assert req.prompt == "hi" + assert req.model == "client/model" + + +def test_chat_response_holds_completion(): + assert ChatResponse(completion="done").completion == "done" + + +def test_enrichment_data_defaults_to_none(): + e = EnrichmentData() + assert e.session_id is None + assert e.extra_metadata is None diff --git a/tests/test_agentapi_decorator.py b/tests/test_agentapi_decorator.py new file mode 100644 index 00000000..affeff5b --- /dev/null +++ b/tests/test_agentapi_decorator.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end tests for the @route decorator.""" + +from __future__ import annotations + +import pytest + +from switchyard.lib.agentapi.decorator import route +from switchyard.lib.agentapi.llm_class import LlmClassifier +from switchyard.lib.agentapi.rand import RandomRouter, WeightedModel + + +class _FakeMessage: + def __init__(self, content: str) -> None: + self.message = type("M", (), {"content": content})() + + +class _FakeResponse: + """Mimics a litellm/OpenAI response: resp.choices[0].message.content.""" + + def __init__(self, content: str) -> None: + self.choices = [_FakeMessage(content)] + + +async def test_rand_calls_wrapped_fn_once_and_rewrites_model(): + calls: list[dict] = [] + + @route(RandomRouter([WeightedModel("frontier/model", 1.0)], rng_seed=1)) + async def chat(model: str, messages: list[dict]) -> _FakeResponse: + calls.append({"model": model, "content": messages[-1]["content"]}) + return _FakeResponse(f"answer from {model}") + + resp = await chat(model="auto", messages=[{"role": "user", "content": "hello"}]) + + assert len(calls) == 1 + assert calls[0]["model"] == "frontier/model" + assert calls[0]["content"] == "hello" + assert resp.choices[0].message.content == "answer from frontier/model" + + +async def test_classifier_calls_wrapped_fn_twice_and_returns_routed_response(): + calls: list[dict] = [] + + @route( + LlmClassifier( + classifier_model="router/clf", + strong_model="big/model", + weak_model="small/model", + threshold=0.5, + ) + ) + async def chat(model: str, messages: list[dict]) -> _FakeResponse: + content = messages[-1]["content"] + calls.append({"model": model, "content": content}) + # The classifier round must return a parseable score; the routed round + # returns the real answer. + reply = "0.9" if model == "router/clf" else f"answer from {model}" + return _FakeResponse(reply) + + resp = await chat(model="auto", messages=[{"role": "user", "content": "prove it"}]) + + assert len(calls) == 2 + # Round 1: classifier model, classifier preamble prompt containing user text. + assert calls[0]["model"] == "router/clf" + assert "prove it" in calls[0]["content"] + assert "frontier model" in calls[0]["content"] + # Round 2: routed (strong) model, original user prompt. + assert calls[1]["model"] == "big/model" + assert calls[1]["content"] == "prove it" + # Returned response is the routed call's, not the classifier's. + assert resp.choices[0].message.content == "answer from big/model" + + +async def test_custom_adapters_drive_non_litellm_shape(): + @route( + RandomRouter([WeightedModel("target/model", 1.0)], rng_seed=1), + get_prompt=lambda kw: kw["text"], + apply_request=lambda kw, req: kw.update(engine=req.model, text=req.prompt), + get_completion=lambda resp: resp, + ) + async def chat(engine: str, text: str) -> str: + return f"{engine}:{text}" + + result = await chat(engine="auto", text="hi") + assert result == "target/model:hi" + + +async def test_default_adapters_drive_prompt_shaped_fn(): + # Exercises the non-`messages` fallbacks: get_prompt reads `prompt`, + # apply_request writes `prompt`, get_completion passes a str through. + calls: list[dict] = [] + + @route(RandomRouter([WeightedModel("target/model", 1.0)], rng_seed=1)) + async def chat(model: str, prompt: str) -> str: + calls.append({"model": model, "prompt": prompt}) + return f"reply to {prompt}" + + result = await chat(model="auto", prompt="hello") + + assert calls == [{"model": "target/model", "prompt": "hello"}] + assert result == "reply to hello" + + +async def test_default_apply_request_preserves_system_message(): + # A system message must survive both rounds; only the last user message's + # content is rewritten (classifier preamble, then the restored prompt). + calls: list[list[dict]] = [] + + @route( + LlmClassifier( + classifier_model="router/clf", + strong_model="big/model", + weak_model="small/model", + threshold=0.5, + ) + ) + async def chat(model: str, messages: list[dict]) -> _FakeResponse: + calls.append(messages) + return _FakeResponse("0.9" if model == "router/clf" else "answer") + + await chat( + model="auto", + messages=[ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "prove it"}, + ], + ) + + assert len(calls) == 2 + for sent in calls: + assert sent[0] == {"role": "system", "content": "be terse"} + # Classifier round rewrote the last user message to the preamble prompt. + assert "prove it" in calls[0][1]["content"] + assert "frontier model" in calls[0][1]["content"] + # Routed round restored the original user prompt. + assert calls[1][1] == {"role": "user", "content": "prove it"} + + +def test_decorating_sync_function_raises(): + with pytest.raises(TypeError): + + @route(RandomRouter([WeightedModel("m", 1.0)], rng_seed=1)) + def chat(model: str, prompt: str) -> str: # not async + return prompt diff --git a/tests/test_agentapi_llm_class.py b/tests/test_agentapi_llm_class.py new file mode 100644 index 00000000..014e7070 --- /dev/null +++ b/tests/test_agentapi_llm_class.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the LLM-classifier router (ports llm_class.rs).""" + +from __future__ import annotations + +import pytest + +from switchyard.lib.agentapi.chat import ChatRequest, ChatResponse, EnrichmentData +from switchyard.lib.agentapi.llm_class import ClassifierTier, LlmClassifier +from switchyard.lib.agentapi.optimizer import ( + ModelInference, + RequestInput, + ResponseInput, + Return, +) + + +def _algorithm(threshold: float) -> LlmClassifier: + return LlmClassifier( + classifier_model="router/classifier", + strong_model="frontier/model", + weak_model="cheap/model", + threshold=threshold, + ) + + +async def _run_flow(threshold: float, user_prompt: str, score_text: str): + opt = _algorithm(threshold).optimizer() + + # 1. Feed the user request; first inference is the classifier call. + await opt.feed(RequestInput(ChatRequest(user_prompt, "client/model")), EnrichmentData()) + classify = await opt.optimize() + assert isinstance(classify, ModelInference) + call = classify.response.requests[0] + assert user_prompt in call.prompt + assert call.model == "router/classifier" + + # 2. Feed the mocked classifier score; next inference is the routed call. + await opt.feed(ResponseInput(ChatResponse(score_text)), EnrichmentData()) + routed = await opt.optimize() + assert isinstance(routed, ModelInference) + decision = routed.response.decision_info + routed_model = routed.response.requests[0].model + + # 3. Feed the routed model response; next optimize returns to the agent. + await opt.feed(ResponseInput(ChatResponse("mocked completion")), EnrichmentData()) + assert isinstance(await opt.optimize(), Return) + return routed_model, decision + + +async def test_score_at_or_above_threshold_routes_strong(): + model, decision = await _run_flow(0.5, "solve this proof", "0.9") + assert model == "frontier/model" + assert decision.tier is ClassifierTier.STRONG + assert decision.score == 0.9 + + +async def test_score_below_threshold_routes_weak(): + model, decision = await _run_flow(0.5, "say hello", "0.2") + assert model == "cheap/model" + assert decision.tier is ClassifierTier.WEAK + assert decision.score == 0.2 + + +async def test_score_exactly_at_threshold_routes_strong(): + model, decision = await _run_flow(0.5, "borderline", "0.5") + assert model == "frontier/model" + assert decision.tier is ClassifierTier.STRONG + + +async def test_unparseable_score_defaults_to_strong(): + model, decision = await _run_flow(0.5, "hi", "not-a-number") + assert model == "frontier/model" + assert decision.tier is ClassifierTier.STRONG + assert decision.score is None + + +async def test_optimize_before_feed_errors(): + opt = _algorithm(0.5).optimizer() + with pytest.raises(ValueError): + await opt.optimize() + + +async def test_optimize_before_classifier_response_errors(): + opt = _algorithm(0.5).optimizer() + await opt.feed(RequestInput(ChatRequest("hi", "client/model")), EnrichmentData()) + assert isinstance(await opt.optimize(), ModelInference) # classifier call emitted + with pytest.raises(ValueError): + await opt.optimize() # no score fed yet diff --git a/tests/test_agentapi_optimizer.py b/tests/test_agentapi_optimizer.py new file mode 100644 index 00000000..92a4c881 --- /dev/null +++ b/tests/test_agentapi_optimizer.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the agentapi optimizer role interfaces.""" + +from __future__ import annotations + +import pytest + +from switchyard.lib.agentapi.chat import ChatRequest, ChatResponse, EnrichmentData +from switchyard.lib.agentapi.optimizer import ( + AgentApiOptimizer, + Decision, + MetadataInput, + ModelInference, + OptimizerResponse, + RequestInput, + ResponseInput, + Return, +) + + +def test_input_variants_carry_payloads(): + assert RequestInput(ChatRequest("hi", "m")).request.prompt == "hi" + assert ResponseInput(ChatResponse("done")).response.completion == "done" + assert MetadataInput({"k": "v"}).metadata == {"k": "v"} + + +def test_decision_variants(): + resp = OptimizerResponse(requests=[ChatRequest("hi", "m")]) + inf = ModelInference(resp) + assert isinstance(inf, Decision) + assert inf.response.requests[0].model == "m" + assert isinstance(Return(), Decision) + + +def test_optimizer_is_abstract(): + with pytest.raises(TypeError): + AgentApiOptimizer() # optimize() is abstract + + +async def test_feed_has_noop_default(): + class Opt(AgentApiOptimizer): + async def optimize(self) -> Decision: + return Return() + + opt = Opt() + # Default feed does nothing and does not raise. + await opt.feed(RequestInput(ChatRequest("hi", "m")), EnrichmentData()) + assert isinstance(await opt.optimize(), Return) diff --git a/tests/test_agentapi_public_api.py b/tests/test_agentapi_public_api.py new file mode 100644 index 00000000..71d1ae21 --- /dev/null +++ b/tests/test_agentapi_public_api.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The agentapi subpackage exports its public API.""" + +from __future__ import annotations + +import switchyard.lib.agentapi as agentapi + + +def test_public_exports_present(): + for name in [ + "route", + "RandomRouter", + "WeightedModel", + "LlmClassifier", + "ClassifierTier", + "ChatRequest", + "ChatResponse", + "EnrichmentData", + "AgentApiOptimizer", + "AgentApiOptAlgorithm", + "Decision", + "ModelInference", + "Return", + ]: + assert hasattr(agentapi, name), name + assert name in agentapi.__all__ diff --git a/tests/test_agentapi_rand.py b/tests/test_agentapi_rand.py new file mode 100644 index 00000000..711756c3 --- /dev/null +++ b/tests/test_agentapi_rand.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the weighted random router (ports rand.rs).""" + +from __future__ import annotations + +import pytest + +from switchyard.lib.agentapi.chat import ChatRequest, EnrichmentData +from switchyard.lib.agentapi.optimizer import ( + AgentApiOptimizer, + ModelInference, + RequestInput, + ResponseInput, + Return, +) +from switchyard.lib.agentapi.rand import RandomRouter, WeightedModel + + +def _algorithm(models: list[WeightedModel], seed: int) -> RandomRouter: + return RandomRouter(models=models, rng_seed=seed) + + +async def _route_once(opt: AgentApiOptimizer) -> str: + await opt.feed(RequestInput(ChatRequest("hi", "client/model")), EnrichmentData()) + decision = await opt.optimize() + assert isinstance(decision, ModelInference) + return decision.response.requests[0].model + + +async def test_all_weight_on_one_model_always_selects_it(): + algo = _algorithm( + [WeightedModel("a/model", 0.0), WeightedModel("b/model", 1.0), WeightedModel("c/model", 0.0)], + seed=7, + ) + opt = algo.optimizer() + for _ in range(25): + assert await _route_once(opt) == "b/model" + + +async def test_selection_frequencies_track_weights(): + algo = _algorithm( + [WeightedModel("a/model", 1.0), WeightedModel("b/model", 3.0), WeightedModel("c/model", 6.0)], + seed=42, + ) + opt = algo.optimizer() + counts = {"a/model": 0, "b/model": 0, "c/model": 0} + draws = 20_000 + for _ in range(draws): + counts[await _route_once(opt)] += 1 + assert abs(counts["a/model"] / draws - 0.1) < 0.02 + assert abs(counts["b/model"] / draws - 0.3) < 0.02 + assert abs(counts["c/model"] / draws - 0.6) < 0.02 + + +async def test_decision_reports_draw_within_total_weight(): + opt = _algorithm([WeightedModel("a/model", 2.0), WeightedModel("b/model", 3.0)], seed=7).optimizer() + await opt.feed(RequestInput(ChatRequest("hi", "client/model")), EnrichmentData()) + decision = await opt.optimize() + assert isinstance(decision, ModelInference) + info = decision.response.decision_info + assert info.total_weight == 5.0 + assert 0.0 <= info.draw < 5.0 + + +async def test_returns_to_agent_after_response_is_fed(): + opt = _algorithm([WeightedModel("frontier/model", 1.0)], seed=7).optimizer() + await opt.feed(RequestInput(ChatRequest("hi", "client/model")), EnrichmentData()) + first = await opt.optimize() + assert isinstance(first, ModelInference) + assert first.response.requests[0].model == "frontier/model" + await opt.feed(ResponseInput_from("mocked completion"), EnrichmentData()) + assert isinstance(await opt.optimize(), Return) + + +async def test_optimize_before_feed_errors(): + opt = _algorithm([WeightedModel("a/model", 1.0)], seed=7).optimizer() + with pytest.raises(ValueError): + await opt.optimize() + + +async def test_no_positive_weight_errors(): + opt = _algorithm([WeightedModel("a/model", 0.0), WeightedModel("b/model", 0.0)], seed=7).optimizer() + await opt.feed(RequestInput(ChatRequest("hi", "client/model")), EnrichmentData()) + with pytest.raises(ValueError): + await opt.optimize() + + +async def test_factory_mints_independent_deterministic_optimizers(): + algo = _algorithm([WeightedModel("a/model", 1.0), WeightedModel("b/model", 1.0)], seed=42) + first, second = algo.optimizer(), algo.optimizer() + + async def first_draw(opt: AgentApiOptimizer) -> float: + await opt.feed(RequestInput(ChatRequest("hi", "client/model")), EnrichmentData()) + decision = await opt.optimize() + assert isinstance(decision, ModelInference) + return decision.response.decision_info.draw + + assert await first_draw(first) == await first_draw(second) + + +def ResponseInput_from(text: str) -> ResponseInput: + from switchyard.lib.agentapi.chat import ChatResponse + + return ResponseInput(ChatResponse(text))