diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs index 059675bbe..f53782555 100644 --- a/crates/libsy-llm-client/src/lib.rs +++ b/crates/libsy-llm-client/src/lib.rs @@ -30,7 +30,7 @@ pub use client::{ModelConfig, TranslatingLlmClient}; pub use error::{LlmClientError, Result}; pub use observation::{LlmCallObservation, RunObservation, RunObserver}; pub use raw::RawResponse; -pub use run::{ClientRouter, run}; +pub use run::{ClientRouter, run, serve_routing_call}; pub use switchyard_translation::RawEventStream; /// Registers process-wide compatibility gauges with the global meter provider. diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index df14f22f8..0932d9e77 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -8,10 +8,10 @@ //! consumer — it drives the stream with [`switchyard_libsy::drive`], hands routing-time calls to //! a [`RoutedLlmClient`], and serves the terminal routing outcome. //! -//! libsy owns the stream mechanics; what this module adds is ordered candidate fallback and the -//! `libsy.client_call` span around each candidate. Each candidate exhausts its backend retry -//! budget before fallback advances, so the worst case is `candidates × (max_retries + 1)` -//! upstream attempts plus every candidate's backoff. +//! libsy owns the stream mechanics; what this module adds is per-target request preparation, +//! ordered candidate fallback, and the `libsy.client_call` span around each candidate. Each +//! candidate exhausts its backend retry budget before fallback advances, so the worst case is +//! `candidates × (max_retries + 1)` upstream attempts plus every candidate's backoff. use std::collections::HashMap; use std::sync::Arc; @@ -23,6 +23,7 @@ use switchyard_libsy::{Algorithm, CallModel, LibsyError, Result, drive}; use switchyard_protocol::{ LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason, }; +use switchyard_translation::prepare_request_for_target; use crate::observation::{LlmCallObservation, RunObservation, RunObserver}; use crate::{metrics, observability}; @@ -84,6 +85,7 @@ pub async fn run( &algorithm_name, &outcome.request, &models, + CallPhase::Completion, &observe, ) .await; @@ -144,22 +146,39 @@ async fn serve( &call.algorithm, &call.request, &call.models, + CallPhase::Routing, &observe, ) .await; call.respond(result) } +/// Serves one model call requested while a libsy algorithm is routing. +/// +/// Candidate failures are returned to the algorithm through the call's response channel. +pub async fn serve_routing_call(clients: ClientRouter, call: CallModel) -> Result<()> { + serve(clients, call, None).await +} + +enum CallPhase { + Routing, + Completion, +} + /// Try candidates in order until one succeeds or a failure stops fallback. async fn call_first_available( clients: &ClientRouter, algorithm: &str, request: &Request, models: &[ModelId], + phase: CallPhase, observe: &(dyn Fn(LlmCallObservation) + Send + Sync), ) -> Result { for (index, target) in models.iter().enumerate() { - let request = request_for(request, target); + let request = match phase { + CallPhase::Routing => clients.prepare_routing_request(request.clone(), target), + CallPhase::Completion => clients.prepare_completion_request(request.clone(), target), + }; match call_one( clients, target, @@ -298,13 +317,6 @@ fn fallback_reason(error: &LibsyError) -> Option { } } -/// Clone a request and stamp the candidate model that should receive it. -fn request_for(request: &Request, target: &ModelId) -> Request { - let mut request = request.clone(); - request.llm_request.model = Some(target.to_string()); - request -} - /// Resolves a routed call's selected model to the client that serves it. /// /// An algorithm routes among named targets; which provider each target lives on is the @@ -315,9 +327,17 @@ fn request_for(request: &Request, target: &ModelId) -> Request { /// Cloning is cheap — the mapping is shared, so one router can serve every request. #[derive(Clone)] pub struct ClientRouter { - routing: Arc, + inner: Arc, } +#[derive(Clone)] +struct ClientRouting { + routing: Routing, + target_prompts: HashMap, + routing_answer_target: Option, +} + +#[derive(Clone)] enum Routing { /// One client serves every model. Single(Arc), @@ -329,7 +349,11 @@ impl ClientRouter { /// Build a router over `model name -> client`, for targets spread across providers. pub fn new(by_model: HashMap>) -> Self { Self { - routing: Arc::new(Routing::ByModel(by_model)), + inner: Arc::new(ClientRouting { + routing: Routing::ByModel(by_model), + target_prompts: HashMap::new(), + routing_answer_target: None, + }), } } @@ -340,10 +364,27 @@ impl ClientRouter { /// only duplicate that. pub fn single(client: Arc) -> Self { Self { - routing: Arc::new(Routing::Single(client)), + inner: Arc::new(ClientRouting { + routing: Routing::Single(client), + target_prompts: HashMap::new(), + routing_answer_target: None, + }), } } + /// Attach system prompts and the target whose routing call may answer the request. + /// Pass `None` when routing only selects a later completion target. + pub fn with_target_prompts( + mut self, + prompts: HashMap, + routing_answer_target: Option, + ) -> Self { + let inner = Arc::make_mut(&mut self.inner); + inner.target_prompts = prompts; + inner.routing_answer_target = routing_answer_target; + self + } + /// The client that serves `model`. /// /// Errors with [`LlmClientError::Configuration`] when the router maps models and has no @@ -352,7 +393,7 @@ impl ClientRouter { &self, model: &ModelId, ) -> std::result::Result<&Arc, LlmClientError> { - match self.routing.as_ref() { + match &self.inner.routing { Routing::Single(client) => Ok(client), Routing::ByModel(by_model) => { by_model @@ -363,6 +404,24 @@ impl ClientRouter { } } } + + /// Prepare a completion candidate with its configured target prompt. + pub fn prepare_completion_request(&self, mut request: Request, target: &ModelId) -> Request { + let prompt = self.inner.target_prompts.get(target).map(String::as_str); + prepare_request_for_target(&mut request.llm_request, target, prompt); + request + } + + /// Prepare a routing call, adding a target prompt only when it generates a candidate answer. + fn prepare_routing_request(&self, mut request: Request, target: &ModelId) -> Request { + let prompt = if self.inner.routing_answer_target.as_ref() == Some(target) { + self.inner.target_prompts.get(target).map(String::as_str) + } else { + None + }; + prepare_request_for_target(&mut request.llm_request, target, prompt); + request + } } impl FromIterator<(ModelId, Arc)> for ClientRouter { @@ -381,8 +440,8 @@ mod tests { use http::StatusCode; use switchyard_libsy::{Driver, RoutingOutcome}; use switchyard_protocol::{ - LlmResponse, LlmResponseChunk, LlmResponseStreamEvent, completion_text, text_request, - text_response, + ContentBlock, LlmResponse, LlmResponseChunk, LlmResponseStreamEvent, completion_text, + text_request, text_response, }; use wiremock::matchers::method; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -449,6 +508,7 @@ mod tests { struct CandidateClient { calls: Mutex>, + requests: Mutex>, first: FirstOutcome, } @@ -457,6 +517,7 @@ mod tests { async fn call(&self, request: Request) -> std::result::Result { let model = request.model_id().unwrap_or_default(); self.calls.lock().push(model.clone()); + self.requests.lock().push(request); if model == "weak" { return match self.first { FirstOutcome::ContextWindow => Err(LlmClientError::ContextWindowExceeded { @@ -516,11 +577,25 @@ mod tests { } } + fn instruction_text(request: &Request) -> Vec<&str> { + request + .llm_request + .instructions + .iter() + .flat_map(|instruction| &instruction.content) + .filter_map(|block| match block { + ContentBlock::Text { text } => Some(text.as_str()), + _ => None, + }) + .collect() + } + async fn run_candidates( first: FirstOutcome, ) -> (Arc, Result<(ModelId, Response)>) { let client = Arc::new(CandidateClient { calls: Mutex::new(Vec::new()), + requests: Mutex::new(Vec::new()), first, }); let algorithm = Arc::new(CandidateAlgorithm { @@ -540,6 +615,7 @@ mod tests { async fn answered_outcome_does_not_make_a_second_model_call() -> Result<()> { let client = Arc::new(CandidateClient { calls: Mutex::new(Vec::new()), + requests: Mutex::new(Vec::new()), first: FirstOutcome::StreamSuccess, }); let observations = Arc::new(Mutex::new(Vec::new())); @@ -569,6 +645,66 @@ mod tests { Ok(()) } + #[tokio::test] + async fn each_fallback_candidate_receives_only_its_own_prompt() -> Result<()> { + let client = Arc::new(CandidateClient { + calls: Mutex::new(Vec::new()), + requests: Mutex::new(Vec::new()), + first: FirstOutcome::ContextWindow, + }); + let clients = ClientRouter::single(client.clone()).with_target_prompts( + HashMap::from([ + ("weak".into(), "weak prompt".to_string()), + ("strong".into(), "strong prompt".to_string()), + ]), + None, + ); + + run( + Arc::new(CandidateAlgorithm { + models: vec!["weak".into(), "strong".into()], + }), + clients, + request(), + None, + ) + .await?; + + let calls = client.requests.lock(); + assert_eq!(calls.len(), 2); + assert_eq!(instruction_text(&calls[0]), ["weak prompt"]); + assert_eq!(instruction_text(&calls[1]), ["strong prompt"]); + Ok(()) + } + + #[tokio::test] + async fn configured_routing_response_target_receives_its_prompt() -> Result<()> { + let client = Arc::new(CandidateClient { + calls: Mutex::new(Vec::new()), + requests: Mutex::new(Vec::new()), + first: FirstOutcome::StreamSuccess, + }); + let clients = ClientRouter::single(client.clone()).with_target_prompts( + HashMap::from([("answer".into(), "answer prompt".to_string())]), + Some("answer".into()), + ); + + run( + Arc::new(AnsweredAlgorithm { + model: "answer".into(), + }), + clients, + request(), + None, + ) + .await?; + + let calls = client.requests.lock(); + assert_eq!(calls.len(), 1); + assert_eq!(instruction_text(&calls[0]), ["answer prompt"]); + Ok(()) + } + #[test] fn fallback_only_accepts_context_and_unavailable_failures() { let error = |source| LibsyError::client_call("target", source); diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index dc4db13ee..7a34b9ce5 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -6,9 +6,9 @@ use std::error::Error; use std::sync::Arc; -use libsy::{Algorithm, CallModel, LibsyError, RoutingOutcome, drive}; +use libsy::{Algorithm, LibsyError, RoutingOutcome, drive}; use serde_json::Value; -use switchyard_llm_client::{ClientRouter, RunObserver, TranslatingLlmClient}; +use switchyard_llm_client::{ClientRouter, RunObserver, TranslatingLlmClient, serve_routing_call}; use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat}; use thiserror::Error; @@ -194,13 +194,17 @@ impl Route { }) } - /// Completes routing-time calls without serving the answer target. + /// Completes routing-time calls and prepares the selected request without serving it. pub async fn decide(&self, request: Request) -> Result { - drive(Arc::clone(&self.algorithm), request, |call| { - serve_decision_dependency(self.clients.clone(), call) + let mut outcome = drive(Arc::clone(&self.algorithm), request, |call| { + serve_routing_call(self.clients.clone(), call) }) .await - .map_err(Into::into) + .map_err(RunnerError::from)?; + outcome.request = self + .clients + .prepare_completion_request(outcome.request, &outcome.selected_model_id); + Ok(outcome) } /// Counts tokens using the configured Anthropic-capable target. @@ -209,6 +213,9 @@ impl Route { .count_tokens_target .as_ref() .ok_or(RunnerError::CountTokensUnsupported)?; + let request = self + .clients + .prepare_completion_request(request, &target.model); target .client .count_tokens(&target.model, request) @@ -216,48 +223,3 @@ impl Route { .map_err(Into::into) } } - -async fn serve_decision_dependency(clients: ClientRouter, call: CallModel) -> libsy::Result<()> { - let mut result = Err(LibsyError::NoTargets); - for (index, model) in call.models.iter().enumerate() { - // The driver stamps only the first candidate, so every fallback must replace it. - let mut request = call.request.clone(); - request.llm_request.model = Some(model.to_string()); - let response = match clients.route(model) { - Ok(client) => client.call(request).await, - Err(source) => Err(source), - }; - match response { - Ok(response) => { - result = Ok(response); - break; - } - Err(source) => { - let try_next = index + 1 < call.models.len() && eligible_routing_fallback(&source); - result = Err(LibsyError::client_call(model.clone(), source)); - if !try_next { - break; - } - } - } - } - call.respond(result) -} - -/// Whether a routing-time candidate failure may fall through to the next model. -fn eligible_routing_fallback(error: &LlmClientError) -> bool { - match error { - LlmClientError::ContextWindowExceeded { .. } - | LlmClientError::Transport { .. } - | LlmClientError::Timeout { .. } => true, - LlmClientError::UpstreamHttp { status, .. } => { - matches!( - *status, - reqwest::StatusCode::FORBIDDEN - | reqwest::StatusCode::REQUEST_TIMEOUT - | reqwest::StatusCode::TOO_MANY_REQUESTS - ) || status.is_server_error() - } - _ => false, - } -} diff --git a/crates/switchyard-runner/tests/route.rs b/crates/switchyard-runner/tests/route.rs index 71ce57c7e..31e5f9bef 100644 --- a/crates/switchyard-runner/tests/route.rs +++ b/crates/switchyard-runner/tests/route.rs @@ -9,8 +9,8 @@ use async_trait::async_trait; use futures_util::StreamExt; use switchyard_llm_client::{ClientRouter, RunObservation}; use switchyard_protocol::{ - LlmClientError, LlmResponse, ModelId, Request, Response, RoutedLlmClient, text_request, - text_response, + ContentBlock, LlmClientError, LlmResponse, ModelId, Request, Response, RoutedLlmClient, + text_request, text_response, }; use switchyard_runner::{AlgorithmSpec, ModelCapabilities, Route}; @@ -29,7 +29,7 @@ impl RoutedLlmClient for StubClient { } } -fn plugin_route(client: Arc) -> Route { +fn plugin_route(client: Arc, target_prompt: Option<&str>) -> Route { let spec = AlgorithmSpec::Passthrough { target: "semantic-target".to_string(), subagents: None, @@ -41,11 +41,19 @@ fn plugin_route(client: Arc) -> Route { let algorithm = spec .build("switchyard", &targets) .expect("identity target map should build"); - let clients = ClientRouter::new( + let mut clients = ClientRouter::new( BTreeMap::from([(ModelId::from("semantic-target"), client)]) .into_iter() .collect(), ); + if let Some(prompt) = target_prompt { + clients = clients.with_target_prompts( + [(ModelId::from("semantic-target"), prompt.to_string())] + .into_iter() + .collect(), + None, + ); + } Route::new( algorithm, clients, @@ -58,7 +66,7 @@ fn plugin_route(client: Arc) -> Route { #[tokio::test] async fn plugin_shaped_route_executes_without_runner_model_or_toml() { - let route = plugin_route(Arc::new(StubClient)); + let route = plugin_route(Arc::new(StubClient), None); let observations = Arc::new(Mutex::new(Vec::new())); let observer = { let observations = Arc::clone(&observations); @@ -101,6 +109,25 @@ async fn plugin_shaped_route_executes_without_runner_model_or_toml() { ); } +#[tokio::test] +async fn decision_prepares_the_selected_target_request() { + let route = plugin_route(Arc::new(StubClient), Some("target prompt")); + let request = Request { + llm_request: text_request(None, "hello"), + ..Request::default() + }; + + let outcome = route + .decide(request) + .await + .expect("passthrough decision should succeed"); + + assert!(matches!( + &outcome.request.llm_request.instructions[0].content[0], + ContentBlock::Text { text } if text == "target prompt" + )); +} + struct LazyStreamClient { polls: Arc, } @@ -124,9 +151,12 @@ impl RoutedLlmClient for LazyStreamClient { #[tokio::test] async fn route_returns_stream_without_polling_it() { let polls = Arc::new(AtomicUsize::new(0)); - let route = plugin_route(Arc::new(LazyStreamClient { - polls: Arc::clone(&polls), - })); + let route = plugin_route( + Arc::new(LazyStreamClient { + polls: Arc::clone(&polls), + }), + None, + ); let request = Request { llm_request: text_request(None, "hello"), ..Request::default()