diff --git a/Cargo.lock b/Cargo.lock index 402d83896..e73214a56 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -416,6 +416,21 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -423,6 +438,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -431,6 +447,17 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.32" @@ -466,6 +493,7 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -821,6 +849,29 @@ 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", + "futures", + "serde_json", + "tokio", + "tokio-stream", +] + +[[package]] +name = "libsy-examples" +version = "0.1.0" +dependencies = [ + "async-trait", + "futures", + "libsy", + "rand 0.8.6", + "tokio", + "tokio-stream", +] + [[package]] name = "libyaml-rs" version = "0.3.0" @@ -1728,6 +1779,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", @@ -1756,6 +1808,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" diff --git a/Cargo.toml b/Cargo.toml index 7c7c18eb4..3295507e0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,8 @@ [workspace] resolver = "2" members = [ + "crates/libsy", + "crates/libsy-examples", "crates/switchyard-components", "crates/switchyard-components-v2", "crates/switchyard-components-v2-macros", diff --git a/crates/libsy-examples/Cargo.toml b/crates/libsy-examples/Cargo.toml new file mode 100644 index 000000000..082bec560 --- /dev/null +++ b/crates/libsy-examples/Cargo.toml @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libsy-examples" +version = "0.1.0" +description = "Reference algorithms and runnable agents built on libsy" +publish = false +authors.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true + +[dependencies] +libsy = { path = "../libsy" } +async-trait = "0.1" +rand = "0.8" +futures = "0.3" +tokio = { version = "1", features = ["full"] } +tokio-stream = "0.1" diff --git a/crates/libsy-examples/examples/research_agent.rs b/crates/libsy-examples/examples/research_agent.rs new file mode 100644 index 000000000..8038ca383 --- /dev/null +++ b/crates/libsy-examples/examples/research_agent.rs @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Minimal research agent using the [`Algorithm::run`] convenience. +//! +//! Every target owns an `LlmClient`, so the agent runs each request to completion with +//! [`Algorithm::run`]: it serves each offloaded call with the routed +//! target's `default_client` and returns the final response — no stream to drive. The +//! multi-step routing (classify -> route) happens inside the classifier algorithm; the +//! agent never sees it. To drive the step stream yourself instead, use +//! `Algorithm::run_stream`. Run with: +//! cargo run -p libsy --example research_agent + +use std::error::Error; +use std::sync::Arc; + +use async_trait::async_trait; +use libsy::{ + Algorithm, Context, LlmClient, LlmRequest, LlmResponse, LlmTarget, LlmTargetSet, Request, + Response, RoutedRequest, +}; +use libsy_examples::llm_class::LlmClassifierOrchAlgo; + +const CLASSIFIER: &str = "classifier/model"; +const STRONG: &str = "strong/model"; +const WEAK: &str = "weak/model"; + +/// Stub transport. Real integrators implement `LlmClient` over their own HTTP. +struct StubClient; + +#[async_trait] +impl LlmClient for StubClient { + async fn call(&self, routed: RoutedRequest) -> Result> { + // The model to call is the routed decision's selection, not the inbound name. + let model = routed.decision.selected_model().to_string(); + println!(" -> model call: {model}"); + // The classifier returns a score; other models return an answer. + let completion = if model == CLASSIFIER { + "0.9".to_string() + } else { + format!("answer from {model}") + }; + Ok(Response { + llm_response: LlmResponse { + completion, + raw_response: None, + }, + metadata: None, + }) + } +} + +fn targets() -> LlmTargetSet { + let client = Arc::new(StubClient) as Arc; + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(client.clone()), + }; + LlmTargetSet::new(vec![target(CLASSIFIER), target(STRONG), target(WEAK)]) +} + +struct ResearchAgent { + algo: Arc, +} + +impl ResearchAgent { + /// Trivial plan: one lookup per question (stub). + fn plan(&self, question: &str) -> Vec { + vec![format!("look up: {question}")] + } + + async fn run(&self, question: &str) -> Result> { + let mut notes = Vec::new(); + for step in self.plan(question) { + let request = Request { + llm_request: LlmRequest { + inbound_model_name: "auto".to_string(), + prompt: step, + }, + raw_request: None, + metadata: None, + }; + + let (_trace, response) = self.algo.clone().run(Context::default(), request).await?; + notes.push(response.llm_response.completion); + } + Ok(notes.join("\n")) + } +} + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Box> { + // Configure routing once: an LLM classifier over three named targets. Swapping + // in `RandomOrchAlgo` needs no change to the agent. + let algo: Arc = Arc::new(LlmClassifierOrchAlgo::new( + CLASSIFIER, + STRONG, + WEAK, + 0.5, + targets(), + )); + + let agent = ResearchAgent { algo }; + println!("{}", agent.run("what is switchyard?").await?); + Ok(()) +} diff --git a/crates/libsy-examples/examples/research_agent_core.rs b/crates/libsy-examples/examples/research_agent_core.rs new file mode 100644 index 000000000..bdb05ae8c --- /dev/null +++ b/crates/libsy-examples/examples/research_agent_core.rs @@ -0,0 +1,118 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Research agent driving the raw `run` stream with **client-less** targets. +//! +//! With no client, every `driver.call_llm_target` is offloaded as a promise the orchestrator +//! surfaces as a `CallLlm` step. The agent makes the "real" model call itself and +//! fulfills the promise — this is the offload/streaming path ("ask, don't call"). +//! The classifier's two steps show up as two `model call:` lines. Run with: +//! cargo run -p libsy --example research_agent_core + +use std::error::Error; +use std::sync::Arc; + +use libsy::{ + Algorithm, Context, Decision, LlmRequest, LlmResponse, LlmTarget, LlmTargetSet, Request, + Response, Step, +}; +use libsy_examples::llm_class::LlmClassifierOrchAlgo; +use tokio_stream::StreamExt; + +const CLASSIFIER: &str = "classifier/model"; +const STRONG: &str = "strong/model"; +const WEAK: &str = "weak/model"; + +/// The "real" model call the agent makes to fulfill a promise. The core never +/// makes the call itself — it hands back a request and waits for the response. +/// The model to call is the routing decision's selection, read off the promise. +async fn call_model(model: &str) -> Response { + println!(" -> model call: {model}"); + let completion = if model == CLASSIFIER { + "0.9".to_string() + } else { + format!("answer from {model}") + }; + Response { + llm_response: LlmResponse { + completion, + raw_response: None, + }, + metadata: None, + } +} + +fn targets() -> LlmTargetSet { + // Client-less targets -> every call is offloaded via a promise. + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: None, + }; + LlmTargetSet::new(vec![target(CLASSIFIER), target(STRONG), target(WEAK)]) +} + +struct ResearchAgent { + algo: Arc, +} + +impl ResearchAgent { + /// Trivial plan: one lookup per question (stub). + fn plan(&self, question: &str) -> Vec { + vec![format!("look up: {question}")] + } + + async fn run(&mut self, question: &str) -> Result> { + let mut notes = Vec::new(); + for step in self.plan(question) { + let request = Request { + llm_request: LlmRequest { + inbound_model_name: "auto".to_string(), + prompt: step, + }, + raw_request: None, + metadata: None, + }; + let stream = self.algo.clone().run_stream(Context::default(), request); + tokio::pin!(stream); + while let Some(update) = stream.next().await { + match update? { + Step::CallLlm(call) => { + // Perform the model call the algorithm asked for, then fulfill. + let response = call_model(call.get_decision()?.selected_model()).await; + call.respond(Ok(response))?; + } + // Decisions stream in as the algorithm makes them. + Step::Decision(decision) => print_decision(decision.as_ref()), + Step::ReturnToAgent(response) => { + notes.push(response.llm_response.completion); + } + } + } + } + Ok(notes.join("\n")) + } +} + +/// Print one decision the algorithm recorded — uniform access via the trait. +fn print_decision(decision: &dyn Decision) { + println!( + " decision: {} ({})", + decision.selected_model(), + decision.reasoning().unwrap_or_default() + ); +} + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Box> { + let algo: Arc = Arc::new(LlmClassifierOrchAlgo::new( + CLASSIFIER, + STRONG, + WEAK, + 0.5, + targets(), + )); + + let mut agent = ResearchAgent { algo }; + println!("{}", agent.run("what is switchyard?").await?); + Ok(()) +} diff --git a/crates/libsy-examples/src/ensemble.rs b/crates/libsy-examples/src/ensemble.rs new file mode 100644 index 000000000..4f4a9fc9d --- /dev/null +++ b/crates/libsy-examples/src/ensemble.rs @@ -0,0 +1,822 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Ensemble router built on the [`Algorithm`] interfaces. +//! +//! Each request is fanned out to a set of candidate models concurrently; a judge +//! model (e.g. Haiku) then picks the best response, which is returned to the +//! agent. The algorithm tallies which candidate the judge preferred across +//! requests and, after `exploration_turns` ensemble turns, commits to the +//! winningest model — every subsequent request routes straight to that one model +//! with no fan-out and no judge call. +//! +//! Unlike the reference routers, this algorithm is **stateful**: the win tally, +//! turn counter, and committed choice live behind a [`std::sync::Mutex`] so one +//! shared `&self` can serve a session's requests concurrently (see the +//! `Algorithm` docs). In a proxy setup one `EnsembleOrchAlgo` is created per session, +//! so this state is per-session. + +use std::collections::BTreeMap; +use std::error::Error; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; + +use libsy::{ + Algorithm, Context, Decision, Driver, LlmRequest, LlmTargetSet, Request, Response, Signals, +}; + +/// Which step of the ensemble flow produced a decision. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EnsemblePhase { + /// A fan-out call to one candidate model during an exploration turn. + Candidate, + /// The judge call that scored the candidate responses. + Judge, + /// The candidate the judge selected on an exploration turn. + Winner, + /// The single model the algorithm committed to after exploration. + Committed, +} + +impl EnsemblePhase { + /// Stable string form of the phase, used in decision reasoning. + pub fn as_str(self) -> &'static str { + match self { + EnsemblePhase::Candidate => "candidate", + EnsemblePhase::Judge => "judge", + EnsemblePhase::Winner => "winner", + EnsemblePhase::Committed => "committed", + } + } +} + +/// Decision produced at each step of the ensemble flow. +pub struct EnsembleDecision { + /// The model this step concerns (a candidate, the judge, the winner, or the + /// committed model). + pub selected_model: String, + /// Human-readable explanation of the step. + pub reasoning: String, + /// Which step of the ensemble flow produced this decision. + pub phase: EnsemblePhase, +} + +impl Decision for EnsembleDecision { + fn selected_model(&self) -> &str { + &self.selected_model + } + fn reasoning(&self) -> Option<&str> { + Some(&self.reasoning) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// Mutable per-session state, guarded by a [`Mutex`] so `&self` can be shared +/// across concurrent requests. Never held across an `await`. +struct EnsembleState { + /// Judge-win count per candidate model. + wins: BTreeMap, + /// Completed ensemble (exploration) turns. + turns: u64, + /// The model committed to once exploration is over; `None` while exploring. + committed: Option, +} + +/// Ensemble router: fan out to candidates, judge the best, then commit. +pub struct EnsembleOrchAlgo { + candidate_models: Vec, + judge_model: String, + /// Number of ensemble turns to run before committing to the best model. + /// `0` disables committing — the algorithm ensembles on every request. + exploration_turns: u64, + target_set: LlmTargetSet, + state: Mutex, +} + +impl EnsembleOrchAlgo { + /// Create an ensemble over `candidate_models`, judged by `judge_model`, + /// exploring for `exploration_turns` before committing to the winningest + /// candidate (`0` = never commit, ensemble every request), routing among + /// `target_set`. Wrap it in an [`Arc`](std::sync::Arc) and drive it with + /// [`run`](libsy::Algorithm::run) or + /// [`run_stream`](libsy::Algorithm::run_stream). + pub fn new( + candidate_models: Vec, + judge_model: impl Into, + exploration_turns: u64, + target_set: LlmTargetSet, + ) -> Self { + Self { + candidate_models, + judge_model: judge_model.into(), + exploration_turns, + target_set, + state: Mutex::new(EnsembleState { + wins: BTreeMap::new(), + turns: 0, + committed: None, + }), + } + } + + /// If exploration is over, return the committed model (committing lazily on + /// the first post-exploration request). `None` means keep ensembling. + /// + /// The lock is taken and dropped here, never across an `await`. Under + /// concurrency two requests may both commit; they pick the same model from + /// the same tally (with a stable tie-break), so the result is identical. + fn resolve_committed(&self) -> Result, Box> { + let mut state = self + .state + .lock() + .map_err(|_| "ensemble state lock poisoned")?; + if let Some(model) = &state.committed { + return Ok(Some(model.clone())); + } + // `exploration_turns == 0` keeps the algorithm in ensemble mode forever. + if self.exploration_turns > 0 && state.turns >= self.exploration_turns { + let best = self.pick_best(&state.wins)?; + state.committed = Some(best.clone()); + return Ok(Some(best)); + } + Ok(None) + } + + /// The candidate with the most judge-wins, breaking ties toward the earlier + /// candidate (stable and deterministic). Errors only if there are no + /// candidates configured. + fn pick_best( + &self, + wins: &BTreeMap, + ) -> Result> { + let mut best = self + .candidate_models + .first() + .ok_or("no candidate models configured")?; + let mut best_wins = wins.get(best).copied().unwrap_or(0); + for model in &self.candidate_models[1..] { + let w = wins.get(model).copied().unwrap_or(0); + if w > best_wins { + best = model; + best_wins = w; + } + } + Ok(best.clone()) + } + + /// Route a request to a single already-chosen model — the committed fast path. + async fn route_committed( + &self, + driver: &Driver, + request: Request, + model: String, + ) -> Result<(Vec>, Response), Box> { + let target = self.target_set.get_target(&model)?; + let decision: Arc = Arc::new(EnsembleDecision { + reasoning: format!( + "committed to '{model}' after {} turns", + self.exploration_turns + ), + selected_model: model.clone(), + phase: EnsemblePhase::Committed, + }); + let routed = Request { + llm_request: LlmRequest { + // The agent's inbound name rides through; the committed model is on + // the decision, not stamped onto the request. + inbound_model_name: request.llm_request.inbound_model_name, + prompt: request.llm_request.prompt, + }, + raw_request: request.raw_request, + metadata: request.metadata, + }; + let response = driver + .call_llm_target(&target, routed, decision.clone()) + .await?; + Ok((vec![decision], response)) + } + + /// One exploration turn: fan out to every candidate, judge the survivors, + /// tally the winner, and return its response. + async fn ensemble_turn( + &self, + driver: &Driver, + request: Request, + ) -> Result<(Vec>, Response), Box> { + let user_prompt = request.llm_request.prompt.clone(); + // The agent's inbound name rides through every sub-call unchanged; the model + // each call hits is carried by its decision, not stamped onto the request. + let inbound = request.llm_request.inbound_model_name.clone(); + + // Fan out to all candidates concurrently with the same user prompt. Each + // call is annotated with its own candidate decision so the caller can see + // which model an offloaded call targets. + let mut candidate_decisions: Vec> = Vec::new(); + let mut calls = Vec::new(); + for model in &self.candidate_models { + let target = self.target_set.get_target(model)?; + let decision: Arc = Arc::new(EnsembleDecision { + selected_model: model.clone(), + reasoning: format!("ensemble candidate '{model}'"), + phase: EnsemblePhase::Candidate, + }); + candidate_decisions.push(decision.clone()); + let call_request = Request { + llm_request: LlmRequest { + inbound_model_name: inbound.clone(), + prompt: user_prompt.clone(), + }, + raw_request: request.raw_request.clone(), + metadata: request.metadata.clone(), + }; + let model = model.clone(); + calls.push(async move { + ( + model, + driver + .call_llm_target(&target, call_request, decision) + .await, + ) + }); + } + let results = futures::future::join_all(calls).await; + + // Keep only successful responses, preserving candidate order. A failed + // candidate is simply excluded from judging rather than failing the turn. + let mut survivors: Vec<(String, Response)> = Vec::new(); + for (model, result) in results { + if let Ok(response) = result { + survivors.push((model, response)); + } + } + if survivors.is_empty() { + return Err("all ensemble candidates failed".into()); + } + + // Pick the winner: judge only when there is a real choice to make. + let (winner_model, winner_response, judge_decision) = if survivors.len() == 1 { + let (model, response) = survivors + .into_iter() + .next() + .ok_or("survivor unexpectedly missing")?; + (model, response, None) + } else { + let judge_prompt = build_judge_prompt(&user_prompt, &survivors); + let judge_target = self.target_set.get_target(&self.judge_model)?; + let judge_decision: Arc = Arc::new(EnsembleDecision { + selected_model: self.judge_model.clone(), + reasoning: format!("judging {} candidate responses", survivors.len()), + phase: EnsemblePhase::Judge, + }); + let judge_request = Request { + llm_request: LlmRequest { + inbound_model_name: inbound.clone(), + prompt: judge_prompt, + }, + raw_request: request.raw_request.clone(), + metadata: request.metadata.clone(), + }; + let judge_response = driver + .call_llm_target(&judge_target, judge_request, judge_decision.clone()) + .await?; + // Fail open: an unparseable pick falls back to the first response. + let choice = parse_choice(&judge_response.llm_response.completion, survivors.len()); + let (model, response) = survivors + .into_iter() + .nth(choice) + .ok_or("judge choice out of range")?; + (model, response, Some(judge_decision)) + }; + + // Record the win and advance the turn counter under the lock (not held + // across any await). + { + let mut state = self + .state + .lock() + .map_err(|_| "ensemble state lock poisoned")?; + *state.wins.entry(winner_model.clone()).or_insert(0) += 1; + state.turns += 1; + } + + let winner_decision: Arc = Arc::new(EnsembleDecision { + reasoning: format!("judge selected '{winner_model}' as best response"), + selected_model: winner_model, + phase: EnsemblePhase::Winner, + }); + + // Trace order: [candidate calls..., judge?, winner]. + let mut trace = candidate_decisions; + if let Some(judge_decision) = judge_decision { + trace.push(judge_decision); + } + trace.push(winner_decision); + Ok((trace, winner_response)) + } +} + +/// Build the judge prompt. Responses are presented anonymously (no model names) +/// so the judge scores on content alone rather than model reputation. +fn build_judge_prompt(user_prompt: &str, survivors: &[(String, Response)]) -> String { + let mut prompt = String::from( + "You are an impartial judge. Choose which response best answers the user request.\n\n", + ); + prompt.push_str("User request:\n"); + prompt.push_str(user_prompt); + prompt.push_str("\n\n"); + for (i, (_model, response)) in survivors.iter().enumerate() { + prompt.push_str(&format!( + "Response {}:\n{}\n\n", + i + 1, + response.llm_response.completion + )); + } + prompt.push_str(&format!( + "Reply with only the number (1-{}) of the best response.", + survivors.len() + )); + prompt +} + +/// Parse the judge's 1-based pick into a 0-based index, failing open to the +/// first response. Reads the first run of digits in the reply, so "2" or +/// "Response 2 is best" both select index 1. +fn parse_choice(completion: &str, count: usize) -> usize { + let mut digits = String::new(); + for c in completion.chars() { + if c.is_ascii_digit() { + digits.push(c); + } else if !digits.is_empty() { + break; + } + } + match digits.parse::() { + Ok(n) if n >= 1 && n <= count => n - 1, + _ => 0, + } +} + +#[async_trait] +impl Algorithm for EnsembleOrchAlgo { + async fn create_run_task( + self: Arc, + _ctx: Context, + driver: Driver, + request: Request, + ) -> Result> { + // Fast path: exploration is over — route straight to the committed model; + // otherwise run a full ensemble turn. Both return a decision trace plus the + // final response. + let (trace, response) = if let Some(model) = self.resolve_committed()? { + self.route_committed(&driver, request, model).await? + } else { + self.ensemble_turn(&driver, request).await? + }; + // Publish the trace to the stream (candidate..., judge?, winner). The + // candidate decisions also rode along on their offloaded `CallLlm` steps. + for decision in trace { + driver.info(decision).await?; + } + Ok(response) + } + + async fn process_signals( + self: Arc, + _signals: Signals, + ) -> Result<(), Box> { + // Success is measured by the judge, not agent-system signals. + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use libsy::{LlmClient, LlmRequest, LlmResponse, LlmTarget, Response, RoutedRequest}; + use std::sync::Mutex as StdMutex; + + /// Mock client that answers candidate calls with `answer from {model}` and, + /// for the judge model, returns the 1-based number of the response whose + /// content mentions `prefer` (so a test controls which candidate "wins"). + /// Records every model it was called with for call-count assertions. + struct JudgingClient { + judge_model: String, + prefer: String, + calls: Arc>>, + } + + #[async_trait] + impl LlmClient for JudgingClient { + async fn call( + &self, + routed: RoutedRequest, + ) -> Result> { + let name = routed.decision.selected_model().to_string(); + self.calls + .lock() + .map_err(|_| "lock poisoned")? + .push(name.clone()); + let completion = if name == self.judge_model { + judge_pick(&routed.request.llm_request.prompt, &self.prefer) + } else { + format!("answer from {name}") + }; + Ok(Response { + llm_response: LlmResponse { + completion, + raw_response: None, + }, + metadata: None, + }) + } + } + + /// Scan a judge prompt for the `Response N:` whose body is `answer from + /// {prefer}` and return `N` as a string; defaults to "1". + fn judge_pick(prompt: &str, prefer: &str) -> String { + let target_line = format!("answer from {prefer}"); + let mut current = 1u32; + for line in prompt.lines() { + if let Some(rest) = line.strip_prefix("Response ") { + if let Ok(num) = rest.trim_end_matches(':').parse::() { + current = num; + } + } else if line == target_line { + return current.to_string(); + } + } + "1".to_string() + } + + /// Build an ensemble algo over `candidates` + a judge, all backed by one + /// judging client that prefers `prefer`. Returns the algo and the shared + /// call log. + fn algo( + candidates: &[&str], + judge: &str, + prefer: &str, + exploration_turns: u64, + ) -> (EnsembleOrchAlgo, Arc>>) { + let calls = Arc::new(StdMutex::new(Vec::new())); + let client = Arc::new(JudgingClient { + judge_model: judge.to_string(), + prefer: prefer.to_string(), + calls: Arc::clone(&calls), + }) as Arc; + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(client.clone()), + }; + let mut targets: Vec = candidates.iter().map(|n| target(n)).collect(); + targets.push(target(judge)); + let algo = EnsembleOrchAlgo::new( + candidates.iter().map(|s| s.to_string()).collect(), + judge.to_string(), + exploration_turns, + LlmTargetSet::new(targets), + ); + (algo, calls) + } + + /// Build an ensemble algo whose candidate + judge targets all share `client`. + /// One such algo models a single session's stateful router. + fn algo_with_client( + candidates: &[&str], + judge: &str, + exploration_turns: u64, + client: Arc, + ) -> EnsembleOrchAlgo { + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(client.clone()), + }; + let mut targets: Vec = candidates.iter().map(|n| target(n)).collect(); + targets.push(target(judge)); + EnsembleOrchAlgo::new( + candidates.iter().map(|s| s.to_string()).collect(), + judge.to_string(), + exploration_turns, + LlmTargetSet::new(targets), + ) + } + + fn request(prompt: &str) -> Request { + Request { + llm_request: LlmRequest { + inbound_model_name: "auto".to_string(), + prompt: prompt.to_string(), + }, + raw_request: None, + metadata: None, + } + } + + /// Wrap an ensemble algo as `Arc` we can drive to completion. + /// Reuse one handle across requests to exercise the algo's per-session state. + fn orch(algo: EnsembleOrchAlgo) -> Arc { + Arc::new(algo) + } + + fn as_ensemble( + d: &Arc, + ) -> Result<&EnsembleDecision, Box> { + d.as_any() + .downcast_ref::() + .ok_or_else(|| "expected an EnsembleDecision".into()) + } + + #[tokio::test] + async fn exploration_turn_fans_out_judges_and_returns_the_winner( + ) -> Result<(), Box> { + // Judge prefers b/model; it should win and be returned. + let (algo, calls) = algo(&["a/model", "b/model"], "judge/haiku", "b/model", 100); + let (trace, response) = orch(algo) + .run(Context::default(), request("solve it")) + .await?; + assert_eq!(response.llm_response.completion, "answer from b/model"); + + // Both candidates and the judge were called. + let calls = calls.lock().map_err(|_| "lock poisoned")?; + assert!(calls.contains(&"a/model".to_string())); + assert!(calls.contains(&"b/model".to_string())); + assert!(calls.contains(&"judge/haiku".to_string())); + + // Trace: [candidate a, candidate b, judge, winner]. + assert_eq!(trace.len(), 4); + assert_eq!(as_ensemble(&trace[0])?.phase, EnsemblePhase::Candidate); + assert_eq!(as_ensemble(&trace[2])?.phase, EnsemblePhase::Judge); + let winner = as_ensemble(&trace[3])?; + assert_eq!(winner.phase, EnsemblePhase::Winner); + assert_eq!(winner.selected_model, "b/model"); + Ok(()) + } + + #[tokio::test] + async fn commits_to_the_winningest_model_after_exploration( + ) -> Result<(), Box> { + // Judge always prefers b/model over 2 exploration turns, so the algo + // commits to b/model even though a/model is listed first. + let (algo, calls) = algo(&["a/model", "b/model"], "judge/haiku", "b/model", 2); + let orch = orch(algo); + + // Two exploration turns. + orch.clone().run(Context::default(), request("t1")).await?; + orch.clone().run(Context::default(), request("t2")).await?; + let judge_calls_after_exploration = calls + .lock() + .map_err(|_| "lock poisoned")? + .iter() + .filter(|c| *c == "judge/haiku") + .count(); + assert_eq!(judge_calls_after_exploration, 2); + + // Third request: committed fast path — routes straight to b/model with no + // fan-out to a/model and no judge call. + let (trace, response) = orch.clone().run(Context::default(), request("t3")).await?; + assert_eq!(response.llm_response.completion, "answer from b/model"); + assert_eq!(trace.len(), 1); + let decision = as_ensemble(&trace[0])?; + assert_eq!(decision.phase, EnsemblePhase::Committed); + assert_eq!(decision.selected_model, "b/model"); + + let calls = calls.lock().map_err(|_| "lock poisoned")?; + // Judge was not called again on the committed turn. + assert_eq!(calls.iter().filter(|c| *c == "judge/haiku").count(), 2); + // a/model was called only on the two exploration turns, not the third. + assert_eq!(calls.iter().filter(|c| *c == "a/model").count(), 2); + Ok(()) + } + + #[tokio::test] + async fn single_candidate_skips_the_judge() -> Result<(), Box> { + let (algo, calls) = algo(&["only/model"], "judge/haiku", "only/model", 100); + let (trace, response) = orch(algo).run(Context::default(), request("hi")).await?; + assert_eq!(response.llm_response.completion, "answer from only/model"); + // No judge call for a lone candidate. + assert!(!calls + .lock() + .map_err(|_| "lock poisoned")? + .contains(&"judge/haiku".to_string())); + // Trace: [candidate, winner] — no judge entry. + assert_eq!(trace.len(), 2); + assert_eq!(as_ensemble(&trace[1])?.phase, EnsemblePhase::Winner); + Ok(()) + } + + #[tokio::test] + async fn zero_exploration_turns_never_commits() -> Result<(), Box> { + // exploration_turns == 0 keeps ensembling forever. + let (algo, calls) = algo(&["a/model", "b/model"], "judge/haiku", "b/model", 0); + let orch = orch(algo); + for _ in 0..3 { + let (trace, _) = orch.clone().run(Context::default(), request("x")).await?; + // Always a full ensemble turn (never a lone Committed decision). + assert_eq!( + as_ensemble(&trace[trace.len() - 1])?.phase, + EnsemblePhase::Winner + ); + } + // Judge ran on every turn. + assert_eq!( + calls + .lock() + .map_err(|_| "lock poisoned")? + .iter() + .filter(|c| *c == "judge/haiku") + .count(), + 3 + ); + Ok(()) + } + + #[tokio::test] + async fn all_candidates_failing_errors() -> Result<(), Box> { + /// Client whose every call fails. + struct FailingClient; + #[async_trait] + impl LlmClient for FailingClient { + async fn call( + &self, + _routed: RoutedRequest, + ) -> Result> { + Err("upstream down".into()) + } + } + let client = Arc::new(FailingClient) as Arc; + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(client.clone()), + }; + let algo = EnsembleOrchAlgo::new( + vec!["a/model".to_string(), "b/model".to_string()], + "judge/haiku", + 100, + LlmTargetSet::new(vec![ + target("a/model"), + target("b/model"), + target("judge/haiku"), + ]), + ); + assert!(orch(algo) + .run(Context::default(), request("x")) + .await + .is_err()); + Ok(()) + } + + #[tokio::test] + async fn process_signals_is_a_noop() -> Result<(), Box> { + let (algo, _) = algo(&["a/model"], "judge/haiku", "a/model", 1); + Arc::new(algo).process_signals(Signals {}).await?; + Ok(()) + } + + #[test] + fn parse_choice_reads_first_number_and_fails_open() { + assert_eq!(parse_choice("2", 3), 1); + assert_eq!(parse_choice("Response 3 is best", 3), 2); + assert_eq!(parse_choice("the winner is 1", 3), 0); + // Out of range and unparseable both fall open to the first response. + assert_eq!(parse_choice("7", 3), 0); + assert_eq!(parse_choice("none", 3), 0); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn two_sessions_process_in_parallel() -> Result<(), Box> { + use std::time::Duration; + use tokio::sync::Barrier; + + // Candidate calls block on a shared barrier; judge calls do not. Each session + // serves its offloaded candidate calls one at a time (the step stream is + // bounded), so a session has exactly one candidate call in flight at once. The + // two sessions run concurrently, so the barrier releases only when both have a + // candidate call pending. If the sessions were serialized, at most one call + // could be pending, the barrier would never reach 2, and the test would time + // out instead of passing. + struct BarrierClient { + barrier: Arc, + judge_model: String, + prefer: String, + } + + #[async_trait] + impl LlmClient for BarrierClient { + async fn call( + &self, + routed: RoutedRequest, + ) -> Result> { + let name = routed.decision.selected_model().to_string(); + let completion = if name == self.judge_model { + // Judge runs after the barrier releases; it must not wait. + judge_pick(&routed.request.llm_request.prompt, &self.prefer) + } else { + // Hold every candidate call until all sessions have fanned out. + self.barrier.wait().await; + format!("answer from {name}") + }; + Ok(Response { + llm_response: LlmResponse { + completion, + raw_response: None, + }, + metadata: None, + }) + } + } + + const SESSIONS: usize = 2; + let barrier = Arc::new(Barrier::new(SESSIONS)); + let client = Arc::new(BarrierClient { + barrier: barrier.clone(), + judge_model: "judge/haiku".to_string(), + prefer: "a/model".to_string(), + }) as Arc; + + // Two independent sessions: separate algo instances, each with its own + // per-session state, sharing only the backend client. + let session_a: Arc = Arc::new(algo_with_client( + &["a/model", "b/model"], + "judge/haiku", + 100, + client.clone(), + )); + let session_b: Arc = Arc::new(algo_with_client( + &["a/model", "b/model"], + "judge/haiku", + 100, + client.clone(), + )); + + let run = |session: Arc, prompt: &'static str| { + tokio::spawn(async move { + session + .run(Context::default(), request(prompt)) + .await + .map(|(_, response)| response.llm_response.completion) + }) + }; + let handle_a = run(session_a, "from A"); + let handle_b = run(session_b, "from B"); + + // The timeout converts a serialization deadlock into a failure, not a hang. + let completion_a = tokio::time::timeout(Duration::from_secs(5), handle_a).await???; + let completion_b = tokio::time::timeout(Duration::from_secs(5), handle_b).await???; + assert_eq!(completion_a, "answer from a/model"); + assert_eq!(completion_b, "answer from a/model"); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn two_parallel_sessions_keep_independent_state( + ) -> Result<(), Box> { + // Two sessions run concurrently with judges that prefer different models: + // session A prefers a/model, session B prefers b/model. Each explores for + // two turns then commits; because the win tally is per-session, they must + // commit to *different* models — proving no state leaks between sessions. + let (session_a, _) = algo(&["a/model", "b/model"], "judge/haiku", "a/model", 2); + let (session_b, _) = algo(&["a/model", "b/model"], "judge/haiku", "b/model", 2); + + // Drive one session's three requests sequentially (so its two exploration + // turns complete before the committing third), returning that third + // request's winning model and decision phase. + let drive = |session: Arc| { + tokio::spawn(async move { + session + .clone() + .run(Context::default(), request("t1")) + .await?; + session + .clone() + .run(Context::default(), request("t2")) + .await?; + let (trace, response) = session + .clone() + .run(Context::default(), request("t3")) + .await?; + let phase = trace + .last() + .and_then(|d| d.as_any().downcast_ref::()) + .map(|d| d.phase) + .ok_or("missing final decision")?; + Ok::<(String, EnsemblePhase), Box>(( + response.llm_response.completion, + phase, + )) + }) + }; + // The two sessions run in parallel; each committed independently. + let handle_a = drive(orch(session_a)); + let handle_b = drive(orch(session_b)); + let (completion_a, phase_a) = handle_a.await??; + let (completion_b, phase_b) = handle_b.await??; + + assert_eq!(phase_a, EnsemblePhase::Committed); + assert_eq!(completion_a, "answer from a/model"); + assert_eq!(phase_b, EnsemblePhase::Committed); + assert_eq!(completion_b, "answer from b/model"); + Ok(()) + } +} diff --git a/crates/libsy-examples/src/lib.rs b/crates/libsy-examples/src/lib.rs new file mode 100644 index 000000000..91874debf --- /dev/null +++ b/crates/libsy-examples/src/lib.rs @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Reference algorithms for [`libsy`], kept out of the core crate but compiled and +//! tested here so they stay current. Each is a worked example of the +//! [`Algorithm`](libsy::Algorithm) trait; the `examples/` directory has runnable agents +//! that drive them. +//! +//! - [`rand::RandomOrchAlgo`] — uniform random over the target set (one call). +//! - [`llm_class::LlmClassifierOrchAlgo`] — classify with one model, then route to a +//! strong/weak model (multi-step). +//! - [`ensemble::EnsembleOrchAlgo`] — fan out to several models, judge, and commit +//! (stateful). + +pub mod ensemble; +pub mod llm_class; +pub mod rand; diff --git a/crates/libsy-examples/src/llm_class.rs b/crates/libsy-examples/src/llm_class.rs new file mode 100644 index 000000000..94e252094 --- /dev/null +++ b/crates/libsy-examples/src/llm_class.rs @@ -0,0 +1,349 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! LLM-classifier router built on the [`Algorithm`] interfaces. +//! +//! Unlike a local ML classifier (which scores a prompt in-process), an LLM +//! classifier needs its own model call to classify the request. On the new +//! interfaces this is just two ordinary `driver.call_llm_target`s inside one +//! `create_run_task`: first the classifier target (to get a score), then the +//! routed strong/weak target. The multi-step nature is invisible to the caller — +//! it is the algorithm's own control flow. + +use std::error::Error; +use std::sync::Arc; + +use async_trait::async_trait; + +use libsy::{ + Algorithm, Context, Decision, Driver, LlmRequest, LlmTargetSet, Request, Response, Signals, +}; + +/// Preamble prepended to the user prompt when asking the classifier target 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 produced at each step of the classifier flow. The classify step +/// leaves `score`/`tier` `None`; the routed step fills them in. +pub struct ClassifierDecision { + /// The model this step selected (the classifier model, then the routed model). + pub selected_model: String, + /// Human-readable explanation of the step. + pub reasoning: String, + /// The classifier score, on the routed step; `None` on the classify step. + pub score: Option, + /// The tier chosen (strong/weak), on the routed step; `None` on the classify step. + pub tier: Option, +} + +impl Decision for ClassifierDecision { + fn selected_model(&self) -> &str { + &self.selected_model + } + fn reasoning(&self) -> Option<&str> { + Some(&self.reasoning) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// LLM-classifier router: classify with one target, then route to strong/weak. +pub struct LlmClassifierOrchAlgo { + classifier_model: String, + strong_model: String, + weak_model: String, + threshold: f64, + target_set: LlmTargetSet, +} + +impl LlmClassifierOrchAlgo { + /// Configure the classifier: the model that scores each request, the strong + /// and weak models to route to, the score `threshold` at or above which the + /// strong model is chosen, and the `target_set` to route among. That set must + /// contain targets named `classifier_model`, `strong_model`, and `weak_model`. + pub fn new( + classifier_model: impl Into, + strong_model: impl Into, + weak_model: impl Into, + threshold: f64, + target_set: LlmTargetSet, + ) -> Self { + Self { + classifier_model: classifier_model.into(), + strong_model: strong_model.into(), + weak_model: weak_model.into(), + threshold, + target_set, + } + } +} + +#[async_trait] +impl Algorithm for LlmClassifierOrchAlgo { + async fn create_run_task( + self: Arc, + _ctx: Context, + driver: Driver, + request: Request, + ) -> Result> { + let user_prompt = request.llm_request.prompt.clone(); + // The agent's inbound name rides through unchanged on every sub-call; the + // model each sub-call actually hits is carried by its decision instead. + let inbound = request.llm_request.inbound_model_name.clone(); + + // 1. Classify: call the classifier target with the score-eliciting prompt. + let classifier_target = self.target_set.get_target(&self.classifier_model)?; + let classify_request = Request { + llm_request: LlmRequest { + inbound_model_name: inbound.clone(), + prompt: format!("{CLASSIFIER_PROMPT_PREAMBLE}{user_prompt}"), + }, + raw_request: request.raw_request.clone(), + metadata: request.metadata.clone(), + }; + let classify_decision: Arc = Arc::new(ClassifierDecision { + selected_model: self.classifier_model.clone(), + reasoning: format!("classifying request via {}", self.classifier_model), + score: None, + tier: None, + }); + driver.info(classify_decision.clone()).await?; + let classify_response = driver + .call_llm_target(&classifier_target, classify_request, classify_decision) + .await?; + let score = classify_response + .llm_response + .completion + .trim() + .parse::() + .ok(); + + // 2. Route: pick strong/weak. Fail open — an unparseable score routes strong. + let (tier, model) = match score { + Some(s) if s >= self.threshold => (ClassifierTier::Strong, self.strong_model.clone()), + Some(_) => (ClassifierTier::Weak, self.weak_model.clone()), + None => (ClassifierTier::Strong, self.strong_model.clone()), + }; + let routed_target = self.target_set.get_target(&model)?; + let route_decision: Arc = Arc::new(ClassifierDecision { + reasoning: format!( + "classifier score {score:?} vs threshold {}; selected {model} ({})", + self.threshold, + tier.as_str() + ), + selected_model: model.clone(), + score, + tier: Some(tier), + }); + let routed_request = Request { + llm_request: LlmRequest { + inbound_model_name: inbound, + prompt: user_prompt, + }, + raw_request: request.raw_request, + metadata: request.metadata, + }; + driver.info(route_decision.clone()).await?; + driver + .call_llm_target(&routed_target, routed_request, route_decision) + .await + } + + async fn process_signals( + self: Arc, + _signals: Signals, + ) -> Result<(), Box> { + // Stateless classification; agent-system signals are ignored. + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use libsy::{LlmClient, LlmRequest, LlmResponse, LlmTarget, Response, RoutedRequest}; + use std::sync::Mutex; + + /// Returns `score` for the classifier target, an answer tagged with the model + /// otherwise; records the requests it saw so a test can inspect the classifier + /// prompt. + struct ScoringClient { + classifier_model: String, + score: String, + seen: Arc>>, + } + + #[async_trait] + impl LlmClient for ScoringClient { + async fn call( + &self, + routed: RoutedRequest, + ) -> Result> { + let name = routed.decision.selected_model().to_string(); + let completion = if name == self.classifier_model { + self.score.clone() + } else { + format!("answer from {name}") + }; + self.seen + .lock() + .map_err(|_| "lock poisoned")? + .push(routed.request); + Ok(Response { + llm_response: LlmResponse { + completion, + raw_response: None, + }, + metadata: None, + }) + } + } + + /// Build a classifier algo whose three targets share a scoring client. + fn algo(threshold: f64, score: &str) -> (LlmClassifierOrchAlgo, Arc>>) { + let seen = Arc::new(Mutex::new(Vec::new())); + let client = Arc::new(ScoringClient { + classifier_model: "router/classifier".to_string(), + score: score.to_string(), + seen: Arc::clone(&seen), + }) as Arc; + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(client.clone()), + }; + let target_set = LlmTargetSet::new(vec![ + target("router/classifier"), + target("frontier/model"), + target("cheap/model"), + ]); + let algo = LlmClassifierOrchAlgo { + classifier_model: "router/classifier".to_string(), + strong_model: "frontier/model".to_string(), + weak_model: "cheap/model".to_string(), + threshold, + target_set, + }; + (algo, seen) + } + + fn request(prompt: &str) -> Request { + Request { + llm_request: LlmRequest { + inbound_model_name: "auto".to_string(), + prompt: prompt.to_string(), + }, + raw_request: None, + metadata: None, + } + } + + /// Wrap a classifier algo as `Arc` we can drive to completion. + fn orch(algo: LlmClassifierOrchAlgo) -> Arc { + Arc::new(algo) + } + + /// Downcast a trace entry to the concrete classifier decision. + fn as_classifier( + d: &Arc, + ) -> Result<&ClassifierDecision, Box> { + d.as_any() + .downcast_ref::() + .ok_or_else(|| "expected a ClassifierDecision".into()) + } + + #[tokio::test] + async fn score_at_or_above_threshold_routes_strong() -> Result<(), Box> + { + let (algo, _) = algo(0.5, "0.9"); + let (trace, response) = orch(algo) + .run(Context::default(), request("solve this proof")) + .await?; + assert_eq!( + response.llm_response.completion, + "answer from frontier/model" + ); + // Trace: [classify, route]. + assert_eq!(trace[0].selected_model(), "router/classifier"); + let routed = as_classifier(&trace[1])?; + assert_eq!(routed.selected_model, "frontier/model"); + assert_eq!(routed.tier, Some(ClassifierTier::Strong)); + assert_eq!(routed.score, Some(0.9)); + Ok(()) + } + + #[tokio::test] + async fn score_below_threshold_routes_weak() -> Result<(), Box> { + let (algo, _) = algo(0.5, "0.2"); + let (trace, response) = orch(algo) + .run(Context::default(), request("say hello")) + .await?; + assert_eq!(response.llm_response.completion, "answer from cheap/model"); + let routed = as_classifier(&trace[1])?; + assert_eq!(routed.tier, Some(ClassifierTier::Weak)); + assert_eq!(routed.score, Some(0.2)); + Ok(()) + } + + #[tokio::test] + async fn score_exactly_at_threshold_routes_strong() -> Result<(), Box> + { + let (algo, _) = algo(0.5, "0.5"); + let (_, response) = orch(algo) + .run(Context::default(), request("borderline")) + .await?; + assert_eq!( + response.llm_response.completion, + "answer from frontier/model" + ); + Ok(()) + } + + #[tokio::test] + async fn unparseable_score_defaults_to_strong() -> Result<(), Box> { + let (algo, _) = algo(0.5, "not-a-number"); + let (trace, response) = orch(algo).run(Context::default(), request("hi")).await?; + assert_eq!( + response.llm_response.completion, + "answer from frontier/model" + ); + let routed = as_classifier(&trace[1])?; + assert_eq!(routed.tier, Some(ClassifierTier::Strong)); + assert_eq!(routed.score, None); + Ok(()) + } + + #[tokio::test] + async fn classifier_prompt_includes_the_user_text() -> Result<(), Box> + { + let (algo, seen) = algo(0.5, "0.9"); + orch(algo) + .run(Context::default(), request("prove it")) + .await?; + let seen = seen.lock().map_err(|_| "lock poisoned")?; + // Two calls: the classifier (preamble + user text), then the routed model. + assert_eq!(seen.len(), 2); + assert!(seen[0].llm_request.prompt.contains("prove it")); + assert!(seen[0].llm_request.prompt.contains("frontier model")); + assert_eq!(seen[1].llm_request.prompt, "prove it"); + Ok(()) + } +} diff --git a/crates/libsy-examples/src/rand.rs b/crates/libsy-examples/src/rand.rs new file mode 100644 index 000000000..9e78f72f8 --- /dev/null +++ b/crates/libsy-examples/src/rand.rs @@ -0,0 +1,232 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Random router built on the [`Algorithm`] interfaces. +//! +//! Selects one target from the set uniformly at random and calls it. This is the +//! simplest possible routing algorithm and the reference for the single-call +//! shape: one `driver.call_llm_target` inside `create_run_task`. (Weighted selection could +//! be layered on later; the set defines the candidates.) + +use std::error::Error; +use std::sync::Arc; + +use async_trait::async_trait; +use rand::seq::SliceRandom; + +use libsy::{Algorithm, Context, Decision, Driver, LlmTargetSet, Request, Response, Signals}; + +/// Decision produced by [`RandomOrchAlgo`]: which target was chosen and why. +pub struct RandomDecision { + /// The randomly selected target/model. + pub selected_model: String, + /// Human-readable explanation of the choice. + pub reasoning: String, +} + +impl Decision for RandomDecision { + fn selected_model(&self) -> &str { + &self.selected_model + } + fn reasoning(&self) -> Option<&str> { + Some(&self.reasoning) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// Uniform random router over a target set. +pub struct RandomOrchAlgo { + target_set: LlmTargetSet, +} + +impl RandomOrchAlgo { + /// Create a router over `target_set`. Wrap it in an + /// [`Arc`](std::sync::Arc) and drive it with + /// [`run`](libsy::Algorithm::run) or + /// [`run_stream`](libsy::Algorithm::run_stream). + pub fn new(target_set: LlmTargetSet) -> Self { + Self { target_set } + } +} + +#[async_trait] +impl Algorithm for RandomOrchAlgo { + async fn create_run_task( + self: Arc, + _ctx: Context, + driver: Driver, + request: Request, + ) -> Result> { + // Select a target uniformly at random. Scope the RNG so the non-Send + // `ThreadRng` is dropped before the await below, keeping the returned + // future `Send` (required by the `Algorithm` bound). + let target = { + let mut rng = rand::thread_rng(); + self.target_set + .targets() + .choose(&mut rng) + .ok_or("no targets available")? + .clone() + }; + + // Route by target semantic name; the caller's client (or offload host) maps + // it to the provider model id when it serves or offloads the call. + let selected = target.semantic_name.clone(); + let decision: Arc = Arc::new(RandomDecision { + reasoning: format!("random routing selected target '{selected}'"), + selected_model: selected, + }); + + // Publish the decision to the stream, then offload the call. + driver.info(decision.clone()).await?; + driver.call_llm_target(&target, request, decision).await + } + + async fn process_signals( + self: Arc, + _signals: Signals, + ) -> Result<(), Box> { + // Random routing is stateless, so agent-system signals are ignored. + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use libsy::{LlmClient, LlmRequest, LlmResponse, LlmTarget, Response, RoutedRequest}; + use std::collections::HashSet; + + /// Echoes back the target name it was called with, so a test can tell which + /// target the algo selected. + struct EchoClient; + + #[async_trait] + impl LlmClient for EchoClient { + async fn call( + &self, + routed: RoutedRequest, + ) -> Result> { + Ok(Response { + llm_response: LlmResponse { + completion: routed.decision.selected_model().to_string(), + raw_response: None, + }, + metadata: None, + }) + } + } + + fn request() -> Request { + Request { + llm_request: LlmRequest { + inbound_model_name: "auto".to_string(), + prompt: "hi".to_string(), + }, + raw_request: None, + metadata: None, + } + } + + /// Build a random-routing algorithm over `names`; every target echoes its name. + fn orch(names: &[&str]) -> Arc { + Arc::new(algo(names)) + } + + fn algo(names: &[&str]) -> RandomOrchAlgo { + let targets: Vec = names + .iter() + .map(|name| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(Arc::new(EchoClient)), + }) + .collect(); + RandomOrchAlgo::new(LlmTargetSet::new(targets)) + } + + #[tokio::test] + async fn single_target_is_always_selected_and_called( + ) -> Result<(), Box> { + let orch = orch(&["only/model"]); + let (trace, response) = orch.clone().run(Context::default(), request()).await?; + assert_eq!(response.llm_response.completion, "only/model"); + assert_eq!(trace.len(), 1); + assert_eq!(trace[0].selected_model(), "only/model"); + Ok(()) + } + + #[tokio::test] + async fn selected_target_is_in_the_set_and_matches_the_trace( + ) -> Result<(), Box> { + let names = ["a/model", "b/model", "c/model"]; + let orch = orch(&names); + for _ in 0..50 { + let (trace, response) = orch.clone().run(Context::default(), request()).await?; + let selected = response.llm_response.completion.clone(); + assert!( + names.contains(&selected.as_str()), + "selected {selected} not in target set" + ); + // The trace records the same target that was actually called. + assert_eq!(trace[0].selected_model(), selected.as_str()); + } + Ok(()) + } + + #[tokio::test] + async fn selection_covers_all_targets_over_many_runs( + ) -> Result<(), Box> { + let orch = orch(&["a/model", "b/model"]); + let mut seen = HashSet::new(); + for _ in 0..100 { + let (_, response) = orch.clone().run(Context::default(), request()).await?; + seen.insert(response.llm_response.completion); + } + // 100 uniform draws over two targets: both should appear (miss ~ 2^-99). + assert_eq!( + seen.len(), + 2, + "expected both targets to be selected, saw {seen:?}" + ); + Ok(()) + } + + #[tokio::test] + async fn empty_target_set_errors() { + let orch = orch(&[]); + assert!(orch + .clone() + .run(Context::default(), request()) + .await + .is_err()); + } + + #[tokio::test] + async fn process_signals_is_a_noop() -> Result<(), Box> { + let algo = algo(&["only/model"]); + Arc::new(algo).process_signals(Signals {}).await?; + Ok(()) + } + + #[tokio::test] + async fn decision_is_inspectable_and_downcasts() -> Result<(), Box> { + let orch = orch(&["only/model"]); + let (trace, _) = orch.clone().run(Context::default(), request()).await?; + let decision = &trace[0]; + // Uniform, algo-agnostic access via the trait — no concrete type needed. + assert_eq!(decision.selected_model(), "only/model"); + assert!(decision + .reasoning() + .unwrap_or_default() + .contains("only/model")); + // Escape hatch: downcast to the concrete decision when the algo is known. + let concrete = decision + .as_any() + .downcast_ref::() + .ok_or("expected a RandomDecision")?; + assert_eq!(concrete.selected_model, "only/model"); + Ok(()) + } +} diff --git a/crates/libsy/Cargo.toml b/crates/libsy/Cargo.toml new file mode 100644 index 000000000..dfd6bda0c --- /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" +serde_json = "1" +futures = "0.3" +tokio = { version = "1", features = ["full"] } +tokio-stream = "0.1" + +[dev-dependencies] +tokio = { version = "1", features = ["macros", "rt"] } diff --git a/crates/libsy/README.md b/crates/libsy/README.md new file mode 100644 index 000000000..652c1b43a --- /dev/null +++ b/crates/libsy/README.md @@ -0,0 +1,196 @@ +# libsy — Switchyard-Lib + +A lightweight, provider-agnostic library for multi-LLM agent optimization, with +**routing** as the first case. libsy *decides* how to serve each request — which +model(s) to call, in what order, how to combine them — and either makes the calls +itself or hands them back for you to make. It owns no HTTP client and no provider SDK, +so it drops into a proxy, gateway, or agent runtime. + +## Example + +Build a target set, pick an algorithm, run a request: + +```rust +use libsy::llm_class::LlmClassifierOrchAlgo; +use libsy::{Algorithm, Context, LlmClient, LlmRequest, LlmTarget, LlmTargetSet, Request}; +use std::sync::Arc; + +// Targets the algorithm routes among, each backed by your LlmClient (see below). +let client = Arc::new(MyClient { /* .. */ }) as Arc; +let target = |name: &str| LlmTarget { semantic_name: name.into(), llm_client: Some(client.clone()) }; + +let algo: Arc = Arc::new(LlmClassifierOrchAlgo::new( + "classifier", "strong", "weak", 0.5, + LlmTargetSet::new(vec![target("classifier"), target("strong"), target("weak")]), +)); + +let req = Request { + llm_request: LlmRequest { inbound_model_name: "auto".into(), prompt: "explain tail latency".into() }, + raw_request: None, + metadata: None, +}; +let (trace, response) = algo.clone().run(Context::default(), req).await?; // calls in, trace + response out +println!("routed to {}", trace.last().unwrap().selected_model()); +``` + +Runnable: [`research_agent`](../libsy-examples/examples/research_agent.rs) (in the `libsy-examples` crate). + +## Requests & responses + +```rust +pub struct Request { + pub llm_request: LlmRequest, // normalized: { inbound_model_name, prompt } + pub raw_request: Option, // original provider body, forwarded verbatim if present + pub metadata: Option, // correlation: session / agent / task / correlation_id / extra +} + +pub struct Response { + pub llm_response: LlmResponse, // normalized: { completion, raw_response? } + pub metadata: Option, +} +``` + +## Targets and clients + +An `LlmTarget` pairs a routing `semantic_name` with an optional `LlmClient`. Mapping that +name to a provider model id is the client's job, not the algorithm's — `LlmClient` is +meant to be implemented by you. + +```rust +struct MyClient { /* http client, base url, key */ } + +#[async_trait::async_trait] +impl LlmClient for MyClient { + async fn call(&self, routed: RoutedRequest) + -> Result> { + let model = routed.decision.selected_model(); // the routed target — map it to a provider id + // routed.request.llm_request.inbound_model_name is the agent's original name (not a call target) + // ... POST to your endpoint, read the completion ... + Ok(Response { llm_response: LlmResponse { completion, raw_response: None }, metadata: None }) + } +} +``` + +A `RoutedRequest` bundles the `request` with the routing `decision` and the target's +`default_client`; the model to call is `decision.selected_model()`, never a mutated +request field. `semantic_name` is the label an algorithm routes by; the client maps it to +the id it calls — they can differ (`"strong"` → `"openai/gpt-4o"`) or coincide. + +## Running a request + +Hold the algorithm as `Arc` and choose one of two entry points: + +```rust +// run: libsy drives the request to completion, serving each call with the target's +// client, and returns (trace, response). Errors if a routed target has no client. +let (trace, response) = algo.clone().run(Context::default(), req).await?; + +// run_stream: "ask, don't call" — you drive the stream and make the calls. +let stream = algo.clone().run_stream(Context::default(), req); +``` + +Under the hood every model call is *offloaded* to the request's `Step` stream; `run` is +the convenience that serves each one via the target's client. The step stream is bounded, +so pulling it paces the algorithm; each run is independent, so many run concurrently. + +## Streaming — you own the model calls (`run_stream`) + +`run_stream` yields `CallLlm` promises you fulfill with your own transport (or the call's +`default_client`), streams each `Decision` as it happens, and ends with `ReturnToAgent`. +Runnable: [`research_agent_core`](../libsy-examples/examples/research_agent_core.rs) (in the `libsy-examples` crate). + +```rust +let stream = algo.clone().run_stream(Context::default(), req); +tokio::pin!(stream); +while let Some(step) = stream.next().await { + match step? { + Step::CallLlm(call) => { + let routed = call.get_routed()?.clone(); // which target, and its default client + let response = call_model(routed.decision.selected_model(), &routed.request).await; // your real call + call.respond(Ok(response))?; // or Err(..) to propagate a failure + } + Step::Decision(decision) => { /* decision.selected_model(), decision.reasoning() */ } + Step::ReturnToAgent(response) => { /* done */ } + } +} +``` + +## Building an algorithm (`Algorithm`) + +Implement `Algorithm` to add a strategy. You write `create_run_task` — one call per +request; `run` / `run_stream` are provided and drive it. Make model calls on the `Driver` +you're handed, and publish a `Decision` for each so consumers (and clients) see *which* +model and *why*. + +```rust +#[async_trait] +pub trait Algorithm: Send + Sync + 'static { + // `self: Arc` (not `&mut`): one algorithm serves requests concurrently — use + // interior mutability for state. Offload calls/decisions on `driver`. + async fn create_run_task(self: Arc, ctx: Context, driver: Driver, request: Request) + -> Result>; + async fn process_signals(self: Arc, signals: Signals) + -> Result<(), Box>; + // provided: run(ctx, request) -> (trace, response), run_stream(ctx, request) -> Stream +} + +pub trait Decision: Send + Sync { + fn selected_model(&self) -> &str; // the model chosen — the client's call target + fn reasoning(&self) -> Option<&str>; // human-readable "why" + fn as_any(&self) -> &dyn std::any::Any; // downcast to the concrete decision +} +``` + +Give it a `new(config.., target_set)` constructor and `Arc`-wrap it — there is no builder. +Example — the LLM classifier (classify, then route; full version in +[`libsy-examples/src/llm_class.rs`](../libsy-examples/src/llm_class.rs)): + +```rust +#[async_trait] +impl Algorithm for LlmClassifierOrchAlgo { + async fn create_run_task(self: Arc, _ctx: Context, driver: Driver, request: Request) + -> Result> { + // 1. Classify: ask the classifier target for a score. + let classifier = self.target_set.get_target(&self.classifier_model)?; + driver.info(classify_decision.clone()).await?; + let score = driver.call_llm_target(&classifier, classify_req, classify_decision).await? + .llm_response.completion.trim().parse::().ok(); + + // 2. Route: strong if score >= threshold, else weak (fail open on None). + let model = if score.map_or(true, |s| s >= self.threshold) { &self.strong_model } else { &self.weak_model }; + let routed = self.target_set.get_target(model)?; + driver.info(route_decision.clone()).await?; + driver.call_llm_target(&routed, routed_req, route_decision).await + } + + async fn process_signals(self: Arc, _s: Signals) -> Result<(), Box> { Ok(()) } +} +``` + +## Explore + +Reference algorithms and runnable agents live in the sibling +[`libsy-examples`](../libsy-examples) crate (kept out of `libsy` itself, but compiled and +tested — `cargo test -p libsy-examples`). + +**Reference algorithms** — implementations to read and route with: + +- [`RandomOrchAlgo`](../libsy-examples/src/rand.rs) — uniform random over the set (one call). +- [`LlmClassifierOrchAlgo`](../libsy-examples/src/llm_class.rs) — classify, then route + strong/weak; fail open to strong. +- [`EnsembleOrchAlgo`](../libsy-examples/src/ensemble.rs) — stateful: fan out to + candidates, judge the best, commit to the winner after N exploration turns. + +**Runnable agents** (`cargo run -p libsy-examples --example `): + +- [`research_agent`](../libsy-examples/examples/research_agent.rs) — client-backed + targets, `run` (libsy makes the calls). +- [`research_agent_core`](../libsy-examples/examples/research_agent_core.rs) — client-less + targets, `run_stream` (the agent makes the calls). + +## Not yet built + +- **`Signals` events** — `process_signals` / `Signals` exist but carry nothing yet. +- **`Context` fields** — the per-request state carrier is an empty placeholder today. +- **Observability** — spans + a metrics sink (`Decision` is the hook). +- **Config-driven construction**, **typed errors** (vs `Box`), **weighted random**. diff --git a/crates/libsy/src/driver.rs b/crates/libsy/src/driver.rs new file mode 100644 index 000000000..410f713d9 --- /dev/null +++ b/crates/libsy/src/driver.rs @@ -0,0 +1,507 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! # driver — a type-erased promise-over-a-stream request pump +//! +//! [`TypeErasedDriver`] is the generic offload primitive; the crate root builds the +//! libsy-typed `Driver` on top of it. This module has no dependency on the rest of the +//! crate — the coupling is one-directional (`lib.rs` → `driver.rs`). +//! +//! A [`TypeErasedDriver`] lets a *producer* (e.g. a routing algorithm) fulfill +//! arbitrary requests by publishing promises onto a stream that a single *consumer* +//! drains. It is the type-erased generalization of the crate's `run_stream` offload: +//! instead of one fixed request/response shape, a producer calls +//! [`fulfill_request`](TypeErasedDriver::fulfill_request) +//! with *any* `REQ` and awaits *any* `RES`. +//! +//! - [`fulfill_request`](TypeErasedDriver::fulfill_request) enqueues a [`DriverStep::Request`] +//! carrying a [`DriverRequest`], then awaits the consumer's response. +//! - [`info`](TypeErasedDriver::info) pushes a fire-and-forget [`DriverStep::Info`] — no promise +//! to await. +//! - [`done`](TypeErasedDriver::done) emits the terminal [`DriverStep::Done`] with a final payload. +//! - [`stream`](TypeErasedDriver::stream) hands the single consumer the [`Stream`] of steps; for +//! each [`DriverStep::Request`] the consumer downcasts the request, computes a +//! response, and writes it back with [`DriverRequest::respond`]. +//! +//! Payloads are erased to `Box`, so one `TypeErasedDriver` serves any request +//! type; the consumer downcasts to the concrete type it expects. `TypeErasedDriver` is `Clone` +//! (many producer tasks may call it concurrently — multi-producer), while the stream +//! is single-consumer. Each request rides its own `oneshot`, so concurrent +//! `fulfill_request` calls never cross responses. +//! +//! ## Pacing (bounded step channel) +//! +//! The step channel has capacity 1, so a producer cannot publish its next step until +//! the consumer has pulled the previous one. The consumer therefore *paces* the +//! algorithm: it advances one step for each `.next().await`. Every producer method is +//! `async` because publishing a step awaits channel capacity. +//! +//! ## Termination +//! +//! There is no explicit stop method — the consumer terminates by **dropping the +//! stream** (and any [`DriverRequest`] it is holding). The producer's next publish +//! (`fulfill_request`/`info`/`done`/`fail`) then resolves to `Err`, and a producer +//! awaiting a response sees `Err` once the promise it handed out is dropped. Either +//! way the algorithm unwinds cooperatively at its next driver interaction. Because the +//! producer runs on a task the driver does not own, hard cancellation (e.g. mid-compute +//! that never touches the driver) is the caller's concern — abort the producer task. + +use std::{ + any::Any, + error::Error, + sync::{Arc, Mutex}, +}; + +use futures::{Stream, StreamExt}; +use tokio::sync::{mpsc, oneshot}; +use tokio_stream::wrappers::ReceiverStream; + +type BoxErr = Box; +type BoxAny = Box; +type StepResult = Result; + +/// One item on the stream returned by [`TypeErasedDriver::stream`]. +pub enum DriverStep { + /// A request awaiting a response. The consumer downcasts it and fulfills the + /// paired promise with [`DriverRequest::respond`]. + Request(DriverRequest), + /// A fire-and-forget payload from a producer; no response is expected. + Info(BoxAny), + /// A producer's terminal result. The consumer treats it as the last meaningful + /// step (the stream itself closes when every [`TypeErasedDriver`] clone drops). + Done(BoxAny), +} + +/// The consumer-facing half of one [`TypeErasedDriver::fulfill_request`] call. +/// +/// Yielded inside [`DriverStep::Request`]. The consumer reads the request via +/// [`request`](Self::request), does whatever work it names, and fulfills the promise +/// with [`respond`](Self::respond) — unblocking the producer's `fulfill_request`. +pub struct DriverRequest { + request: BoxAny, + // Fulfilled exactly once — `respond` consumes `self` to send, so no `Option`. + tx: oneshot::Sender>, +} + +impl DriverRequest { + /// Borrow the request payload as `REQ`. Errors if the producer enqueued a + /// different type than the consumer expected. + pub fn request(&self) -> Result<&REQ, BoxErr> { + self.request + .downcast_ref::() + .ok_or_else(|| "driver: request type mismatch".into()) + } + + /// Fulfill the promise with a typed response, or an `Err` to propagate a failure + /// back to the producer. Consumes `self`: a promise is fulfilled exactly once. + pub fn respond(self, res: Result) -> Result<(), BoxErr> { + // Erase the response so the single stream item type can carry any RES; the + // producer downcasts it back in `fulfill_request`. + let boxed: Result = res.map(|r| Box::new(r) as BoxAny); + self.tx + .send(boxed) + .map_err(|_| "driver: response receiver dropped".into()) + } +} + +/// Internal shared state: the step channel plus the single, take-once receiver. +struct DriverInner { + // Multi-producer: cloned into every `TypeErasedDriver`, so many tasks can enqueue steps. + // Capacity 1: a producer blocks publishing its next step until the consumer pulls + // the previous one, so the consumer paces the algorithm. + step_tx: mpsc::Sender, + // Single-consumer: taken out (once) by `stream`. `None` after the first take. + step_rx: Mutex>>, +} + +/// A promise-over-a-stream request pump. See the [module docs](self) for the model. +/// +/// Cheap to clone (shares one `Arc`); clone it to hand a producer handle to another +/// task. The consumer calls [`stream`](Self::stream) exactly once to drain steps. +#[derive(Clone)] +pub struct TypeErasedDriver { + inner: Arc, +} + +impl TypeErasedDriver { + /// Build an empty driver with its step channel ready. Take the consumer stream + /// with [`stream`](Self::stream); enqueue work with the other methods. + pub fn new() -> Self { + let (step_tx, step_rx) = mpsc::channel(1); + TypeErasedDriver { + inner: Arc::new(DriverInner { + step_tx, + step_rx: Mutex::new(Some(step_rx)), + }), + } + } + + /// Enqueue `req` as a [`DriverStep::Request`], await the consumer's response, and + /// downcast it to `RES`. Errors if the stream is closed, the promise is dropped + /// unfulfilled, the consumer responded with `Err`, or the response was not a `RES`. + pub async fn fulfill_request(&self, req: REQ) -> Result + where + REQ: Any + Send + 'static, + RES: Any + Send + 'static, + { + let (tx, rx) = oneshot::channel::>(); + let promise = DriverRequest { + request: Box::new(req), + tx, + }; + self.inner + .step_tx + .send(Ok(DriverStep::Request(promise))) + .await + .map_err(|_| "driver: stream closed")?; + + // Outer error: the promise was dropped without a response. Inner error: the + // consumer fulfilled it with an explicit `Err` — propagate it as-is. + let response = match rx.await { + Ok(result) => result?, + Err(_) => return Err("driver: promise dropped without a response".into()), + }; + response + .downcast::() + .map(|boxed| *boxed) + .map_err(|_| "driver: response type mismatch".into()) + } + + /// Push a fire-and-forget [`DriverStep::Info`] payload; there is no promise to + /// await for a response. Awaits channel capacity (the consumer pacing the stream) + /// and errors only if the stream is closed. + pub async fn info(&self, info: INFO) -> Result<(), BoxErr> + where + INFO: Any + Send + 'static, + { + self.inner + .step_tx + .send(Ok(DriverStep::Info(Box::new(info)))) + .await + .map_err(|_| "driver: stream closed".into()) + } + + /// Emit the terminal [`DriverStep::Done`] with a final payload. Does not close the + /// stream (that happens when every `TypeErasedDriver` clone drops); the consumer treats it + /// as the last meaningful step. Awaits channel capacity and errors only if the + /// stream is closed. + pub async fn done(&self, payload: T) -> Result<(), BoxErr> + where + T: Any + Send + 'static, + { + self.inner + .step_tx + .send(Ok(DriverStep::Done(Box::new(payload)))) + .await + .map_err(|_| "driver: stream closed".into()) + } + + /// Terminate the stream with an error item — the producer-side way to surface a + /// failure to the consumer (mirrors how the crate's `run_stream` yields an `Err` + /// step). Awaits channel capacity and errors only if the stream is + /// already closed. + pub async fn fail(&self, err: BoxErr) -> Result<(), BoxErr> { + self.inner + .step_tx + .send(Err(err)) + .await + .map_err(|_| "driver: stream closed".into()) + } + + /// Take the single consumer stream of [`DriverStep`]s. Callable once: a second + /// call yields a one-item stream carrying an `Err`, since the receiver is gone. + pub fn stream(&self) -> impl Stream> { + // Take the receiver out; `None` means already taken (or the lock was poisoned). + let taken = match self.inner.step_rx.lock() { + Ok(mut guard) => guard.take(), + Err(_) => None, + }; + match taken { + Some(rx) => ReceiverStream::new(rx).left_stream(), + None => futures::stream::once(async { + Err::("driver: stream already taken".into()) + }) + .right_stream(), + } + } +} + +impl Default for TypeErasedDriver { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + + #[tokio::test] + async fn fulfill_request_round_trips_typed_values() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + + // Producer asks for a u32 -> String on its own task. + let producer = driver.clone(); + let handle = + tokio::spawn(async move { producer.fulfill_request::(7u32).await }); + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + DriverStep::Request(promise) => { + let req = *promise.request::()?; + assert_eq!(req, 7); + promise.respond::(Ok(format!("got {req}")))?; + } + _ => return Err("expected a Request step".into()), + } + + assert_eq!(handle.await??, "got 7"); + Ok(()) + } + + #[tokio::test] + async fn info_pushes_a_typed_payload() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + driver.info(42u64).await?; + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + DriverStep::Info(payload) => { + let value = payload.downcast::().map_err(|_| "wrong info type")?; + assert_eq!(*value, 42); + } + _ => return Err("expected an Info step".into()), + } + Ok(()) + } + + #[tokio::test] + async fn done_emits_the_terminal_payload() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + driver.done("finished".to_string()).await?; + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + DriverStep::Done(payload) => { + let value = payload + .downcast::() + .map_err(|_| "wrong done type")?; + assert_eq!(*value, "finished"); + } + _ => return Err("expected a Done step".into()), + } + Ok(()) + } + + #[tokio::test] + async fn respond_error_propagates_to_the_producer() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + + let producer = driver.clone(); + let handle = tokio::spawn(async move { producer.fulfill_request::(1u32).await }); + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + DriverStep::Request(promise) => { + promise.respond::(Err("upstream failed".into()))?; + } + _ => return Err("expected a Request step".into()), + } + + match handle.await? { + Ok(_) => Err("expected the error to propagate".into()), + Err(err) => { + assert!(err.to_string().contains("upstream failed")); + Ok(()) + } + } + } + + #[tokio::test] + async fn response_type_mismatch_errors() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + + // Producer expects a String back. + let producer = driver.clone(); + let handle = + tokio::spawn(async move { producer.fulfill_request::(1u32).await }); + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + DriverStep::Request(promise) => { + // But the consumer responds with a u32. + promise.respond::(Ok(99u32))?; + } + _ => return Err("expected a Request step".into()), + } + + match handle.await? { + Ok(_) => Err("expected a response type mismatch".into()), + Err(err) => { + assert!(err.to_string().contains("response type mismatch")); + Ok(()) + } + } + } + + #[tokio::test] + async fn request_downcast_to_wrong_type_errors() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + + let producer = driver.clone(); + let handle = tokio::spawn(async move { producer.fulfill_request::(5u32).await }); + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + DriverStep::Request(promise) => { + assert!(promise.request::().is_err()); + // Unblock the producer so its task can finish. + promise.respond::(Ok(5u32))?; + } + _ => return Err("expected a Request step".into()), + } + + assert_eq!(handle.await??, 5); + Ok(()) + } + + #[tokio::test] + async fn closed_stream_errors_on_send() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + // Drop the consumer stream (and its receiver) before producing anything. + drop(driver.stream()); + + assert!(driver.fulfill_request::(1u32).await.is_err()); + assert!(driver.info(1u32).await.is_err()); + assert!(driver.done(1u32).await.is_err()); + Ok(()) + } + + #[tokio::test] + async fn promise_dropped_without_response_errors() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + + let producer = driver.clone(); + let handle = tokio::spawn(async move { producer.fulfill_request::(1u32).await }); + + tokio::pin!(stream); + match stream.next().await.ok_or("no step")?? { + // Drop the promise without responding. + DriverStep::Request(_promise) => {} + _ => return Err("expected a Request step".into()), + } + + assert!(handle.await?.is_err()); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_producers_do_not_cross_responses() -> Result<(), BoxErr> { + const N: usize = 8; + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + + // N producers each fulfill their own request concurrently. + let mut handles = Vec::new(); + for i in 0..N { + let producer = driver.clone(); + handles.push(( + i, + tokio::spawn(async move { producer.fulfill_request::(i).await }), + )); + } + + // Single consumer responds `req * 10` to each request. + tokio::pin!(stream); + let mut served = 0; + while served < N { + match stream.next().await.ok_or("stream ended early")?? { + DriverStep::Request(promise) => { + let req = *promise.request::()?; + promise.respond::(Ok(req * 10))?; + served += 1; + } + _ => return Err("expected a Request step".into()), + } + } + + // Each producer must see exactly its own response, not another's. + for (i, handle) in handles { + assert_eq!(handle.await??, i * 10); + } + Ok(()) + } + + #[tokio::test] + async fn stream_taken_twice_yields_an_error_item() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let _first = driver.stream(); + let second = driver.stream(); + + tokio::pin!(second); + match second.next().await.ok_or("expected an item")? { + Err(err) => { + assert!(err.to_string().contains("already taken")); + Ok(()) + } + Ok(_) => Err("expected an error item".into()), + } + } + + #[tokio::test] + async fn fail_surfaces_an_error_item_on_the_stream() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + let stream = driver.stream(); + driver.fail("kaboom".into()).await?; + + tokio::pin!(stream); + match stream.next().await.ok_or("no item")? { + Err(err) => { + assert!(err.to_string().contains("kaboom")); + Ok(()) + } + Ok(_) => Err("expected an error item".into()), + } + } + + #[tokio::test] + async fn dropping_the_stream_terminates_the_producer() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + // Box::pin so the stream is owned here and `drop` actually drops the receiver. + // (`tokio::pin!` would rebind to a `Pin<&mut _>`, making `drop` a no-op.) + let mut stream = Box::pin(driver.stream()); + + // Producer publishes paced steps until the consumer goes away, then reports + // how many it managed to send. + let producer = driver.clone(); + let handle = tokio::spawn(async move { + let mut sent = 0usize; + while producer.info(sent).await.is_ok() { + sent += 1; + } + sent + }); + + // Pace two steps, then terminate by dropping the stream. + for _ in 0..2 { + match stream.next().await.ok_or("stream ended early")?? { + DriverStep::Info(_) => {} + _ => return Err("expected an Info step".into()), + } + } + drop(stream); + + // With the consumer gone, the producer's next publish errors and it stops. + let sent = handle.await?; + assert!(sent >= 2, "producer should have published the paced steps"); + Ok(()) + } +} diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs new file mode 100644 index 000000000..67b6536cc --- /dev/null +++ b/crates/libsy/src/lib.rs @@ -0,0 +1,784 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! # libsy — multi-LLM agent optimization (routing first) +//! +//! `libsy` decides, per request, *how* to serve an LLM call: which model(s) to +//! invoke, in what order, and how to combine the results. Routing is the first +//! and simplest case; the same interfaces also express classifier routing, +//! ensembles, cascades, and other optimizations. The library owns no HTTP client +//! and no provider SDK — it decides, and the host makes (or is asked to make) the +//! actual calls — so it embeds cleanly in a proxy, gateway, or agent runtime. +//! +//! ## The model +//! +//! - An [`Algorithm`] is the optimization *algorithm*. Its +//! [`create_run_task`](Algorithm::create_run_task) runs once per request +//! and makes as many model calls as it needs — via [`Driver::call_llm_target`], which look +//! like ordinary calls — publishes its [`Decision`]s with [`Driver::info`], and +//! returns the final [`Response`]. The provided +//! [`run_stream`](Algorithm::run_stream) drives that on its own task and hands +//! back a stream of [`Step`]s; [`run`](Algorithm::run) runs +//! it to completion with the targets' default clients. +//! - An [`LlmTarget`] names a routing target by its [`semantic_name`](LlmTarget::semantic_name). +//! Every call is *offloaded* to the request's stream as a [`Step::CallLlm`]; the +//! target's [`LlmClient`], if any, rides along as +//! [`RoutedRequest::default_client`] so the host can serve it by default or +//! override it (see below). +//! +//! ## Running a request +//! +//! Hold the algorithm as `Arc` and call one of two provided methods: +//! +//! - [`run`](Algorithm::run) — run to completion, serving each +//! offloaded call via its [`RoutedRequest::default_client`], and return the decision +//! trace plus the final [`Response`]. The simplest integration; use it when the +//! algorithm holds the model clients (it errors if a routed target has no client). +//! - [`run_stream`](Algorithm::run_stream) — return a stream of [`Step`]s. Each +//! model call is offloaded: the stream yields a [`Step::CallLlm`] carrying a promise; +//! the host performs the real model call (optionally via the promise's +//! `default_client`) and fulfills it with [`CallLlmRequest::respond`]. Decisions +//! arrive as [`Step::Decision`] as the algorithm makes them, and the run ends with a +//! [`Step::ReturnToAgent`] carrying the final response. The step stream is bounded, +//! so pulling it paces the algorithm one step at a time — an "ask, don't call" mode +//! that lets a host that owns its transport keep control of every call. +//! +//! ## Concurrency +//! +//! [`Algorithm::create_run_task`] takes `self: Arc`, so one shared +//! `Arc` (no lock) serves many requests in parallel. Each +//! [`run_stream`](Algorithm::run_stream) call builds its own [`Driver`], so +//! offloaded calls never cross between concurrent requests. An algorithm is +//! responsible for its own thread-safety — stateless (like the reference routers) or +//! interior mutability over just its own state. +//! +//! ## Reference algorithms +//! +//! Worked implementations — a random router, an LLM classifier, and a stateful +//! ensemble — plus runnable agents live in the `libsy-examples` crate. + +mod driver; + +use std::{error::Error, pin::Pin, sync::Arc}; + +use async_trait::async_trait; +use futures::{Stream, StreamExt}; + +use crate::driver::{DriverRequest, DriverStep, TypeErasedDriver}; + +/// Shorthand for the crate's boxed, thread-safe error type. +type BoxErr = Box; + +/// A boxed, `Send` stream of [`Step`]s — the output of +/// [`Algorithm::run_stream`]. Boxed so the trait method that produces it keeps +/// `Arc` object-safe. +pub type StepStream = Pin> + Send>>; + +/// Correlation and routing metadata attached to a request or response. +/// +/// All fields are optional; algorithms and observers use whichever are present +/// (e.g. to key per-session state or emit correlated telemetry). `extra_metadata` +/// is a free-form escape hatch for host-specific keys. +#[derive(Clone)] +pub struct Metadata { + /// Stable id for a multi-request session/conversation. + pub session_id: Option, + /// Id of the agent making the request. + pub agent_id: Option, + /// Id of the task the request belongs to. + pub task_id: Option, + /// External trace/request id for joining with the host's telemetry. + pub correlation_id: Option, + /// Arbitrary host-defined key/value metadata. + pub extra_metadata: Option>, +} + +/// The normalized model request an algorithm reasons over and hands to a target. +/// +/// Deliberately minimal: a target model name and the user prompt. The full +/// provider-shaped request (messages, params, tools) rides on +/// [`Request::raw_request`] when a host needs to forward it losslessly. +#[derive(Clone)] +pub struct LlmRequest { + /// The model to call. Algorithms rewrite this as they route. + pub inbound_model_name: String, + /// The user prompt an algorithm inspects (e.g. to classify) and sends. + pub prompt: String, +} + +/// A request entering the orchestrator: the normalized [`LlmRequest`] plus the +/// original provider payload and correlation [`Metadata`]. +#[derive(Clone)] +pub struct Request { + /// The normalized request an algorithm routes. + pub llm_request: LlmRequest, + /// The original provider-shaped request body, if the host wants to forward it + /// verbatim (e.g. a proxy preserving messages/params). libsy does not read it. + pub raw_request: Option, + /// Correlation metadata carried through the request. + pub metadata: Option, +} + +/// Agentic-stack events fed to an algorithm out of band via +/// [`Algorithm::process_signals`] (e.g. tool results, budget updates). +/// +/// A placeholder today; a stateful algorithm can begin consuming signals as the +/// enum grows without changing the orchestrator contract. +#[derive(Clone)] +pub struct Signals {} + +/// The neutral model response returned by a target. +#[derive(Clone)] +pub struct LlmResponse { + /// The model's completion text — what an algorithm inspects (e.g. a + /// classifier score) or returns. + pub completion: String, + /// Optional raw provider response body, so a host (e.g. a proxy) can forward + /// the upstream response losslessly instead of rebuilding it from `completion`. + pub raw_response: Option, +} + +/// A response leaving the orchestrator: the neutral [`LlmResponse`] plus optional +/// correlation [`Metadata`]. +#[derive(Clone)] +pub struct Response { + /// The neutral model response. + pub llm_response: LlmResponse, + /// Correlation metadata carried through the response. + pub metadata: Option, +} + +/// A decision/trace object produced by an algorithm. +/// +/// Carried as a trait object (not a generic parameter) so a stream consumer can +/// inspect any algorithm's decision through this common interface without +/// knowing the concrete type. `as_any` is the escape hatch for a consumer that +/// *does* know the algo and wants to downcast to the concrete decision. +pub trait Decision: Send + Sync { + /// The model this decision selected (e.g. the routed target's name). + fn selected_model(&self) -> &str; + /// A human-readable explanation of the decision, for logs and traces. + fn reasoning(&self) -> Option<&str>; + /// Downcast handle: a consumer that knows the algorithm can recover the + /// concrete decision type via `as_any().downcast_ref::()`. + fn as_any(&self) -> &dyn std::any::Any; +} + +/// A request paired with the routing [`Decision`] that produced it — the unit an +/// [`LlmClient`] (or an offload host) is handed to serve. +/// +/// The two model identifiers live in separate, unambiguous places: the model to +/// call is [`decision.selected_model()`](Decision::selected_model), while +/// `request.llm_request.inbound_model_name` is the *inbound* name the agent asked +/// for (libsy never overwrites it). A client maps `selected_model()` to the +/// provider model id it hits. +#[derive(Clone)] +pub struct RoutedRequest { + /// The request to serve; its `inbound_model_name` is the agent's original name. + pub request: Request, + /// The routing decision behind this call; `selected_model()` is the model to hit. + pub decision: Arc, + /// The client that serves this call by default, or `None` when the routed target + /// had no client. Rides along on the offloaded call so a host driving the stream + /// can serve it by default or override it with its own transport. + pub default_client: Option>, +} + +/// The host-facing half of an offloaded model call, surfaced inside [`Step::CallLlm`]. +/// +/// Wraps a [`DriverRequest`] whose payload is a [`RoutedRequest`]. The host reads the +/// routed request ([`get_routed`](Self::get_routed)) and the decision behind it +/// ([`get_decision`](Self::get_decision)), performs (or delegates) the model call, and +/// fulfills it with [`respond`](Self::respond) — unblocking the algorithm's +/// [`Driver::call_llm`] on the other side. +pub struct CallLlmRequest { + inner: DriverRequest, +} + +impl CallLlmRequest { + /// Wrap a driver request whose payload is a [`RoutedRequest`]. + fn new(inner: DriverRequest) -> Self { + Self { inner } + } + + /// The routed request the host should serve. Its + /// [`default_client`](RoutedRequest::default_client) serves the call by default, + /// and its `decision.selected_model()` names the model to hit. Errors if the + /// promise payload was not a [`RoutedRequest`]. + pub fn get_routed(&self) -> Result<&RoutedRequest, BoxErr> { + self.inner.request::() + } + + /// The model request to perform (the [`Request`] inside the routed request). + pub fn get_request(&self) -> Result<&Request, BoxErr> { + Ok(&self.get_routed()?.request) + } + + /// The decision that led to this call — its `selected_model()` is the model to hit. + pub fn get_decision(&self) -> Result<&dyn Decision, BoxErr> { + Ok(self.get_routed()?.decision.as_ref()) + } + + /// Fulfill the promise with the caller's model-call result. Pass `Err(..)` to + /// propagate a failed model call back to the algorithm. Consumes the promise: it + /// can only be fulfilled once. + pub fn respond(self, result: Result) -> Result<(), BoxErr> { + self.inner.respond::(result) + } +} + +/// The offload channel handed to an algorithm's +/// [`create_run_task`](Algorithm::create_run_task). The algorithm makes model calls +/// with [`call_llm_target`](Self::call_llm_target) (or [`call_llm`](Self::call_llm)) and +/// publishes its [`Decision`]s with [`info`](Self::info); each call is offloaded to the +/// request's [`Step`] stream and awaits the consumer's response. The step channel is +/// bounded, so the consumer paces the algorithm one step at a time. +#[derive(Clone)] +pub struct Driver { + driver: TypeErasedDriver, +} + +impl Driver { + /// Build an empty driver with its step channel ready. Created per call by + /// [`run_stream`](Algorithm::run_stream). + pub(crate) fn new() -> Self { + Self { + driver: TypeErasedDriver::new(), + } + } + + /// Offload a model call: publish `routed` as a [`Step::CallLlm`] and await the + /// consumer's [`Response`]. Errors if the stream is closed or the call failed. + pub async fn call_llm(&self, routed: RoutedRequest) -> Result { + self.driver + .fulfill_request::(routed) + .await + } + + /// Offload a call to `target`: pair `request` with `decision` and the target's + /// default client into a [`RoutedRequest`], then publish it (see + /// [`call_llm`](Self::call_llm)). The convenience most algorithms use; + /// `decision.selected_model()` names the model to hit, and `request`'s + /// `inbound_model_name` is left untouched. + pub async fn call_llm_target( + &self, + target: &LlmTarget, + request: Request, + decision: Arc, + ) -> Result { + self.call_llm(RoutedRequest { + request, + decision, + default_client: target.llm_client.clone(), + }) + .await + } + + /// Publish a routing [`Decision`] as a [`Step::Decision`] on the stream. + pub async fn info(&self, decision: Arc) -> Result<(), BoxErr> { + self.driver.info(decision).await + } + + /// Emit the terminal step: [`Step::ReturnToAgent`] on `Ok`, or an `Err` stream + /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream) + /// when the algorithm finishes. + pub(crate) async fn finish(&self, result: Result) -> Result<(), BoxErr> { + match result { + Ok(response) => self.driver.done(response).await, + Err(err) => self.driver.fail(err).await, + } + } + + /// Transform the raw driver stream into a stream of [`Step`]s. Internal: the + /// consumer stream is taken (once) by [`run_stream`](Algorithm::run_stream). A + /// payload that does not match the expected type for its step becomes an `Err` item. + pub(crate) fn stream(&self) -> impl Stream> { + self.driver.stream().map(|item| match item? { + DriverStep::Request(req) => Ok(Step::CallLlm(CallLlmRequest::new(req))), + DriverStep::Info(payload) => payload + .downcast::>() + .map(|decision| Step::Decision(*decision)) + .map_err(|_| "driver: info payload was not a Decision".into()), + DriverStep::Done(payload) => payload + .downcast::() + .map(|response| Step::ReturnToAgent(*response)) + .map_err(|_| "driver: done payload was not a Response".into()), + }) + } +} + +impl Default for Driver { + fn default() -> Self { + Self::new() + } +} + +/// Per-request state threaded to an algorithm alongside its [`Driver`]. A placeholder +/// for cross-cutting state (correlation ids, budgets, deadlines) an algorithm will +/// read; empty today. It does not carry the offload driver, so it is safe to share. +#[derive(Clone, Default)] +pub struct Context {} + +impl Context { + /// Build an empty context. + pub fn new() -> Self { + Self {} + } +} + +/// One item in the stream returned by [`Driver::stream`] / [`Algorithm::run_stream`]. +pub enum Step { + /// The algorithm needs this model call performed. The host serves it (optionally + /// via [`RoutedRequest::default_client`]) and fulfills it with + /// [`CallLlmRequest::respond`]. + CallLlm(CallLlmRequest), + /// A routing decision the algorithm made, published via [`Driver::info`] as it + /// happens (rather than collected into a trace returned at the end). + Decision(Arc), + /// The algorithm finished with its final response — the last step of a run. + ReturnToAgent(Response), +} + +/// Performs the actual model call for a target. This is the one piece of I/O +/// `libsy` does not own — a host implements it over its own transport (HTTP SDK, +/// in-process model, mock). It serves a call the stream consumer chose not to +/// override, reached as [`RoutedRequest::default_client`] (see [`Algorithm::run_stream`]). +#[async_trait] +pub trait LlmClient: Send + Sync { + /// Serve `routed`, returning the model's response. Call the model named by + /// [`routed.decision.selected_model()`](Decision::selected_model) — the target + /// the algorithm routed to — mapping it to whatever provider model id this + /// client hits. `routed.request.llm_request.inbound_model_name` is the agent's + /// original name, carried through for reference, not a call target. + async fn call(&self, request: RoutedRequest) -> Result>; +} + +/// A named routing target: a `semantic_name` an algorithm routes by, and an optional +/// [`LlmClient`] to serve its calls. An algorithm hands a target to +/// [`Driver::call_llm_target`]; the client rides along as +/// [`RoutedRequest::default_client`] for the stream consumer to serve or override. +#[derive(Clone)] +pub struct LlmTarget { + /// The routing label an algorithm selects this target by — a logical tier like + /// `"strong"`, or the model id when they coincide. Mapping it to a provider model + /// id is the client's concern, never the algorithm's. + pub semantic_name: String, + /// The client that serves this target's calls by default, or `None` (then the + /// stream consumer must serve them). + pub llm_client: Option>, +} + +/// The set of targets an algorithm may route among. An algorithm is constructed +/// with one and picks targets by position ([`targets`](Self::targets)) or by name +/// ([`get_target`](Self::get_target)). +#[derive(Clone)] +pub struct LlmTargetSet { + targets: Vec, +} + +impl LlmTargetSet { + /// Build a target set from a list of targets. + pub fn new(targets: Vec) -> Self { + Self { targets } + } + + /// All targets in the set — e.g. for an algorithm to select among. + pub fn targets(&self) -> &[LlmTarget] { + &self.targets + } + + /// Look up a target by name; errors if no target has that name. + pub fn get_target(&self, name: &str) -> Result> { + self.targets + .iter() + .find(|t| t.semantic_name == name) + .cloned() + .ok_or(format!("Target {} not found", name).into()) + } +} + +/// An optimization strategy. Implement [`create_run_task`](Self::create_run_task); +/// callers drive it with the provided [`run`](Self::run) (serve calls, get the answer) +/// or [`run_stream`](Self::run_stream) (drive the [`Step`] stream yourself). +/// +/// Methods take `self: Arc`: one algorithm (`Arc`) is shared across +/// requests and run concurrently, so it owns its thread-safety. Stateless algorithms +/// (the reference routers) get this for free; a stateful one uses interior mutability +/// over just its own state. +#[async_trait] +pub trait Algorithm: Send + Sync + 'static { + /// Run one request to completion: make model calls with [`Driver::call_llm_target`], + /// publish [`Decision`]s with [`Driver::info`], and return the final [`Response`]. + /// The method an algorithm implements; [`run`](Self::run) / [`run_stream`](Self::run_stream) + /// drive it. `ctx` carries cross-cutting request state (empty today). + async fn create_run_task( + self: Arc, + ctx: Context, + driver: Driver, + request: Request, + ) -> Result>; + + /// Feed the algorithm agentic-stack events (tool results, budgets, etc.). The + /// reference algorithms ignore signals; a stateful algorithm updates its own + /// (interior-mutable) state. Takes `self: Arc` like the other run methods. + async fn process_signals( + self: Arc, + signals: Signals, + ) -> Result<(), Box>; + + /// Run one request as a stream of [`Step`]s (provided). The algorithm runs on its + /// own task; drive the stream: serve each [`Step::CallLlm`] (via its + /// [`default_client`](RoutedRequest::default_client) or your own transport) and read + /// [`Step::Decision`]s until the final [`Step::ReturnToAgent`]. The step channel is + /// bounded, so pulling paces the algorithm; each call is independent, so many run + /// concurrently. + fn run_stream(self: Arc, ctx: Context, request: Request) -> StepStream { + // This call's own driver: take its consumer stream, hand a producer-side clone to + // the algorithm task, and keep one to emit the terminal step. The task blocks + // publishing a step until the consumer pulls the previous one. + let driver = Driver::new(); + let stream = driver.stream(); + tokio::spawn(async move { + let outcome = self.create_run_task(ctx, driver.clone(), request).await; + let _ = driver.finish(outcome).await; + }); + Box::pin(stream) + } + + /// Run one request to completion, serving each offloaded call with its + /// [`RoutedRequest::default_client`], and return the decision trace plus the final + /// [`Response`]. Provided: drives [`run_stream`](Self::run_stream) internally, + /// collecting each [`Step::Decision`]. Use it when the algorithm holds its own model + /// clients and the host wants the answer (and the decisions behind it); drive + /// [`run_stream`](Self::run_stream) instead to serve the calls yourself. Errors + /// if a routed target has no client to serve its call, or the algorithm fails. + async fn run( + self: Arc, + ctx: Context, + request: Request, + ) -> Result<(Vec>, Response), Box> { + let stream = self.run_stream(ctx, request); + tokio::pin!(stream); + let mut trace: Vec> = Vec::new(); + while let Some(item) = stream.next().await { + match item? { + Step::CallLlm(call) => { + // Serve the call with the target's default client, or error if the + // routed target had none. + let routed = call.get_routed()?.clone(); + let client = routed.default_client.clone().ok_or_else(|| { + format!( + "run: target '{}' has no client to serve the call", + routed.decision.selected_model() + ) + })?; + call.respond(client.call(routed).await)?; + } + Step::Decision(decision) => trace.push(decision), + // The terminal step: return as soon as the algorithm finishes, rather + // than draining the stream until it closes. + Step::ReturnToAgent(response) => return Ok((trace, response)), + } + } + Err("run: stream ended without a final response".into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + + /// Mock client that echoes back the target name it was called with. + struct EchoClient; + + #[async_trait] + impl LlmClient for EchoClient { + async fn call( + &self, + routed: RoutedRequest, + ) -> Result> { + // Echo back the model the algorithm routed to (the decision's selection). + Ok(Response { + llm_response: LlmResponse { + completion: routed.decision.selected_model().to_string(), + raw_response: None, + }, + metadata: None, + }) + } + } + + /// Trivial decision + algo used only to exercise the orchestrator: calls the + /// first target and returns its response with a one-item trace. + struct TestDecision { + model: String, + } + + impl Decision for TestDecision { + fn selected_model(&self) -> &str { + &self.model + } + fn reasoning(&self) -> Option<&str> { + None + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + struct TestAlgo { + target_set: LlmTargetSet, + } + + #[async_trait] + impl Algorithm for TestAlgo { + async fn create_run_task( + self: Arc, + _ctx: Context, + driver: Driver, + request: Request, + ) -> Result> { + let target = self + .target_set + .targets() + .first() + .ok_or("no targets")? + .clone(); + let decision: Arc = Arc::new(TestDecision { + model: target.semantic_name.clone(), + }); + driver.info(decision.clone()).await?; + driver.call_llm_target(&target, request, decision).await + } + + async fn process_signals( + self: Arc, + _signals: Signals, + ) -> Result<(), Box> { + Ok(()) + } + } + + /// Build a shared `TestAlgo` over the given target set. + fn orch(target_set: LlmTargetSet) -> Arc { + Arc::new(TestAlgo { target_set }) + } + + fn request() -> Request { + Request { + llm_request: LlmRequest { + inbound_model_name: "auto".to_string(), + prompt: "hi".to_string(), + }, + raw_request: None, + metadata: None, + } + } + + /// `(name, has_client)` — `has_client: false` builds a target with no default client. + fn target_set(names: &[(&str, bool)]) -> LlmTargetSet { + let targets = names + .iter() + .map(|(name, has_client)| LlmTarget { + semantic_name: name.to_string(), + llm_client: has_client.then(|| Arc::new(EchoClient) as Arc), + }) + .collect(); + LlmTargetSet::new(targets) + } + + #[tokio::test] + async fn run_offloads_via_promise_then_returns_to_agent( + ) -> Result<(), Box> { + // A client-less target -> its call is offloaded via a promise the + // orchestrator surfaces as a `CallLlm` step for us to fulfill. + let stream = + orch(target_set(&[("offload/model", false)])).run_stream(Context::default(), request()); + tokio::pin!(stream); + + let mut saw_call = false; + let mut final_completion = None; + while let Some(step) = stream.next().await { + match step? { + Step::CallLlm(call) => { + saw_call = true; + // The decision rode along with the promise. + assert_eq!(call.get_decision()?.selected_model(), "offload/model"); + // Fulfilling the promise is the "real" model call the caller makes. + call.respond(Ok(Response { + llm_response: LlmResponse { + completion: "fulfilled".to_string(), + raw_response: None, + }, + metadata: None, + }))?; + } + Step::Decision(decision) => { + assert_eq!(decision.selected_model(), "offload/model"); + } + Step::ReturnToAgent(response) => { + final_completion = Some(response.llm_response.completion); + } + } + } + + assert!(saw_call, "expected a CallLlm step before ReturnToAgent"); + assert_eq!( + final_completion.ok_or("no ReturnToAgent step")?, + "fulfilled" + ); + Ok(()) + } + + #[tokio::test] + async fn client_backed_target_offloads_with_a_default_client( + ) -> Result<(), Box> { + // Every call now offloads to the stream; a client-backed target rides its + // client along as `default_client` so the consumer can serve it by default. + let stream = + orch(target_set(&[("direct/model", true)])).run_stream(Context::default(), request()); + tokio::pin!(stream); + + let mut final_completion = None; + while let Some(step) = stream.next().await { + match step? { + Step::CallLlm(call) => { + let routed = call.get_routed()?.clone(); + let client = routed + .default_client + .clone() + .ok_or("expected a default client")?; + let result = client.call(routed).await; + call.respond(result)?; + } + Step::Decision(_) => {} + Step::ReturnToAgent(response) => { + final_completion = Some(response.llm_response.completion); + } + } + } + + // EchoClient echoes the model name back as the completion. + assert_eq!(final_completion.ok_or("no ReturnToAgent")?, "direct/model"); + Ok(()) + } + + #[tokio::test] + async fn run_returns_the_response_when_all_targets_have_clients( + ) -> Result<(), Box> { + // Every target has a client, so run serves every call via the + // default client and returns the trace + final response. + let (trace, response) = orch(target_set(&[("direct/model", true)])) + .run(Context::default(), request()) + .await?; + // TestAlgo calls the first target; EchoClient echoes its name. + assert_eq!(response.llm_response.completion, "direct/model"); + assert_eq!(trace[0].selected_model(), "direct/model"); + Ok(()) + } + + #[tokio::test] + async fn run_errors_when_a_target_lacks_a_client() -> Result<(), Box> { + // A client-less target has no default client to serve its offloaded call, so + // driving it to completion errors. + assert!(orch(target_set(&[("offload/model", false)])) + .run(Context::default(), request()) + .await + .is_err()); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn requests_are_processed_in_parallel() -> Result<(), Box> { + use std::time::Duration; + use tokio::sync::Barrier; + + const N: usize = 4; + + // A client that blocks until all N concurrent calls have arrived. If + // requests were serialized (one algorithm behind a `Mutex`), only one + // call could be in flight, the barrier would never reach N, and the test + // would time out. It passes only because the shared algorithm is driven + // concurrently across requests. + struct BarrierClient { + barrier: Arc, + } + + #[async_trait] + impl LlmClient for BarrierClient { + async fn call( + &self, + routed: RoutedRequest, + ) -> Result> { + self.barrier.wait().await; + Ok(Response { + llm_response: LlmResponse { + completion: routed.decision.selected_model().to_string(), + raw_response: None, + }, + metadata: None, + }) + } + } + + let barrier = Arc::new(Barrier::new(N)); + let targets = LlmTargetSet::new(vec![LlmTarget { + semantic_name: "m".to_string(), + llm_client: Some(Arc::new(BarrierClient { + barrier: barrier.clone(), + })), + }]); + // One shared algorithm driven by many concurrent requests. + let algo = orch(targets); + + let mut handles = Vec::new(); + for _ in 0..N { + let algo = algo.clone(); + handles.push(tokio::spawn(async move { + algo.run(Context::default(), request()) + .await + .map(|(_, response)| response.llm_response.completion) + })); + } + + for handle in handles { + // The timeout turns a serialization deadlock into a failure, not a hang. + let completion = tokio::time::timeout(Duration::from_secs(5), handle).await???; + assert_eq!(completion, "m"); + } + Ok(()) + } + + #[tokio::test] + async fn offload_error_propagates_back_to_the_algorithm( + ) -> Result<(), Box> { + // A client-less target offloads its call; we fulfill the promise with an + // Err, which must flow back through `call_llm_target` into the algorithm and + // out as an error step — not a response. + let stream = + orch(target_set(&[("offload/model", false)])).run_stream(Context::default(), request()); + tokio::pin!(stream); + + let mut saw_error = false; + while let Some(step) = stream.next().await { + match step { + Ok(Step::CallLlm(call)) => { + call.respond(Err("upstream model call failed".into()))?; + } + Ok(Step::Decision(_)) => {} + Ok(Step::ReturnToAgent(..)) => { + return Err("expected the offload error to propagate, got a response".into()); + } + Err(err) => { + // The algorithm's `call_llm_target` saw the error via the promise. + assert!(err.to_string().contains("upstream model call failed")); + saw_error = true; + } + } + } + + assert!(saw_error, "expected an error step"); + Ok(()) + } +}