From 57c015db2b32fb8817a4caaca5e65941871b4be2 Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Thu, 16 Jul 2026 14:30:12 -0400 Subject: [PATCH] chore(libsy): move Decision/RoutedLlmClient traits to protocol Signed-off-by: Greg Clark --- Cargo.lock | 1 + .../libsy-examples/examples/research_agent.rs | 21 ++-- crates/libsy-examples/src/ensemble.rs | 36 +++--- crates/libsy-examples/src/llm_class.rs | 17 ++- crates/libsy-protocol/Cargo.toml | 1 + crates/libsy-protocol/src/client.rs | 53 ++++++++ crates/libsy-protocol/src/envelope.rs | 9 +- crates/libsy-protocol/src/lib.rs | 2 + crates/libsy/README.md | 26 ++-- crates/libsy/src/algorithms/rand.rs | 15 ++- crates/libsy/src/lib.rs | 115 ++++++++---------- 11 files changed, 177 insertions(+), 119 deletions(-) create mode 100644 crates/libsy-protocol/src/client.rs diff --git a/Cargo.lock b/Cargo.lock index 266129c7..18749f29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1633,6 +1633,7 @@ dependencies = [ name = "switchyard-protocol" version = "0.1.0" dependencies = [ + "async-trait", "futures", "serde", "serde_json", diff --git a/crates/libsy-examples/examples/research_agent.rs b/crates/libsy-examples/examples/research_agent.rs index 221a37b6..cba976b5 100644 --- a/crates/libsy-examples/examples/research_agent.rs +++ b/crates/libsy-examples/examples/research_agent.rs @@ -3,7 +3,7 @@ //! Minimal research agent using the [`Algorithm::run`] convenience. //! -//! Every target owns an `LlmClient`, so the agent runs each request to completion with +//! Every target owns an `RoutedLlmClient`, 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 @@ -16,8 +16,8 @@ use std::sync::Arc; use async_trait::async_trait; use libsy::{ - Algorithm, Context, LlmClient, LlmResponse, LlmTarget, LlmTargetSet, Request, Response, - RoutedRequest, + Algorithm, Context, Decision, LlmResponse, LlmTarget, LlmTargetSet, Request, Response, + RoutedLlmClient, }; use libsy_examples::llm_class::LlmClassifierOrchAlgo; use switchyard_protocol::{completion_text, text_request, text_response}; @@ -26,14 +26,19 @@ 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. +/// Stub transport. Real integrators implement `RoutedLlmClient` over their own HTTP. struct StubClient; #[async_trait] -impl LlmClient for StubClient { - async fn call(&self, routed: RoutedRequest) -> Result> { +impl RoutedLlmClient for StubClient { + async fn call( + &self, + _ctx: Context, + _request: Request, + decision: Arc, + ) -> Result> { // The model to call is the routed decision's selection, not the inbound name. - let model = routed.decision.selected_model().to_string(); + let model = decision.selected_model().to_string(); println!(" -> model call: {model}"); // The classifier returns a score; other models return an answer. let completion = if model == CLASSIFIER { @@ -49,7 +54,7 @@ impl LlmClient for StubClient { } fn targets() -> LlmTargetSet { - let client = Arc::new(StubClient) as Arc; + let client = Arc::new(StubClient) as Arc; let target = |name: &str| LlmTarget { semantic_name: name.to_string(), llm_client: Some(client.clone()), diff --git a/crates/libsy-examples/src/ensemble.rs b/crates/libsy-examples/src/ensemble.rs index 333c8e59..102bcb52 100644 --- a/crates/libsy-examples/src/ensemble.rs +++ b/crates/libsy-examples/src/ensemble.rs @@ -395,7 +395,7 @@ impl Algorithm for EnsembleOrchAlgo { #[cfg(test)] mod tests { use super::*; - use libsy::{LlmClient, LlmResponse, LlmTarget, Response, RoutedRequest, Signals}; + use libsy::{LlmResponse, LlmTarget, Response, RoutedLlmClient, Signals}; use std::sync::Mutex as StdMutex; use switchyard_protocol::text_response; @@ -410,18 +410,20 @@ mod tests { } #[async_trait] - impl LlmClient for JudgingClient { + impl RoutedLlmClient for JudgingClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + request: Request, + decision: Arc, ) -> Result> { - let name = routed.decision.selected_model().to_string(); + let name = decision.selected_model().to_string(); self.calls .lock() .map_err(|_| "lock poisoned")? .push(name.clone()); let completion = if name == self.judge_model { - judge_pick(&prompt_text(&routed.request.llm_request), &self.prefer) + judge_pick(&prompt_text(&request.llm_request), &self.prefer) } else { format!("answer from {name}") }; @@ -463,7 +465,7 @@ mod tests { judge_model: judge.to_string(), prefer: prefer.to_string(), calls: Arc::clone(&calls), - }) as Arc; + }) as Arc; let target = |name: &str| LlmTarget { semantic_name: name.to_string(), llm_client: Some(client.clone()), @@ -485,7 +487,7 @@ mod tests { candidates: &[&str], judge: &str, exploration_turns: u64, - client: Arc, + client: Arc, ) -> EnsembleOrchAlgo { let target = |name: &str| LlmTarget { semantic_name: name.to_string(), @@ -653,15 +655,17 @@ mod tests { /// Client whose every call fails. struct FailingClient; #[async_trait] - impl LlmClient for FailingClient { + impl RoutedLlmClient for FailingClient { async fn call( &self, - _routed: RoutedRequest, + _ctx: Context, + _request: Request, + _decision: Arc, ) -> Result> { Err("upstream down".into()) } } - let client = Arc::new(FailingClient) as Arc; + let client = Arc::new(FailingClient) as Arc; let target = |name: &str| LlmTarget { semantic_name: name.to_string(), llm_client: Some(client.clone()), @@ -718,15 +722,17 @@ mod tests { } #[async_trait] - impl LlmClient for BarrierClient { + impl RoutedLlmClient for BarrierClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + request: Request, + decision: Arc, ) -> Result> { - let name = routed.decision.selected_model().to_string(); + let name = decision.selected_model().to_string(); let completion = if name == self.judge_model { // Judge runs after the barrier releases; it must not wait. - judge_pick(&prompt_text(&routed.request.llm_request), &self.prefer) + judge_pick(&prompt_text(&request.llm_request), &self.prefer) } else { // Hold every candidate call until all sessions have fanned out. self.barrier.wait().await; @@ -746,7 +752,7 @@ mod tests { barrier: barrier.clone(), judge_model: "judge/haiku".to_string(), prefer: "a/model".to_string(), - }) as Arc; + }) as Arc; // Two independent sessions: separate algo instances, each with its own // per-session state, sharing only the backend client. diff --git a/crates/libsy-examples/src/llm_class.rs b/crates/libsy-examples/src/llm_class.rs index a0612dd6..67e9336a 100644 --- a/crates/libsy-examples/src/llm_class.rs +++ b/crates/libsy-examples/src/llm_class.rs @@ -175,7 +175,7 @@ impl Algorithm for LlmClassifierOrchAlgo { #[cfg(test)] mod tests { use super::*; - use libsy::{LlmClient, LlmResponse, LlmTarget, Response, RoutedRequest}; + use libsy::{LlmResponse, LlmTarget, Response, RoutedLlmClient}; use std::sync::Mutex; use switchyard_protocol::text_response; @@ -189,21 +189,20 @@ mod tests { } #[async_trait] - impl LlmClient for ScoringClient { + impl RoutedLlmClient for ScoringClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + request: Request, + decision: Arc, ) -> Result> { - let name = routed.decision.selected_model().to_string(); + let name = 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); + self.seen.lock().map_err(|_| "lock poisoned")?.push(request); Ok(Response { llm_response: LlmResponse::Agg(text_response(None, completion)), metadata: None, @@ -218,7 +217,7 @@ mod tests { classifier_model: "router/classifier".to_string(), score: score.to_string(), seen: Arc::clone(&seen), - }) as Arc; + }) as Arc; let target = |name: &str| LlmTarget { semantic_name: name.to_string(), llm_client: Some(client.clone()), diff --git a/crates/libsy-protocol/Cargo.toml b/crates/libsy-protocol/Cargo.toml index 6e4ec82c..9e4b891e 100644 --- a/crates/libsy-protocol/Cargo.toml +++ b/crates/libsy-protocol/Cargo.toml @@ -17,6 +17,7 @@ keywords = ["llm", "protocol", "openai", "anthropic"] publish = ["crates-io"] [dependencies] +async-trait = "0.1" futures = "0.3" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/crates/libsy-protocol/src/client.rs b/crates/libsy-protocol/src/client.rs new file mode 100644 index 00000000..e6692406 --- /dev/null +++ b/crates/libsy-protocol/src/client.rs @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The routed-call server trait and the routing decision it carries. +//! +//! [`RoutedLlmClient`] is the one piece of I/O the protocol does not own: a host +//! implements it to actually perform a model call. [`Decision`] is the routing +//! decision that produced the call, carried alongside so the client and any +//! observer can see which model was chosen and why. Both live here — rather than +//! in libsy's orchestration crate — so a client crate that depends only on the +//! protocol can serve routed calls without pulling in the orchestrator. + +use std::error::Error; +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::{Context, Request, Response}; + +/// 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; +} + +/// Performs the actual model call for a target. This is the one piece of I/O the +/// library 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 a routed request's `default_client`. +#[async_trait] +pub trait RoutedLlmClient: Send + Sync { + /// Serve the call, returning the model's response. Call the model named by + /// [`decision.selected_model()`](Decision::selected_model) — the target the algorithm + /// routed to — mapping it to whatever provider model id this client hits. + /// `request.llm_request.model` is the agent's original name, carried through for + /// reference, not a call target. `ctx` carries the request's cross-cutting state. + async fn call( + &self, + ctx: Context, + request: Request, + decision: Arc, + ) -> Result>; +} diff --git a/crates/libsy-protocol/src/envelope.rs b/crates/libsy-protocol/src/envelope.rs index af218251..9906bd86 100644 --- a/crates/libsy-protocol/src/envelope.rs +++ b/crates/libsy-protocol/src/envelope.rs @@ -4,13 +4,16 @@ //! The request/response envelope: the normalized [`LlmRequest`]/[`LlmResponse`] paired //! with the original provider payload and correlation [`Metadata`]. -use crate::{LlmRequest, LlmResponse}; +use crate::{LlmRequest, LlmResponse, WireFormat}; +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. #[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Context {} +pub struct Context { + pub values: HashMap, +} /// Correlation and routing metadata attached to a request or response. /// @@ -31,6 +34,8 @@ pub struct Metadata { pub extra_metadata: Option>, /// HTTP headers to attach when forwarding the request/response, if any. pub http_headers: Option>, + /// The wire format the request/response was originally encoded in, if known. + pub wire_format: Option, } /// A request an algorithm routes: the normalized [`LlmRequest`] plus the original diff --git a/crates/libsy-protocol/src/lib.rs b/crates/libsy-protocol/src/lib.rs index f2f3076b..a1ae2cf9 100644 --- a/crates/libsy-protocol/src/lib.rs +++ b/crates/libsy-protocol/src/lib.rs @@ -20,11 +20,13 @@ //! //! [`Algorithm`]: https://docs.rs/libsy +pub mod client; pub mod envelope; pub mod format; pub mod llm; pub mod stream; +pub use client::*; pub use envelope::*; pub use format::*; pub use llm::*; diff --git a/crates/libsy/README.md b/crates/libsy/README.md index 423ab070..36f28495 100644 --- a/crates/libsy/README.md +++ b/crates/libsy/README.md @@ -11,13 +11,13 @@ so it drops into a proxy, gateway, or agent runtime. Build a target set, pick an algorithm, run a request: ```rust -use libsy::{Algorithm, Context, LlmClient, LlmTarget, LlmTargetSet, Request}; +use libsy::{Algorithm, Context, RoutedLlmClient, LlmTarget, LlmTargetSet, Request}; use libsy_examples::llm_class::LlmClassifierOrchAlgo; use switchyard_protocol::{completion_text, text_request}; use std::sync::Arc; -// Targets the algorithm routes among, each backed by your LlmClient (see below). -let client = Arc::new(MyClient { /* .. */ }) as Arc; +// Targets the algorithm routes among, each backed by your RoutedLlmClient (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( @@ -75,19 +75,19 @@ fidelity. ## 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 +An `LlmTarget` pairs a routing `semantic_name` with an optional `RoutedLlmClient`. Mapping that +name to a provider model id is the client's job, not the algorithm's — `RoutedLlmClient` 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) +impl RoutedLlmClient for MyClient { + async fn call(&self, ctx: Context, request: Request, decision: Arc) -> Result> { - let model = routed.decision.selected_model(); // the routed target — map it to a provider id - // routed.request.llm_request.model is the agent's original name (not a call target) + let model = decision.selected_model(); // the routed target — map it to a provider id + // request.llm_request.model is the agent's original name (not a call target) // ... POST to your endpoint, read the completion ... let completion = String::from("provider response text"); Ok(Response { @@ -104,9 +104,9 @@ To stream instead, return `LlmResponse::Stream(..)` — a boxed `Stream>` — and emit chunks as they arrive from your upstream. See [Streaming responses](#streaming-responses) below. -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 +A `RoutedRequest` bundles the `request` with the routing `decision`, the target's +`default_client`, and the request's `ctx`; 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 @@ -129,7 +129,7 @@ so pulling it paces the algorithm; each run is independent, so many run concurre ## Streaming responses A `Response` carries an `LlmResponse` that is either buffered (`Agg`) or a live token -stream (`Stream`). An `LlmClient` — or an algorithm — chooses: return +stream (`Stream`). An `RoutedLlmClient` — or an algorithm — chooses: return `LlmResponse::Agg(..)` for a whole answer, or `LlmResponse::Stream(..)` to forward tokens as they arrive. libsy never buffers a stream on your behalf; it flows through the algorithm untouched, so `run` / `run_stream` hand back whatever was produced and **the caller diff --git a/crates/libsy/src/algorithms/rand.rs b/crates/libsy/src/algorithms/rand.rs index b80a84a2..adc2b366 100644 --- a/crates/libsy/src/algorithms/rand.rs +++ b/crates/libsy/src/algorithms/rand.rs @@ -86,27 +86,26 @@ impl Algorithm for RandomAlgo { #[cfg(test)] mod tests { + use super::*; use std::collections::HashSet; use switchyard_protocol::{completion_text, text_request, text_response}; - use super::*; - use crate::{LlmClient, LlmResponse, LlmTarget, RoutedRequest, Signals}; + use crate::{LlmResponse, LlmTarget, Request, RoutedLlmClient, Signals}; /// Echoes the selected target so tests can inspect which target was called. struct EchoClient; #[async_trait] - impl LlmClient for EchoClient { + impl RoutedLlmClient for EchoClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + _request: Request, + decision: Arc, ) -> Result> { Ok(Response { - llm_response: LlmResponse::Agg(text_response( - None, - routed.decision.selected_model(), - )), + llm_response: LlmResponse::Agg(text_response(None, decision.selected_model())), metadata: None, }) } diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index d4b81ee3..9d107233 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -22,7 +22,7 @@ //! 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 +//! target's [`RoutedLlmClient`], if any, rides along as //! [`RoutedRequest::default_client`] so the host can serve it by default or //! override it (see below). //! @@ -74,8 +74,8 @@ use futures::{Stream, StreamExt}; /// [`LlmResponseChunk`] is one streaming event; [`LlmResponse`] is the streamed response /// (a live [`LlmResponseStream`] or the terminal aggregate). pub use switchyard_protocol::{ - AggLlmResponse, Context, LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStream, - Metadata, Request, Response, + AggLlmResponse, Context, Decision, LlmRequest, LlmResponse, LlmResponseChunk, + LlmResponseStream, Metadata, Request, Response, RoutedLlmClient, }; use crate::driver::{DriverRequest, DriverStep, TypeErasedDriver}; @@ -96,24 +96,8 @@ pub type StepStream = Pin> + Send>>; #[derive(Clone)] pub struct Signals {} -/// 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. +/// A request paired with the routing [`Decision`] that produced it — the offload +/// payload a host reads (via [`CallLlmRequest::get_routed`]) to serve the call. /// /// The two model identifiers live in separate, unambiguous places: the model to /// call is [`decision.selected_model()`](Decision::selected_model), while @@ -129,7 +113,11 @@ pub struct RoutedRequest { /// 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>, + pub default_client: Option>, + /// The request's cross-cutting context, carried through the offload so whoever + /// serves the call (libsy's own `run`, or a host driving the stream) hands it to + /// [`RoutedLlmClient::call`]. + pub ctx: Context, } /// The host-facing half of an offloaded model call, surfaced inside [`Step::CallLlm`]. @@ -203,10 +191,11 @@ impl Driver { } /// 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, ctx: Context, routed: RoutedRequest) -> Result { + /// consumer's [`Response`]. The call's context travels inside + /// [`routed.ctx`](RoutedRequest::ctx). Errors if the stream is closed or the call failed. + pub async fn call_llm(&self, routed: RoutedRequest) -> Result { self.driver - .fulfill_request::(ctx, routed) + .fulfill_request::(routed.ctx.clone(), routed) .await } @@ -222,14 +211,12 @@ impl Driver { request: Request, decision: Arc, ) -> Result { - self.call_llm( + self.call_llm(RoutedRequest { + request, + decision, + default_client: target.llm_client.clone(), ctx, - RoutedRequest { - request, - decision, - default_client: target.llm_client.clone(), - }, - ) + }) .await } @@ -289,22 +276,8 @@ pub enum Step { ReturnToAgent(Box), } -/// 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.model` 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 +/// [`RoutedLlmClient`] 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)] @@ -315,7 +288,7 @@ pub struct LlmTarget { 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>, + pub llm_client: Option>, } /// The set of targets an algorithm may route among. An algorithm is constructed @@ -430,7 +403,11 @@ pub trait Algorithm: Send + Sync + 'static { routed.decision.selected_model() ) })?; - call.respond(client.call(routed).await) + call.respond( + client + .call(routed.ctx, routed.request, routed.decision) + .await, + ) } let stream = self.run_stream(ctx, request); @@ -480,16 +457,18 @@ mod tests { struct EchoClient; #[async_trait] - impl LlmClient for EchoClient { + impl RoutedLlmClient for EchoClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + _request: Request, + decision: Arc, ) -> Result> { // Echo back the model the algorithm routed to (the decision's selection). Ok(Response { llm_response: LlmResponse::Agg(text_response( None, - routed.decision.selected_model().to_string(), + decision.selected_model().to_string(), )), metadata: None, }) @@ -561,7 +540,7 @@ mod tests { .iter() .map(|(name, has_client)| LlmTarget { semantic_name: name.to_string(), - llm_client: has_client.then(|| Arc::new(EchoClient) as Arc), + llm_client: has_client.then(|| Arc::new(EchoClient) as Arc), }) .collect(); LlmTargetSet::new(targets) @@ -574,10 +553,12 @@ mod tests { } #[async_trait] - impl LlmClient for StreamingClient { + impl RoutedLlmClient for StreamingClient { async fn call( &self, - _routed: RoutedRequest, + _ctx: Context, + _request: Request, + _decision: Arc, ) -> Result> { let stream = futures::stream::iter(self.chunks.clone().into_iter().map(Ok)).boxed(); Ok(Response { @@ -591,7 +572,7 @@ mod tests { fn streaming_orch(chunks: Vec) -> Arc { let target = LlmTarget { semantic_name: "stream/model".to_string(), - llm_client: Some(Arc::new(StreamingClient { chunks }) as Arc), + llm_client: Some(Arc::new(StreamingClient { chunks }) as Arc), }; orch(LlmTargetSet::new(vec![target])) } @@ -718,7 +699,9 @@ mod tests { .default_client .clone() .ok_or("expected a default client")?; - let result = client.call(routed).await; + let result = client + .call(routed.ctx, routed.request, routed.decision) + .await; call.respond(result)?; } Step::Decision(_) => {} @@ -788,16 +771,18 @@ mod tests { } #[async_trait] - impl LlmClient for BarrierClient { + impl RoutedLlmClient for BarrierClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + _request: Request, + decision: Arc, ) -> Result> { self.barrier.wait().await; Ok(Response { llm_response: LlmResponse::Agg(text_response( None, - routed.decision.selected_model().to_string(), + decision.selected_model().to_string(), )), metadata: None, }) @@ -891,10 +876,12 @@ mod tests { } #[async_trait] - impl LlmClient for ProbeClient { + impl RoutedLlmClient for ProbeClient { async fn call( &self, - routed: RoutedRequest, + _ctx: Context, + _request: Request, + decision: Arc, ) -> Result> { let now = self.current.fetch_add(1, Ordering::SeqCst) + 1; self.max.fetch_max(now, Ordering::SeqCst); @@ -906,7 +893,7 @@ mod tests { Ok(Response { llm_response: LlmResponse::Agg(text_response( None, - routed.decision.selected_model().to_string(), + decision.selected_model().to_string(), )), metadata: None, }) @@ -958,7 +945,7 @@ mod tests { max: max.clone(), entered: entered_tx, gate: gate.clone(), - }) as Arc; + }) as Arc; let target = LlmTarget { semantic_name: "m".to_string(), llm_client: Some(client),