diff --git a/crates/libsy/src/algorithms.rs b/crates/libsy/src/algorithms.rs index 6288e4aec..4f1522931 100644 --- a/crates/libsy/src/algorithms.rs +++ b/crates/libsy/src/algorithms.rs @@ -6,10 +6,12 @@ //! Reach for them by name — `use switchyard_libsy::algorithms::Random` — rather than through the //! per-algorithm submodules. +pub mod fall_through; pub mod llm_class; pub mod noop; pub mod rand; +pub use fall_through::{FallThrough, FallThroughDecision}; pub use llm_class::{ClassifierDecision, ClassifierTier, LlmClassifier}; pub use noop::{Noop, NoopDecision}; pub use rand::{Random, RandomDecision}; diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs new file mode 100644 index 000000000..b42de8912 --- /dev/null +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -0,0 +1,416 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Fall-through classifier routing: a stateful [`Algorithm`] that routes each turn +//! through a processor chain and a classifier cascade. +//! +//! Each turn: request-side [`Processor`]s fold facts into the session [`State`]; the +//! [`Classifier`] cascade is consulted in order and the first to score decides the target +//! (its `argmax`); the [`Decision`] is published and then replayed to the processors so +//! stateful ones (latch, affinity) can bind it. +//! +//! [`FallThrough`] itself is stateless — one shared instance serves every session. The +//! per-session [`State`] rides in the threaded [`Context`] the caller passes into +//! each turn, so routing can depend on accumulated history rather than the current request +//! alone. The state is held under a lock, so concurrent turns of the same session serialize +//! on it. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::core::{Classifier, Event, Processor, Score, SharedState}; +use crate::{Algorithm, Context, Decision, Driver, LlmTargetSet, Request, Response}; + +/// Boxed, thread-safe error type used across the SDK. +type BoxErr = Box; + +/// The decision a fall-through run produces: the selected model plus a human-readable reason. +pub struct FallThroughDecision { + model: String, + reason: String, +} + +impl Decision for FallThroughDecision { + fn selected_model(&self) -> &str { + &self.model + } + + fn reasoning(&self) -> Option<&str> { + Some(&self.reason) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// Processor chain → classifier cascade → routed model call. See the [module docs](self). +/// +/// Stateless: the per-session [`State`] lives in the threaded [`Context`], so one +/// shared instance serves every session (see the module docs). +pub struct FallThrough { + processors: Vec>, + classifiers: Vec>, + targets: LlmTargetSet, +} + +impl FallThrough { + /// Creates an empty router over `targets`; add components with the `with_*` builders. + pub fn new(targets: LlmTargetSet) -> Self { + Self { + processors: Vec::new(), + classifiers: Vec::new(), + targets, + } + } + + /// Appends a processor to the head-of-request chain. + pub fn with_processor(mut self, processor: Arc) -> Self { + self.processors.push(processor); + self + } + + /// Appends a classifier to the cascade. + pub fn with_classifier(mut self, classifier: Arc) -> Self { + self.classifiers.push(classifier); + self + } +} +#[async_trait] +impl Algorithm for FallThrough { + async fn create_run_task( + self: Arc, + ctx: Context, + driver: Driver, + request: Request, + ) -> Result { + // Hold the session state (carried in the threaded context) for the whole + // request→decision fold, so a turn's fact accumulation is atomic and concurrent + // turns of the same session serialize on it. + let mut state = ctx.state.lock().await; + + // 1. Processor chain accumulates request-side facts into the session State. + for processor in &self.processors { + processor + .process(&mut state, Event::Request(&request)) + .await?; + } + + // 2. Fall through the cascade: the first classifier to score decides (argmax). The + // per-request driver is offered to each — driver-backed classifiers use it. + let mut winner: Option = None; + for classifier in &self.classifiers { + let scores = classifier + .score(&mut state, &request, Some(&driver)) + .await?; + if let Some(score) = scores.argmax(false)? { + winner = Some(score); + break; + } + } + let Some(winner) = winner else { + return Err("fall-through: every classifier abstained".into()); + }; + + // 3. Resolve the target and publish the decision. + let target = self.targets.get_target(&winner.target)?; + let decision: Arc = Arc::new(FallThroughDecision { + model: winner.target.clone(), + reason: format!( + "fall-through selected {} (confidence {:.3})", + winner.target, winner.confidence + ), + }); + driver.info(ctx.without_state(), decision.clone()).await?; + + // 4. Replay the decision to the processors so stateful ones (latch / affinity) bind + // it into the session State. + for processor in &self.processors { + processor + .process(&mut state, Event::Decision(decision.as_ref())) + .await?; + } + + driver + .call_llm_target(ctx.without_state(), &target, request, decision) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{Classification, State}; + + use switchyard_protocol::{ + completion_text, text_response, LlmRequest, Message, Metadata, Role, + }; + + use crate::{LlmResponse, LlmTarget, RoutedLlmClient}; + + // --- fixtures ---------------------------------------------------------------------- + + /// A client that echoes the routed model name back as the completion. + struct EchoClient; + + #[async_trait] + impl RoutedLlmClient for EchoClient { + async fn call( + &self, + _ctx: Context, + _request: Request, + decision: Arc, + ) -> Result { + Ok(Response { + llm_response: LlmResponse::Agg(text_response( + None, + decision.selected_model().to_string(), + )), + metadata: None, + }) + } + } + + /// A target set whose targets all serve via [`EchoClient`]. + fn target_set(names: &[&str]) -> LlmTargetSet { + LlmTargetSet::new( + names + .iter() + .map(|name| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(Arc::new(EchoClient) as Arc), + }) + .collect(), + ) + } + + /// A classifier that emits fixed scores (empty = abstain). + struct FixedClassifier(Vec); + + #[async_trait] + impl Classifier for FixedClassifier { + async fn score( + &self, + _state: &mut State, + _request: &Request, + _driver: Option<&Driver>, + ) -> Result { + Ok(Classification::Scores( + self.0 + .iter() + .map(|s| Score { + confidence: s.confidence, + target: s.target.clone(), + }) + .collect(), + )) + } + } + + fn score(target: &str, confidence: f64) -> Score { + Score { + confidence, + target: target.to_string(), + } + } + + fn fixed(scores: Vec) -> Arc { + Arc::new(FixedClassifier(scores)) + } + + fn request() -> Request { + Request { + llm_request: LlmRequest { + model: Some("auto".to_string()), + messages: vec![Message::text(Role::User, "hi")], + ..LlmRequest::default() + }, + raw_request: None, + metadata: Some(Metadata { + session_id: Some("session-1".to_string()), + ..Metadata::default() + }), + } + } + + /// Drives a shared router through one turn on the given session context, returning the + /// completion text + trace. Reuse the same `ctx` across calls to model a session. + async fn run_turn( + router: &Arc, + ctx: Context, + ) -> Result<(String, Vec>), BoxErr> { + let (trace, response) = router.clone().run(ctx, request()).await?; + let text = response + .llm_response + .into_agg() + .await + .map(|agg| completion_text(&agg))?; + Ok((text, trace)) + } + + /// Drives a fresh router through one turn on a fresh session. + async fn run(router: FallThrough) -> Result<(String, Vec>), BoxErr> { + run_turn(&Arc::new(router), Context::default()).await + } + + // --- tests ------------------------------------------------------------------------- + + #[tokio::test] + async fn argmax_picks_the_highest_confidence_target() -> Result<(), BoxErr> { + let router = FallThrough::new(target_set(&["strong", "weak"])) + .with_classifier(fixed(vec![score("weak", 0.2), score("strong", 0.9)])); + let (model, trace) = run(router).await?; + assert_eq!(model, "strong"); + assert_eq!(trace.len(), 1); + assert_eq!(trace[0].selected_model(), "strong"); + Ok(()) + } + + #[tokio::test] + async fn falls_through_the_first_abstaining_classifier() -> Result<(), BoxErr> { + // First classifier abstains (empty); the second decides. + let router = FallThrough::new(target_set(&["strong", "weak"])) + .with_classifier(fixed(vec![])) + .with_classifier(fixed(vec![score("weak", 1.0)])); + let (model, _) = run(router).await?; + assert_eq!(model, "weak"); + Ok(()) + } + + #[tokio::test] + async fn first_deciding_classifier_wins_the_cascade() -> Result<(), BoxErr> { + // The first classifier decides; the second is never consulted. + let router = FallThrough::new(target_set(&["strong", "weak"])) + .with_classifier(fixed(vec![score("strong", 0.6)])) + .with_classifier(fixed(vec![score("weak", 1.0)])); + let (model, _) = run(router).await?; + assert_eq!(model, "strong"); + Ok(()) + } + + #[tokio::test] + async fn all_abstaining_is_an_error() -> Result<(), BoxErr> { + let router = + FallThrough::new(target_set(&["strong", "weak"])).with_classifier(fixed(vec![])); + assert!(run(router).await.is_err()); + Ok(()) + } + + #[tokio::test] + async fn classifiers_receive_the_per_request_driver() -> Result<(), BoxErr> { + // A classifier that only decides when handed a driver — proving the cascade offers + // the per-request driver to every classifier (driver-backed ones need it). + struct NeedsDriver; + + #[async_trait] + impl Classifier for NeedsDriver { + async fn score( + &self, + _state: &mut State, + _request: &Request, + driver: Option<&Driver>, + ) -> Result { + match driver { + Some(_) => Ok(Classification::Scores(vec![score("strong", 1.0)])), + None => Err("expected a driver".into()), + } + } + } + + let router = FallThrough::new(target_set(&["strong", "weak"])) + .with_classifier(Arc::new(NeedsDriver)); + let (model, _) = run(router).await?; + assert_eq!(model, "strong"); + Ok(()) + } + + #[tokio::test] + async fn processor_observes_request_and_decision() -> Result<(), BoxErr> { + use std::sync::Mutex; + + // Records which event kinds it saw, proving the request-then-decision replay. + struct RecordingProcessor(Arc>>); + + #[async_trait] + impl Processor for RecordingProcessor { + async fn process(&self, _state: &mut State, event: Event<'_>) -> Result<(), BoxErr> { + let kind = match event { + Event::Request(_) => "request", + Event::Decision(_) => "decision", + _ => "other", + }; + self.0.lock().map_err(|_| "lock poisoned")?.push(kind); + Ok(()) + } + } + + let seen = Arc::new(Mutex::new(Vec::new())); + let router = FallThrough::new(target_set(&["strong", "weak"])) + .with_processor(Arc::new(RecordingProcessor(seen.clone()))) + .with_classifier(fixed(vec![score("strong", 1.0)])); + run(router).await?; + + assert_eq!( + *seen.lock().map_err(|_| "lock poisoned")?, + vec!["request", "decision"] + ); + Ok(()) + } + + #[tokio::test] + async fn state_persists_across_turns() -> Result<(), BoxErr> { + // Increments the session turn count on every request. + struct CountingProcessor; + + #[async_trait] + impl Processor for CountingProcessor { + async fn process(&self, state: &mut State, event: Event<'_>) -> Result<(), BoxErr> { + if let Event::Request(_) = event { + state.turn_count += 1; + } + Ok(()) + } + } + + // Routes to "weak" until the session has accumulated >= 2 turns, then "strong" — + // its decision depends only on state carried over from earlier turns. + struct ThresholdClassifier; + + #[async_trait] + impl Classifier for ThresholdClassifier { + async fn score( + &self, + state: &mut State, + _request: &Request, + _driver: Option<&Driver>, + ) -> Result { + let target = if state.turn_count >= 2 { + "strong" + } else { + "weak" + }; + Ok(Classification::Scores(vec![score(target, 1.0)])) + } + } + + // One shared, stateless router; the session lives in the context we thread through. + let router = Arc::new( + FallThrough::new(target_set(&["strong", "weak"])) + .with_processor(Arc::new(CountingProcessor)) + .with_classifier(Arc::new(ThresholdClassifier)), + ); + + // One session context, reused across three turns. The turn counter accumulates in + // its `State`, so the classifier crosses its threshold on turn 2. + let ctx = Context::::default(); + let (turn1, _) = run_turn(&router, ctx.clone()).await?; + let (turn2, _) = run_turn(&router, ctx.clone()).await?; + let (turn3, _) = run_turn(&router, ctx.clone()).await?; + + assert_eq!(turn1, "weak"); // count 1 — below threshold + assert_eq!(turn2, "strong"); // count 2 — state carried over from turn 1 + assert_eq!(turn3, "strong"); // count 3 — still above threshold + Ok(()) + } +} diff --git a/crates/libsy/src/core.rs b/crates/libsy/src/core.rs index a46cf2f4c..e3499ae10 100644 --- a/crates/libsy/src/core.rs +++ b/crates/libsy/src/core.rs @@ -6,7 +6,14 @@ //! re-exported at the crate root; algorithm implementations live in //! [`crate::algorithms`]. -pub mod algorithm; mod driver; +pub mod algorithm; +pub mod classifier; +pub mod processor; +pub mod state; + pub use algorithm::*; +pub use classifier::*; +pub use processor::*; +pub use state::*; diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 6313ef3ff..b8eccc0c4 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -55,7 +55,7 @@ pub struct RoutedRequest { /// 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 +/// 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 @@ -196,7 +196,7 @@ impl Default for Driver { } } -/// One item in the stream returned by [`Driver::stream`] / [`Algorithm::run_stream`]. +/// 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 @@ -268,17 +268,25 @@ impl LlmTargetSet { /// /// 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. +/// (the reference routers) get this for free; a stateful one keeps its per-session state +/// in the threaded [`Context`], not in the shared algorithm. +/// +/// The `S` type parameter is the per-session state carried in `Context`. It defaults +/// to `()`, so stateless algorithms just write `impl Algorithm for MyRouter` and callers +/// keep using `Arc`. A stateful algorithm picks its own state type (e.g. +/// `impl Algorithm for FallThrough`). #[async_trait] -pub trait Algorithm: Send + Sync + 'static { +pub trait Algorithm: Send + Sync + 'static +where + S: Clone + 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). + /// drive it. `ctx` carries the request's cross-cutting values and per-session `state`. async fn create_run_task( self: Arc, - ctx: Context, + ctx: Context, driver: Driver, request: Request, ) -> Result>; @@ -297,7 +305,7 @@ pub trait Algorithm: Send + Sync + 'static { /// Process a request to completion, returning a stream of [`Step`]s. /// Each [`Step::CallLlm`] is an offloaded model call the consumer must serve. /// The stream ends with a [`Step::ReturnToAgent`] on success, or an `Err` item on failure. - fn run_stream(self: Arc, ctx: Context, request: Request) -> StepStream { + fn run_stream(self: Arc, ctx: Context, request: Request) -> StepStream { let driver = Driver::new(); let task_driver = driver.clone(); let task_ctx = ctx.clone(); @@ -308,7 +316,8 @@ pub trait Algorithm: Send + Sync + 'static { let abort_guard = AbortOnDrop(handle.abort_handle()); let finish_driver = driver.clone(); - let finish_ctx = ctx.clone(); + // The terminal step carries no session state; hand `finish` the base context. + let finish_ctx = ctx.without_state(); let tail: StepStream = Box::pin( futures::stream::once(async move { let result = match handle.await { @@ -332,7 +341,7 @@ pub trait Algorithm: Send + Sync + 'static { /// [`Decision`]s the algorithm made along the way. async fn run( self: Arc, - ctx: Context, + ctx: Context, request: Request, ) -> Result<(Vec>, Response), Box> { // Serve up to this many offloaded calls concurrently, so an algorithm that diff --git a/crates/libsy/src/core/classifier.rs b/crates/libsy/src/core/classifier.rs new file mode 100644 index 000000000..2097885ea --- /dev/null +++ b/crates/libsy/src/core/classifier.rs @@ -0,0 +1,197 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::{Driver, State}; +use async_trait::async_trait; +use switchyard_protocol::Request; + +/// One classifier's recommendation of a routing `target`, with a `[0.0, 1.0]` confidence. +#[derive(Debug, Clone, PartialEq)] +pub struct Score { + /// `[0.0, 1.0]` confidence in `target`. + pub confidence: f64, + /// The target (model / tier) being recommended. + pub target: String, +} + +type BoxErr = Box; + +/// A classifier's verdict for a request: a set of target [`Score`]s, flagged by how +/// confident the classifier is that they are decisive. +pub enum Classification { + /// Definite recommendations; [`argmax`](Self::argmax) always yields the top target. + Scores(Vec), + /// Recommendations the classifier considers ambiguous; [`argmax`](Self::argmax) yields + /// nothing unless the caller opts to ignore ambiguity. + Ambiguous(Vec), +} + +impl Classification { + /// The top-scoring [`Score`], or `None` when the classifier abstained (an empty set). + /// + /// An [`Ambiguous`](Self::Ambiguous) classification also yields `None` unless + /// `ignore_ambiguous` is set, in which case it falls back to the plain argmax. + /// Errors if any confidence is `NaN` (an unorderable score the caller should surface). + pub fn argmax(&self, ignore_ambiguous: bool) -> Result, BoxErr> { + match self { + Classification::Scores(scores) => argmax(scores), + Classification::Ambiguous(scores) => { + if ignore_ambiguous { + argmax(scores) + } else { + Ok(None) + } + } + } + } +} + +/// The highest-confidence score, or `None` when the set is empty (the classifier abstained). +/// Ties keep the first. Errors on a `NaN` confidence, +/// which has no defined ordering. +fn argmax(scores: &[Score]) -> Result, BoxErr> { + let mut best: Option<&Score> = None; + for score in scores.iter() { + if score.confidence.is_nan() { + return Err(format!("argmax: NaN confidence for target '{}'", score.target).into()); + } + match best { + Some(cur_best) if score.confidence > cur_best.confidence => best = Some(score), + None => best = Some(score), + _ => {} + } + } + Ok(best.cloned()) +} + +/// Scores each of the classifier's targets given State the current Request +#[async_trait] +pub trait Classifier: Send + Sync { + /// Score the classifier's targets given the current state and request. + /// driver is optional. It is used to offload model calls + async fn score( + &self, + state: &mut State, + request: &Request, + driver: Option<&Driver>, + ) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::StateValue; + use switchyard_protocol::text_request; + + /// Terse `Score` builder for the assertions below. + fn score(target: &str, confidence: f64) -> Score { + Score { + target: target.to_string(), + confidence, + } + } + + #[test] + fn argmax_picks_the_highest_confidence_score() -> Result<(), BoxErr> { + let scores = vec![score("weak", 0.2), score("strong", 0.9), score("mid", 0.5)]; + let best = Classification::Scores(scores).argmax(false)?; + assert_eq!(best, Some(score("strong", 0.9))); + Ok(()) + } + + #[test] + fn argmax_breaks_ties_by_cascade_order() -> Result<(), BoxErr> { + // Equal confidence: the earlier target in cascade order wins the tie. + let scores = vec![score("first", 0.7), score("second", 0.7)]; + let best = Classification::Scores(scores).argmax(false)?; + assert_eq!(best.map(|s| s.target), Some("first".to_string())); + Ok(()) + } + + #[test] + fn argmax_on_an_empty_set_abstains() -> Result<(), BoxErr> { + // No scores means the classifier abstained — no choice to make. + assert_eq!(Classification::Scores(vec![]).argmax(false)?, None); + assert_eq!(Classification::Ambiguous(vec![]).argmax(true)?, None); + Ok(()) + } + + #[test] + fn argmax_errors_on_nan_confidence() { + // A NaN confidence has no defined ordering — surface it rather than guess. + let scores = vec![score("weak", 0.3), score("strong", f64::NAN)]; + assert!(Classification::Scores(scores).argmax(false).is_err()); + // A lone NaN errors too, even with nothing to compare it against. + assert!(Classification::Scores(vec![score("only", f64::NAN)]) + .argmax(false) + .is_err()); + } + + #[test] + fn ambiguous_without_ignore_makes_no_choice() -> Result<(), BoxErr> { + // Ambiguous means "don't pick" unless the caller opts to ignore ambiguity. + let scores = vec![score("strong", 0.9)]; + assert_eq!(Classification::Ambiguous(scores).argmax(false)?, None); + Ok(()) + } + + #[test] + fn ambiguous_with_ignore_falls_back_to_argmax() -> Result<(), BoxErr> { + let scores = vec![score("weak", 0.3), score("strong", 0.8)]; + let best = Classification::Ambiguous(scores).argmax(true)?; + assert_eq!(best, Some(score("strong", 0.8))); + Ok(()) + } + + #[test] + fn scores_variant_ignores_the_ambiguous_flag() -> Result<(), BoxErr> { + // A definitive classification always yields its argmax, regardless of the flag. + let scores = vec![score("a", 0.4), score("b", 0.6)]; + let with_ignore = Classification::Scores(scores.clone()).argmax(true)?; + let without_ignore = Classification::Scores(scores).argmax(false)?; + assert_eq!(with_ignore, without_ignore); + assert_eq!(with_ignore, Some(score("b", 0.6))); + Ok(()) + } + + /// Scores the request's requested model at full confidence and records that it ran. + struct RecordingClassifier; + + #[async_trait] + impl Classifier for RecordingClassifier { + async fn score( + &self, + state: &mut State, + request: &Request, + _driver: Option<&Driver>, + ) -> Result { + // Stash a marker in `extra` to prove state is threaded mutably. + state.extra.insert("ran".to_string(), StateValue::Count(1)); + let target = request.requested_model().unwrap_or("auto").to_string(); + Ok(Classification::Scores(vec![Score { + target, + confidence: 1.0, + }])) + } + } + + #[tokio::test] + async fn classifier_reads_request_and_mutates_state() -> Result<(), BoxErr> { + let mut state = State::default(); + let request = Request { + llm_request: text_request(Some("strong".to_string()), "hi"), + raw_request: None, + metadata: None, + }; + // A `None` driver is valid: the classifier scored without offloading a model call. + let classification = RecordingClassifier + .score(&mut state, &request, None) + .await?; + assert_eq!( + classification.argmax(false)?.map(|s| s.target), + Some("strong".to_string()) + ); + assert!(matches!(state.extra.get("ran"), Some(StateValue::Count(1)))); + Ok(()) + } +} diff --git a/crates/libsy/src/core/driver.rs b/crates/libsy/src/core/driver.rs index 093b6e306..ef57a3fa1 100644 --- a/crates/libsy/src/core/driver.rs +++ b/crates/libsy/src/core/driver.rs @@ -57,12 +57,16 @@ use crate::Context; use futures::{Stream, StreamExt}; use tokio::sync::{mpsc, oneshot}; +use tokio::time::{timeout, Duration}; use tokio_stream::wrappers::ReceiverStream; type BoxErr = Box; type BoxAny = Box; type StepResult = Result; +// TODO make request timeout configurable +const FULFILL_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); + /// One item on the stream returned by [`TypeErasedDriver::stream`]. pub enum DriverStep { /// A request awaiting a response. The consumer downcasts it and fulfills the @@ -139,6 +143,19 @@ impl TypeErasedDriver { } } + fn ensure_started(&self) -> Result<(), BoxErr> { + let guard = self + .inner + .step_rx + .lock() + .map_err(|_| "driver: lock poisoned")?; + if guard.is_none() { + Ok(()) + } else { + Err("get this driver's stream with driver.stream() before calling this method".into()) + } + } + /// 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`. @@ -147,6 +164,8 @@ impl TypeErasedDriver { REQ: Any + Send + 'static, RES: Any + Send + 'static, { + self.ensure_started()?; + let (tx, rx) = oneshot::channel::>(); let promise = DriverRequest { request: Box::new(req), @@ -160,10 +179,10 @@ impl TypeErasedDriver { // 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()), - }; + let response = timeout(FULFILL_REQUEST_TIMEOUT, rx) + .await + .map_err(|_| "driver: response timed out")? + .map_err(|_| "driver: error receiving response")??; response .downcast::() .map(|boxed| *boxed) @@ -177,6 +196,8 @@ impl TypeErasedDriver { where INFO: Any + Send + 'static, { + self.ensure_started()?; + self.inner .step_tx .send(Ok(DriverStep::Info(Box::new(info)))) @@ -192,6 +213,8 @@ impl TypeErasedDriver { where T: Any + Send + 'static, { + self.ensure_started()?; + self.inner .step_tx .send(Ok(DriverStep::Done(Box::new(payload)))) @@ -204,6 +227,8 @@ impl TypeErasedDriver { /// step). Awaits channel capacity and errors only if the stream is /// already closed. pub async fn fail(&self, _ctx: Context, err: BoxErr) -> Result<(), BoxErr> { + self.ensure_started()?; + self.inner .step_tx .send(Err(err)) @@ -534,4 +559,24 @@ mod tests { assert!(sent >= 2, "producer should have published the paced steps"); Ok(()) } + + #[test] + fn ensure_started_requires_the_stream_be_taken_first() -> Result<(), BoxErr> { + let driver = TypeErasedDriver::new(); + + // Before the consumer takes the stream, the receiver is still present, so producing + // is refused — with a message that points at the fix. + match driver.ensure_started() { + Ok(()) => return Err("expected ensure_started to reject an untaken stream".into()), + Err(err) => assert!( + err.to_string().contains("driver.stream()"), + "error should point at driver.stream(), got: {err}" + ), + } + + // Taking the stream claims the receiver (`Some` -> `None`), so the driver is started. + let _stream = driver.stream(); + assert!(driver.ensure_started().is_ok()); + Ok(()) + } } diff --git a/crates/libsy/src/core/processor.rs b/crates/libsy/src/core/processor.rs new file mode 100644 index 000000000..3fbd317ef --- /dev/null +++ b/crates/libsy/src/core/processor.rs @@ -0,0 +1,141 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::State; +use async_trait::async_trait; +use switchyard_protocol::{AggLlmResponse, Decision, Request, Signals}; + +/// An event observed by the algorithm. Events are consumed by [`Processor`] to mutate [`State`] +pub enum Event<'a> { + /// The inbound request that begins a turn. + Request(&'a Request), + /// An out-of-band agentic-stack signal (tool results, budget updates, …). + Signal(&'a Signals), + /// A routing decision the algorithm just made. + Decision(&'a dyn Decision), + /// A request about to be sent to a model. + ModelRequest(&'a Request), + /// A buffered response received back from a model. + ModelResponse(&'a AggLlmResponse), +} + +type BoxErr = Box; + +/// Collects events as the algorithm runs and mutates [`State`] +#[async_trait] +pub trait Processor: Send + Sync { + /// Process an event, accumulating facts into `state`. + async fn process(&self, state: &mut State, event: Event<'_>) -> Result<(), BoxErr>; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::StateValue; + use switchyard_protocol::{text_request, text_response}; + + /// The `State::extra` key each event variant tallies under. + fn event_key(event: &Event<'_>) -> &'static str { + match event { + Event::Request(_) => "requests", + Event::Signal(_) => "signals", + Event::Decision(_) => "decisions", + Event::ModelRequest(_) => "model_requests", + Event::ModelResponse(_) => "model_responses", + } + } + + /// Reads a `StateValue::Count` from `extra`, treating a missing key as zero. + fn count(state: &State, key: &str) -> u32 { + match state.extra.get(key) { + Some(StateValue::Count(n)) => *n, + _ => 0, + } + } + + /// Tallies each event variant under its own key in [`State::extra`]. + struct CountingProcessor; + + #[async_trait] + impl Processor for CountingProcessor { + async fn process(&self, state: &mut State, event: Event<'_>) -> Result<(), BoxErr> { + let entry = state + .extra + .entry(event_key(&event).to_string()) + .or_insert(StateValue::Count(0)); + if let StateValue::Count(n) = entry { + *n += 1; + } + Ok(()) + } + } + + /// Minimal [`Decision`] so an `Event::Decision` can be constructed. + struct TestDecision; + + impl Decision for TestDecision { + fn selected_model(&self) -> &str { + "test/model" + } + fn reasoning(&self) -> Option<&str> { + None + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + fn request() -> Request { + Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata: None, + } + } + + #[tokio::test] + async fn processor_tallies_each_event_variant_into_state() -> Result<(), BoxErr> { + let processor = CountingProcessor; + let mut state = State::default(); + let req = request(); + let response = text_response(None, "ok"); + let decision = TestDecision; + let signals = Signals {}; + + // Feed one of every event variant through the processor. + processor.process(&mut state, Event::Request(&req)).await?; + processor + .process(&mut state, Event::ModelRequest(&req)) + .await?; + processor + .process(&mut state, Event::ModelResponse(&response)) + .await?; + processor + .process(&mut state, Event::Decision(&decision)) + .await?; + processor + .process(&mut state, Event::Signal(&signals)) + .await?; + + assert_eq!(count(&state, "requests"), 1); + assert_eq!(count(&state, "signals"), 1); + assert_eq!(count(&state, "decisions"), 1); + assert_eq!(count(&state, "model_requests"), 1); + assert_eq!(count(&state, "model_responses"), 1); + Ok(()) + } + + #[tokio::test] + async fn process_accumulates_state_across_repeated_events() -> Result<(), BoxErr> { + let processor = CountingProcessor; + let mut state = State::default(); + let req = request(); + + for _ in 0..3 { + processor.process(&mut state, Event::Request(&req)).await?; + } + + assert_eq!(count(&state, "requests"), 3); + Ok(()) + } +} diff --git a/crates/libsy/src/core/state.rs b/crates/libsy/src/core/state.rs new file mode 100644 index 000000000..330ca3e7f --- /dev/null +++ b/crates/libsy/src/core/state.rs @@ -0,0 +1,26 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; +use std::sync::Arc; + +use tokio::sync::Mutex; + +/// A value in a session's [`State`]. +#[derive(Debug, Clone)] +pub enum StateValue { + String(String), + Count(u32), + Int(i32), + Scalar(f32), +} + +/// State maitaineed by [`Algorithm`]s +#[derive(Debug, Clone, Default)] +pub struct State { + pub turn_count: u32, + pub extra: HashMap, +} + +/// State shared accross tunrs / sessions +pub type SharedState = Arc>; diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 2f8eebca8..60221beac 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -56,10 +56,10 @@ //! //! Concrete algorithms live in [`algorithms`]: //! -//! [`algorithms::Random`](crate::algorithms::Random) provides uniform random routing. +//! [`algorithms::Random`] provides uniform random routing. //! -//! [`algorithms::LlmClassifier`](crate::algorithms::LlmClassifier) classifies -//! with one model, then routes to a strong/weak model depending on the classifier's choice. +//! [`algorithms::LlmClassifier`] classifies with one model, then routes to a strong/weak +//! model depending on the classifier's choice. mod core; pub use core::*; diff --git a/crates/protocol/src/envelope.rs b/crates/protocol/src/envelope.rs index 7899fd523..b97175e0f 100644 --- a/crates/protocol/src/envelope.rs +++ b/crates/protocol/src/envelope.rs @@ -7,12 +7,24 @@ use crate::{LlmRequest, LlmResponse, Metadata}; use std::collections::HashMap; -/// Per-request state threaded to an algorithm alongside a request. A placeholder -/// for cross-cutting state (correlation ids, budgets, deadlines) an algorithm will -/// read; empty today. +/// [`Context`] for an algorithm or llm client. +/// This struct is fed through the entire process and provides unified information to each step #[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Context { +pub struct Context { + /// caller specific values pub values: HashMap, + /// Caller-defined per-session state; `()` when a consumer carries no session state. + pub state: S, +} + +impl Context { + /// Create a new [`Context`] with the given state and an empty values map. + pub fn without_state(&self) -> Context<()> { + Context { + values: self.values.clone(), + state: (), + } + } } /// Agentic-stack events fed to an algorithm out of band (e.g. tool results, budget