From 55d2a22823a6acc877b3dd24aedd4e34bf2a7b6f Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 17 Aug 2026 11:57:29 -0700 Subject: [PATCH 1/3] feat(translation): prepare requests for routed targets Signed-off-by: Alex Fournier --- crates/switchyard-translation/src/lib.rs | 3 +- crates/switchyard-translation/src/util.rs | 27 ++++++++- .../tests/request_translation.rs | 60 ++++++++++++++++++- 3 files changed, 87 insertions(+), 3 deletions(-) diff --git a/crates/switchyard-translation/src/lib.rs b/crates/switchyard-translation/src/lib.rs index d07a3497d..10da7ab08 100644 --- a/crates/switchyard-translation/src/lib.rs +++ b/crates/switchyard-translation/src/lib.rs @@ -31,5 +31,6 @@ pub use llm::*; pub use policy::*; pub use stream::*; pub use util::{ - PRESERVATION_METADATA_KEY, normalize_anthropic_tool_use_ids, sanitize_anthropic_tool_use_id, + PRESERVATION_METADATA_KEY, normalize_anthropic_tool_use_ids, prepare_request_for_target, + sanitize_anthropic_tool_use_id, }; diff --git a/crates/switchyard-translation/src/util.rs b/crates/switchyard-translation/src/util.rs index 9fcf769fd..11a5492d1 100644 --- a/crates/switchyard-translation/src/util.rs +++ b/crates/switchyard-translation/src/util.rs @@ -6,11 +6,12 @@ use std::collections::BTreeMap; use serde_json::{Map, Value, json}; +use switchyard_protocol::ModelId; use crate::diagnostic::TranslationDiagnostic; use crate::error::{Result, TranslationError}; use crate::format::FormatId; -use crate::llm::{ContentBlock, LlmRequest, Message, PreservationMetadata}; +use crate::llm::{ContentBlock, InstructionBlock, LlmRequest, Message, PreservationMetadata, Role}; use crate::policy::{ LossyConversionPolicy, PreservationPolicy, TranslationPolicy, UnknownFieldPolicy, }; @@ -271,6 +272,30 @@ pub fn exact_preserved_response( .flatten() } +/// Applies a selected target model and optionally prepends its system prompt. +/// +/// Adding a prompt invalidates preserved provider bodies because they predate the mutation. +/// Call this once per candidate using a request that has not already received a target prompt. +pub fn prepare_request_for_target( + request: &mut LlmRequest, + target: &ModelId, + prompt: Option<&str>, +) { + request.model = Some(target.to_string()); + if let Some(prompt) = prompt { + request.instructions.insert( + 0, + InstructionBlock { + role: Role::System, + content: vec![ContentBlock::Text { + text: prompt.to_string(), + }], + }, + ); + request.preservation.requests.clear(); + } +} + /// Embeds preservation metadata into a translated wire body when requested. pub fn embed_preservation( mut body: Value, diff --git a/crates/switchyard-translation/tests/request_translation.rs b/crates/switchyard-translation/tests/request_translation.rs index b7c4121a9..8310e7439 100644 --- a/crates/switchyard-translation/tests/request_translation.rs +++ b/crates/switchyard-translation/tests/request_translation.rs @@ -9,11 +9,69 @@ use pretty_assertions::assert_eq; use serde_json::{Value, json}; use switchyard_translation::{ LossyConversionPolicy, TranslationEngine, TranslationPolicy, WireFormat, + prepare_request_for_target, }; use common::{REASONING_MODEL, normalized_policy, shell_tool_call}; -type TestResult = std::result::Result<(), Box>; +type TestResult = std::result::Result>; + +// A target prompt makes every preserved provider body stale. +#[test] +fn preparing_a_target_prompt_invalidates_exact_replay() -> TestResult { + let engine = TranslationEngine::default(); + let policy = TranslationPolicy::default(); + let body = json!({ + "model": "route", + "messages": [ + {"role": "system", "name": "caller", "content": "client prompt"}, + {"role": "user", "content": "hi"} + ] + }); + let mut request = engine + .decode_request(WireFormat::OpenAiChat, &body, &policy)? + .request; + + prepare_request_for_target( + &mut request, + &"selected/model".into(), + Some("target prompt"), + ); + + assert!(request.preservation.requests.is_empty()); + let encoded = engine + .encode_request(WireFormat::OpenAiChat, &request, &policy)? + .body; + assert_eq!(encoded["model"], "selected/model"); + assert_eq!(encoded["messages"][0]["content"], "target prompt"); + assert_eq!(encoded["messages"][1]["content"], "client prompt"); + assert!(encoded["messages"][1].get("name").is_none()); + Ok(()) +} + +// Stamping only the normalized target does not invalidate exact replay. +#[test] +fn preparing_without_a_prompt_preserves_exact_replay() -> TestResult { + let engine = TranslationEngine::default(); + let policy = TranslationPolicy::default(); + let body = json!({ + "model": "route", + "messages": [{"role": "user", "content": "hi"}], + "provider_field": true + }); + let mut request = engine + .decode_request(WireFormat::OpenAiChat, &body, &policy)? + .request; + + prepare_request_for_target(&mut request, &"selected/model".into(), None); + + assert_eq!(request.model.as_deref(), Some("selected/model")); + assert_eq!( + request.preservation.requests[&WireFormat::OpenAiChat.into()], + body + ); + Ok(()) +} // Verifies Anthropic-only request fields are dropped or mapped for OpenAI Chat. #[test] From e9a30fba57f6affcfe245ba39dda26402fa19eba Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 17 Aug 2026 11:59:55 -0700 Subject: [PATCH 2/3] feat(libsy): prepare requests for routed candidates Signed-off-by: Alex Fournier --- Cargo.lock | 2 + crates/libsy-llm-client/src/run.rs | 59 +++-- crates/libsy/Cargo.toml | 1 + crates/libsy/README.md | 9 + crates/libsy/src/algorithms/advisor_gate.rs | 2 +- .../src/algorithms/advisor_gate/tests.rs | 21 +- crates/libsy/src/algorithms/llm_class.rs | 40 ++- crates/libsy/src/algorithms/stage.rs | 60 +++-- crates/libsy/src/algorithms/util/prompts.rs | 79 ++---- crates/libsy/src/core.rs | 2 + crates/libsy/src/core/algorithm.rs | 250 +++++++++++++++++- crates/libsy/src/core/target_prompts.rs | 32 +++ crates/libsy/src/core/testing.rs | 15 +- crates/libsy/src/lib.rs | 7 +- crates/switchyard-py/Cargo.toml | 1 + crates/switchyard-py/src/libsy_bindings.rs | 76 ++++-- switchyard_rust/libsy.py | 6 + tests/test_libsy_minimal_bindings.py | 19 +- 18 files changed, 527 insertions(+), 154 deletions(-) create mode 100644 crates/libsy/src/core/target_prompts.rs diff --git a/Cargo.lock b/Cargo.lock index 405dda9da..82d822a87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2280,6 +2280,7 @@ dependencies = [ "serde", "serde_json", "switchyard-protocol", + "switchyard-translation", "thiserror 2.0.18", "tokio", "tokio-stream", @@ -2331,6 +2332,7 @@ version = "0.2.0" dependencies = [ "futures", "http", + "parking_lot", "pyo3", "pyo3-async-runtimes", "pythonize", diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index 73d897c39..6398a9bf8 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -62,17 +62,17 @@ pub async fn run( .and_then(|outcome| outcome.response.as_ref()) .and_then(Response::served_model); emit_routing_observations(&observer, &routing_observations, answered_model); - let outcome = outcome?; + let mut outcome = outcome?; let overhead = run_started.elapsed(); metrics::record_routing_overhead(&algorithm_name, overhead); - let selected_model_id = outcome.selected_model_id; - let (result, answer_duration) = if let Some(response) = outcome.response { + let selected_model_id = outcome.selected_model_id.clone(); + let (result, answer_duration) = if let Some(response) = outcome.response.take() { (Ok(response), None) } else { let mut models = Vec::with_capacity(1 + outcome.fallback_models.len()); models.push(selected_model_id.clone()); - models.extend(outcome.fallback_models); + models.extend(outcome.fallback_models.iter().cloned()); let answer_started = Instant::now(); let observe = |observation| { if let Some(observer) = &observer { @@ -82,8 +82,8 @@ pub async fn run( let result = call_first_available( &clients, &algorithm_name, - &outcome.request, &models, + move |target| outcome.request_for(target), &observe, ) .await; @@ -142,8 +142,8 @@ async fn serve( let result = call_first_available( &clients, &call.algorithm, - &call.request, &call.models, + |target| call.request_for(target), &observe, ) .await; @@ -154,12 +154,12 @@ async fn serve( async fn call_first_available( clients: &ClientRouter, algorithm: &str, - request: &Request, models: &[ModelId], + request_for: impl Fn(&ModelId) -> Result + Send, observe: &(dyn Fn(LlmCallObservation) + Send + Sync), ) -> Result { for (index, target) in models.iter().enumerate() { - let request = request_for(request, target); + let request = request_for(target)?; match call_one( clients, target, @@ -298,13 +298,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 @@ -379,10 +372,10 @@ mod tests { use async_trait::async_trait; use futures::StreamExt; use http::StatusCode; - use switchyard_libsy::{Driver, RoutingOutcome}; + use switchyard_libsy::{Driver, RoutingOutcome, TargetPrompts, with_target_prompts}; 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}; @@ -429,7 +422,7 @@ mod tests { request: Request, ) -> Result { let response = driver - .call_model(request.clone(), vec![self.model.clone()]) + .call_answer_model(request.clone(), self.model.clone()) .await?; Ok(RoutingOutcome::answered( self.model.clone(), @@ -449,6 +442,7 @@ mod tests { struct CandidateClient { calls: Mutex>, + prompts: Mutex>>, first: FirstOutcome, } @@ -457,6 +451,18 @@ 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.prompts.lock().push( + request + .llm_request + .instructions + .iter() + .flat_map(|instruction| &instruction.content) + .filter_map(|block| match block { + ContentBlock::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(), + ); if model == "weak" { return match self.first { FirstOutcome::ContextWindow => Err(LlmClientError::ContextWindowExceeded { @@ -521,11 +527,16 @@ mod tests { ) -> (Arc, Result<(ModelId, Response)>) { let client = Arc::new(CandidateClient { calls: Mutex::new(Vec::new()), + prompts: Mutex::new(Vec::new()), first, }); - let algorithm = Arc::new(CandidateAlgorithm { + let inner: Arc = Arc::new(CandidateAlgorithm { models: vec!["weak".into(), "strong".into()], }); + let prompts = TargetPrompts::default() + .with("weak", "weak prompt") + .with("strong", "strong prompt"); + let algorithm = with_target_prompts(inner, prompts); let result = run( algorithm, ClientRouter::single(client.clone()), @@ -540,6 +551,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()), + prompts: Mutex::new(Vec::new()), first: FirstOutcome::StreamSuccess, }); let observations = Arc::new(Mutex::new(Vec::new())); @@ -621,6 +633,13 @@ mod tests { &*client.calls.lock(), &[ModelId::from("weak"), "strong".into()] ); + assert_eq!( + &*client.prompts.lock(), + &[ + vec!["weak prompt".to_string()], + vec!["strong prompt".to_string()] + ] + ); assert_eq!( response .llm_response diff --git a/crates/libsy/Cargo.toml b/crates/libsy/Cargo.toml index 369cb6160..1749400c4 100644 --- a/crates/libsy/Cargo.toml +++ b/crates/libsy/Cargo.toml @@ -32,6 +32,7 @@ parking_lot.workspace = true rand.workspace = true regex.workspace = true switchyard-protocol.workspace = true +switchyard-translation.workspace = true thiserror.workspace = true tokio.workspace = true tokio-stream = "0.1" diff --git a/crates/libsy/README.md b/crates/libsy/README.md index b69c69754..fd0f81f14 100644 --- a/crates/libsy/README.md +++ b/crates/libsy/README.md @@ -38,6 +38,15 @@ fallbacks, rewritten request, and an optional response already produced while ro makes no network calls itself — `switchyard-llm-client`'s `run` is a ready-made consumer that drives the stream and performs the terminal answer call, retries, and fallback over HTTP. +[`RoutingOutcome`]'s `request` field is ready for the selected answer target. A custom host +trying the selected target or a fallback should call [`RoutingOutcome::request_for`]; that +prepares the candidate's model and any prompt configured with [`with_target_prompts`] as one +operation. + +Routing-time [`CallModel`] requests are likewise ready for their first candidate. Hosts trying +a later classifier or judge candidate should use [`CallModel::request_for`] so exact provider +bodies receive the candidate model together with the normalized request. + The provider-neutral [`Request`], [`Response`], [`Usage`], and [`LlmResponse`] contracts come from `switchyard-protocol`. diff --git a/crates/libsy/src/algorithms/advisor_gate.rs b/crates/libsy/src/algorithms/advisor_gate.rs index 63cabdd16..e0bf093ef 100644 --- a/crates/libsy/src/algorithms/advisor_gate.rs +++ b/crates/libsy/src/algorithms/advisor_gate.rs @@ -334,7 +334,7 @@ impl AdvisorGate { // Gated phase: generate the turn once, fully buffered, so the gate // can inspect it before the client sees anything. let response = driver - .call_model(request.clone(), vec![self.executor.clone()]) + .call_answer_model(request.clone(), self.executor.clone()) .await?; let turn = buffer_turn(self.executor.as_str(), response).await?; diff --git a/crates/libsy/src/algorithms/advisor_gate/tests.rs b/crates/libsy/src/algorithms/advisor_gate/tests.rs index e57495668..7c55b164c 100644 --- a/crates/libsy/src/algorithms/advisor_gate/tests.rs +++ b/crates/libsy/src/algorithms/advisor_gate/tests.rs @@ -16,6 +16,7 @@ use switchyard_protocol::{ use super::transcript::{NO_TEXT_PLACEHOLDER, TRUNCATION_MARKER, middle_drop}; use super::*; use crate::core::testing::{reply, test_drive}; +use crate::{TargetPrompts, with_target_prompts}; const EXECUTOR: &str = "executor"; const ADVISOR: &str = "advisor"; @@ -268,7 +269,12 @@ async fn tool_call_turn_replays_without_review() { #[tokio::test] async fn approved_terminal_turn_returns_buffered_body() { let script = Script::new(); - let gate = gate(AdvisorGateConfig::default()); + let gate = with_target_prompts( + gate(AdvisorGateConfig::default()), + TargetPrompts::default() + .with(EXECUTOR, "executor prompt") + .with(ADVISOR, "answer-only advisor prompt"), + ); let serve = script.serve("APPROVE", |_| reply("all done")); let (selected_model, response) = test_drive(gate, task_request(), serve) .await @@ -279,6 +285,19 @@ async fn approved_terminal_turn_returns_buffered_body() { ); assert_eq!(completion_text(&agg_of(response).await), "all done"); assert_eq!(selected_model, EXECUTOR); + let executor = script.call(0); + assert_eq!( + executor.llm_request.instructions[0].content, + vec![ContentBlock::Text { + text: "executor prompt".to_string(), + }] + ); + let advisor = script.call(1); + assert!(!advisor.llm_request.instructions.iter().any(|instruction| { + instruction.content.iter().any(|block| { + matches!(block, ContentBlock::Text { text } if text == "answer-only advisor prompt") + }) + })); } #[tokio::test] diff --git a/crates/libsy/src/algorithms/llm_class.rs b/crates/libsy/src/algorithms/llm_class.rs index 589a20158..633bafa62 100644 --- a/crates/libsy/src/algorithms/llm_class.rs +++ b/crates/libsy/src/algorithms/llm_class.rs @@ -522,7 +522,7 @@ impl Classifier for EscalationClassifier { "escalation classifier selected efficient tier" ); let efficient_response = match driver - .call_model(request.clone(), vec![self.efficient.clone()]) + .call_answer_model(request.clone(), self.efficient.clone()) .await { Ok(r) => r, @@ -1777,6 +1777,16 @@ mod tests { } } + /// Reports whether a request contains `expected` as an instruction text block. + fn has_instruction(request: &Request, expected: &str) -> bool { + request + .llm_request + .instructions + .iter() + .flat_map(|instruction| &instruction.content) + .any(|block| matches!(block, ContentBlock::Text { text } if text == expected)) + } + /// Returns a stream that emits partial content before failing during aggregation. fn streamed_then_error(error: LlmClientError) -> Response { Response { @@ -1814,10 +1824,28 @@ mod tests { // Judge: no escalation. Expect the efficient response to be returned directly. let judge = Queue::new([r#"{"escalate":false,"reason":"progressing"}"#]); let model = Queue::new(["efficient answer"]); - let router = escalation_router()?; + let replies = queued(model, judge); + let prompted = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&prompted); + let serve = move |target: ModelId, request: Request| { + let expected = if target == "judge" { + "answer-only judge prompt" + } else { + "efficient prompt" + }; + recorded + .lock() + .push((target.clone(), has_instruction(&request, expected))); + replies.serve(target, request) + }; + let router = crate::with_target_prompts( + escalation_router()?, + crate::TargetPrompts::default() + .with("efficient", "efficient prompt") + .with("judge", "answer-only judge prompt"), + ); - let (selected_model, response) = - test_drive(router, classify_request(), queued(model, judge)).await?; + let (selected_model, response) = test_drive(router, classify_request(), serve).await?; // The efficient model is the serving target, and the response comes from its call. assert_eq!(selected_model, "efficient"); @@ -1825,6 +1853,10 @@ mod tests { response.llm_response.as_agg().map(completion_text), Some("efficient answer".to_string()) ); + assert_eq!( + &*prompted.lock(), + &[(ModelId::from("efficient"), true), ("judge".into(), false)] + ); Ok(()) } diff --git a/crates/libsy/src/algorithms/stage.rs b/crates/libsy/src/algorithms/stage.rs index 00ec3e59e..4ea0207d7 100644 --- a/crates/libsy/src/algorithms/stage.rs +++ b/crates/libsy/src/algorithms/stage.rs @@ -19,7 +19,6 @@ use async_trait::async_trait; use super::fall_through::{DefaultTarget, FallThrough}; use super::llm_class::{LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig}; -use super::util::prompts::{SystemPromptProcessor, TargetPrompts}; use super::util::stage::{ DecisionSource, HandoffNoteConfig, PickerMode, StageClassifier, StageTargets, record_decision_source, record_routing_decision, @@ -28,7 +27,7 @@ use super::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignalProcessor}; use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::{Classification, Classifier}; use crate::core::state::State; -use crate::{LibsyError, Result}; +use crate::{LibsyError, Result, TargetPrompts}; use switchyard_protocol::{ModelId, Request, Response}; /// Telemetry name for a router this module assembles. @@ -118,6 +117,7 @@ impl StageRouterConfig { /// the picker's default tier closes the cascade so a turn is never left unrouted. pub struct StageRouter { route: FallThrough, + tier_prompts: Option>, } impl StageRouter { @@ -126,9 +126,16 @@ impl StageRouter { /// routing destination. /// /// Errors if either threshold in `config` is outside `[0.0, 1.0]`. - pub fn new(capable: ModelId, efficient: ModelId, config: StageRouterConfig) -> Result { + pub fn new( + capable: ModelId, + efficient: ModelId, + mut config: StageRouterConfig, + ) -> Result { + let tier_prompts = std::mem::take(&mut config.tier_prompts); + let tier_prompts = (!tier_prompts.is_empty()).then(|| Arc::new(tier_prompts)); Ok(Self { route: build_route(capable, efficient, config)?, + tier_prompts, }) } } @@ -144,7 +151,11 @@ impl Algorithm for StageRouter { driver: Driver, request: Request, ) -> Result { - self.route.execute(driver, request).await + let outcome = self.route.execute(driver, request).await?; + Ok(match &self.tier_prompts { + Some(prompts) => outcome.with_target_prompts(Arc::clone(prompts)), + None => outcome, + }) } } @@ -200,10 +211,6 @@ fn build_route( inner: Arc::new(DefaultTarget::new(fall_open)), source: DecisionSource::FallOpen, })); - // Runs on the post-decision hook, so it applies to the target the cascade - // settled on, whichever classifier picked it. With no prompts configured it - // is a no-op, so there is nothing to branch on. - router = router.with_processor(Arc::new(SystemPromptProcessor::new(config.tier_prompts))); Ok(router) } @@ -337,6 +344,7 @@ mod tests { #[derive(Clone, Debug)] struct Call { target: String, + instructions: Vec, messages: Vec, } @@ -367,6 +375,16 @@ mod tests { let target = target.to_string(); recorder.calls.lock().push(Call { target: target.clone(), + instructions: request + .llm_request + .instructions + .iter() + .flat_map(|instruction| &instruction.content) + .filter_map(|block| match block { + ContentBlock::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(), messages: request .llm_request .messages @@ -493,22 +511,28 @@ mod tests { } #[tokio::test] - async fn the_judge_decides_a_turn_the_signals_leave_undecided() -> Result<()> { + async fn the_judge_selects_a_target_without_receiving_its_answer_prompt() -> Result<()> { let recorder = Arc::new(Recorder::default()); - let router = recording_router(config_with_judge(&recorder, 0.1))?; + let mut config = config_with_judge(&recorder, 0.1); + config.tier_prompts = TargetPrompts::default().with("strong", "answer-only strong prompt"); + let router = recording_router(config)?; - let (selected_model, _) = - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + let (selected_model, _) = test_drive(router, turn_request(false), recorder.serve()).await?; let calls = recorder.calls.lock(); + let Some(judge) = calls.iter().find(|call| call.target == JUDGE) else { + panic!("the judge was never called"); + }; assert!( - calls.iter().any(|call| call.target == JUDGE), - "the judge should be recorded as a routing side call" - ); - assert!( - calls.iter().any(|call| call.target == "strong"), - "the selected target should be recorded as an answer call" + !judge + .instructions + .iter() + .any(|instruction| instruction == "answer-only strong prompt") ); + let Some(strong) = calls.iter().find(|call| call.target == "strong") else { + panic!("the selected target was never called"); + }; + assert_eq!(strong.instructions, ["answer-only strong prompt"]); drop(calls); assert_eq!(selected_model, "strong"); Ok(()) diff --git a/crates/libsy/src/algorithms/util/prompts.rs b/crates/libsy/src/algorithms/util/prompts.rs index 37314d43c..fbb4d7591 100644 --- a/crates/libsy/src/algorithms/util/prompts.rs +++ b/crates/libsy/src/algorithms/util/prompts.rs @@ -1,35 +1,20 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Adding text to a request on its way to the model it was routed to. +//! Text added to requests by routing policies. //! -//! Two shapes, both target-agnostic — any algorithm routing between named -//! targets can use them, and neither writes anything back into the caller's -//! conversation: +//! [`append_note`] adds a one-turn conversation note. [`TargetPrompts`] stores standing +//! instructions by answer target; [`SystemPromptProcessor`] remains the low-level processor for +//! existing single-candidate fall-through compositions. //! -//! * [`append_note`] — a one-off note in the conversation itself, for telling -//! the model something about *this* turn. -//! * [`SystemPromptProcessor`] — standing instructions per target, applied on -//! every turn that target serves. -//! -//! Which text, and when, is the caller's policy; this module only knows how to -//! place it so the provider accepts it and the prompt cache survives. -//! -//! **Anything added here must call [`drop_exact_replay`].** Both shapes above -//! mutate the normalized request, and a codec asked to encode for the format the -//! request arrived in replays the body captured at decode instead of reading that -//! request — so an addition that leaves exact replay in place never reaches the -//! model. This is not enforced: a future processor that mutates the request and -//! forgets the call reintroduces SWITCH-1224, silently and without a failing -//! test. - -use std::collections::BTreeMap; +//! Any request mutation must also update or discard exact preserved provider bodies. Otherwise a +//! same-format encode can replay the body captured at decode and silently omit the mutation. use async_trait::async_trait; -use switchyard_protocol::{ContentBlock, InstructionBlock, Message, ModelId, Request, Role}; +use switchyard_protocol::{ContentBlock, Message, Request, Role}; -use crate::Result; use crate::core::processor::{Event, Processor}; +use crate::{Result, TargetPrompts}; /// Appends `note` to the request as conversation text. /// @@ -70,33 +55,10 @@ pub(crate) fn drop_exact_replay(request: &mut Request) { request.llm_request.preservation.requests.clear(); } -/// System prompts keyed by routing target. A target left unset is routed -/// untouched. -#[derive(Clone, Debug, Default)] -pub struct TargetPrompts { - by_target: BTreeMap, -} - -impl TargetPrompts { - /// Hand `target` this prompt on every turn it serves. - pub fn with(mut self, target: impl Into, prompt: impl Into) -> Self { - self.by_target.insert(target.into(), prompt.into()); - self - } - - /// The prompt configured for `target`, if any. - pub fn get(&self, target: &ModelId) -> Option<&str> { - self.by_target.get(target).map(String::as_str) - } - - /// Whether any target has a prompt, so a caller can skip wiring the - /// processor when none does. - pub fn is_empty(&self) -> bool { - self.by_target.is_empty() - } -} - /// Prepends the routed target's system prompt to the outbound request. +/// +/// This low-level processor remains for existing single-candidate `FallThrough` +/// compositions. Fallback-capable algorithms should use [`crate::with_target_prompts`]. pub struct SystemPromptProcessor { prompts: TargetPrompts, } @@ -124,18 +86,11 @@ impl Processor for SystemPromptProcessor { let Some(prompt) = self.prompts.get(selected_model_id) else { return Ok(()); }; - // Ahead of the client's own instructions, so this framing is what the - // model reads first. - request.llm_request.instructions.insert( - 0, - InstructionBlock { - role: Role::System, - content: vec![ContentBlock::Text { - text: prompt.to_string(), - }], - }, + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + selected_model_id, + Some(prompt), ); - drop_exact_replay(request); Ok(()) } } @@ -143,7 +98,7 @@ impl Processor for SystemPromptProcessor { #[cfg(test)] mod tests { use super::*; - use switchyard_protocol::{LlmRequest, ModelId, ToolResult, text_request}; + use switchyard_protocol::{InstructionBlock, LlmRequest, ModelId, ToolResult, text_request}; const NOTE: &str = "recovering from an error"; const STRONG_PROMPT: &str = "diagnose before you edit"; @@ -296,7 +251,7 @@ mod tests { assert_eq!(instructions(&request), vec![expected]); assert!( !replays_exactly(&request), - "{target}: a same-format hop would replay the body captured before the prompt" + "{target}: a same-format hop must rebuild after adding the prompt" ); } Ok(()) diff --git a/crates/libsy/src/core.rs b/crates/libsy/src/core.rs index b22ccc2ee..224357651 100644 --- a/crates/libsy/src/core.rs +++ b/crates/libsy/src/core.rs @@ -12,3 +12,5 @@ pub mod algorithm; pub mod classifier; pub mod processor; pub mod state; +mod target_prompts; +pub use target_prompts::TargetPrompts; diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 47ddfb668..502684669 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -19,8 +19,9 @@ use tracing::Instrument; /// [`switchyard_protocol::LlmResponseStreamEvent`] is its host/algorithm envelope; and /// [`switchyard_protocol::LlmResponse`] carries either a live /// [`switchyard_protocol::LlmResponseStream`] or the terminal aggregate. -use switchyard_protocol::{ModelId, Request, Response}; +use switchyard_protocol::{LlmRequest, ModelId, Request, Response}; +use super::TargetPrompts; use crate::{DriverError, LibsyError, Result, observability}; /// A boxed, `Send` stream of [`Step`]s — the output of @@ -28,6 +29,10 @@ use crate::{DriverError, LibsyError, Result, observability}; /// `Arc` object-safe. pub type StepStream = Pin> + Send>>; +fn target_prompt<'a>(prompts: &'a [Arc], target: &ModelId) -> Option<&'a str> { + prompts.iter().find_map(|prompts| prompts.get(target)) +} + /// An offloaded model call, surfaced inside [`Step::CallModel`]. /// /// The host reads the public fields, performs (or delegates) the model call, and fulfills it @@ -50,6 +55,14 @@ pub struct CallModel { } impl CallModel { + /// Build the request for one candidate in this routing-time call. + pub fn request_for(&self, target: &ModelId) -> Result { + ensure_model_is_target(&self.models, target)?; + let mut request = self.request.clone(); + switchyard_translation::prepare_request_for_target(&mut request.llm_request, target, None); + Ok(request) + } + /// 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. @@ -61,6 +74,10 @@ impl CallModel { } /// The terminal result of routing. +/// +/// Hosts making the answer call should use [`request_for`](Self::request_for) for the selected +/// model and every fallback. This keeps target-specific request preparation inside libsy. +#[non_exhaustive] pub struct RoutingOutcome { /// The model selected by the algorithm and tried first by the client. pub selected_model_id: ModelId, @@ -70,6 +87,11 @@ pub struct RoutingOutcome { pub request: Request, /// A response produced while routing, or `None` when the client must make the answer call. pub response: Option, + // Request state before the selected target's prompt was applied. Retained only when a + // prompted selection has fallbacks, so a fallback cannot inherit the selected prompt. + base_llm_request: Option>, + // Outer prompt layers precede inner layers, so deployment policy overrides router defaults. + target_prompts: Vec>, } impl RoutingOutcome { @@ -81,25 +103,82 @@ impl RoutingOutcome { fallback_models: Vec, mut request: Request, ) -> Self { - request.llm_request.model = Some(selected_model_id.to_string()); + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + &selected_model_id, + None, + ); Self { selected_model_id, fallback_models, request, response: None, + base_llm_request: None, + target_prompts: Vec::new(), } } /// Algorithm generated the response as part of the routing decision. Here it is. /// The `request` will have the `selected_model_id` written into it by this function. pub fn answered(selected_model_id: ModelId, mut request: Request, response: Response) -> Self { - request.llm_request.model = Some(selected_model_id.to_string()); + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + &selected_model_id, + None, + ); Self { selected_model_id, fallback_models: Vec::new(), request, response: Some(response), + base_llm_request: None, + target_prompts: Vec::new(), + } + } + + /// Build the answer request for the selected target or one of its fallbacks. + pub fn request_for(&self, target: &ModelId) -> Result { + if target != &self.selected_model_id && !self.fallback_models.contains(target) { + return Err(LibsyError::TargetNotFound { + target: target.clone(), + }); + } + if target == &self.selected_model_id { + return Ok(self.request.clone()); + } + let mut request = match &self.base_llm_request { + Some(base) => Request { + llm_request: base.as_ref().clone(), + raw_request: self.request.raw_request.clone(), + metadata: self.request.metadata.clone(), + }, + None => self.request.clone(), + }; + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + target, + target_prompt(&self.target_prompts, target), + ); + Ok(request) + } + + pub(crate) fn with_target_prompts(mut self, prompts: Arc) -> Self { + self.target_prompts.insert(0, prompts); + self + } + + fn prepare_selected_request(&mut self) { + let Some(prompt) = target_prompt(&self.target_prompts, &self.selected_model_id) else { + return; + }; + if !self.fallback_models.is_empty() { + self.base_llm_request = Some(Box::new(self.request.llm_request.clone())); } + switchyard_translation::prepare_request_for_target( + &mut self.request.llm_request, + &self.selected_model_id, + Some(prompt), + ); } } @@ -109,6 +188,8 @@ pub struct Driver { step_tx: mpsc::Sender>, /// The owning algorithm's telemetry label, stamped onto every call this driver publishes. algorithm: String, + // Prompt policy is shared with answer calls made while an algorithm is still routing. + target_prompts: Vec>, } impl Driver { @@ -124,11 +205,17 @@ impl Driver { Self { step_tx, algorithm: algorithm.to_string(), + target_prompts: Vec::new(), }, step_rx, ) } + pub(crate) fn with_target_prompts(mut self, prompts: Arc) -> Self { + self.target_prompts.push(prompts); + self + } + /// Publish a model call and await the consumer's response. /// /// Errors if the stream is closed or the call failed. @@ -137,13 +224,45 @@ impl Driver { /// response resolves when its stream handle arrives); latency, outcome, and /// token usage are recorded when it resolves. The provider call itself is the /// host's, and is instrumented by whoever makes it. + pub async fn call_model(&self, mut request: Request, models: Vec) -> Result { + let Some(selected_model_id) = models.first().cloned() else { + return Err(LibsyError::NoTargets); + }; + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + &selected_model_id, + None, + ); + self.call_prepared_model(request, models, selected_model_id) + .await + } + + /// Publish a single model call whose response may become the client-visible answer. + /// + /// Classifier and judge calls should use [`call_model`](Self::call_model), which deliberately + /// does not receive an answer target's prompt. + pub async fn call_answer_model( + &self, + mut request: Request, + model: ModelId, + ) -> Result { + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + &model, + target_prompt(&self.target_prompts, &model), + ); + let selected_model_id = model.clone(); + self.call_prepared_model(request, vec![model], selected_model_id) + .await + } + #[tracing::instrument( target = "libsy", name = "libsy.llm_call", skip_all, fields( algorithm = self.algorithm, - selected_model = %models.first().map(ModelId::as_str).unwrap_or("NoTargets"), + selected_model = %selected_model_id, openinference.span.kind = "CHAIN", outcome = tracing::field::Empty, error = tracing::field::Empty, @@ -153,11 +272,12 @@ impl Driver { reasoning_tokens = tracing::field::Empty, ) )] - pub async fn call_model(&self, mut request: Request, models: Vec) -> Result { - let Some(selected_model_id) = models.first().cloned() else { - return Err(LibsyError::NoTargets); - }; - request.llm_request.model = Some(selected_model_id.to_string()); + async fn call_prepared_model( + &self, + request: Request, + models: Vec, + selected_model_id: ModelId, + ) -> Result { let started = Instant::now(); let (reply, response) = oneshot::channel::>(); let call = CallModel { @@ -357,7 +477,9 @@ pub trait Algorithm: Send + Sync + 'static { fn name(&self) -> &str; /// Run one request to completion: make routing-time model calls with - /// [`Driver::call_model`] and return the terminal [`RoutingOutcome`]. + /// [`Driver::call_model`] and return the terminal [`RoutingOutcome`]. A model call whose + /// response may be returned in [`RoutingOutcome::response`] must use + /// [`Driver::call_answer_model`] so answer-target policy is applied before the call. /// The method an algorithm implements; [`run_stream`](Self::run_stream) drives it. async fn route(self: Arc, driver: Driver, request: Request) -> Result; @@ -387,7 +509,11 @@ pub trait Algorithm: Send + Sync + 'static { }) }) }) - .await; + .await + .map(|mut outcome| { + outcome.prepare_selected_request(); + outcome + }); let _ = driver.finish(result).await; } @@ -403,6 +529,44 @@ pub trait Algorithm: Send + Sync + 'static { } } +struct TargetPromptAlgorithm { + inner: Arc, + prompts: Arc, +} + +#[async_trait] +impl Algorithm for TargetPromptAlgorithm { + fn name(&self) -> &str { + self.inner.name() + } + + async fn route(self: Arc, driver: Driver, request: Request) -> Result { + let prompts = Arc::clone(&self.prompts); + let outcome = Arc::clone(&self.inner) + .route(driver.with_target_prompts(Arc::clone(&prompts)), request) + .await?; + Ok(outcome.with_target_prompts(prompts)) + } +} + +/// Decorates `inner` with answer-target prompt policy. +/// +/// The policy applies to answer calls made during routing and to terminal selected and fallback +/// requests. An algorithm that produces an answer while routing must make that call with +/// [`Driver::call_answer_model`]. +pub fn with_target_prompts( + inner: Arc, + prompts: TargetPrompts, +) -> Arc { + if prompts.is_empty() { + return inner; + } + Arc::new(TargetPromptAlgorithm { + inner, + prompts: Arc::new(prompts), + }) +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -410,8 +574,9 @@ mod tests { use super::*; use crate::core::testing::{Serve, ServeResult, echo, reply, test_drive}; use futures::StreamExt; + use serde_json::json; use switchyard_protocol::{ - LlmResponse, LlmResponseChunk, completion_text, text_request, text_response, + ContentBlock, LlmResponse, LlmResponseChunk, completion_text, text_request, text_response, }; #[derive(Debug, thiserror::Error)] @@ -445,7 +610,7 @@ mod tests { .ok_or(LibsyError::NoTargets)? .clone(); let response = driver - .call_model(request.clone(), vec![target.clone()]) + .call_answer_model(request.clone(), target.clone()) .await?; Ok(RoutingOutcome::answered(target, request, response)) } @@ -509,6 +674,65 @@ mod tests { names.iter().map(|name| ModelId::from(*name)).collect() } + 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() + } + + #[tokio::test] + async fn target_prompts_follow_selected_and_fallback_targets() -> Result<()> { + let algorithm: Arc = Arc::new(crate::Random::new( + target_set(&["weak", "strong", "plain", "bare"]), + Some(vec![1.0, 0.0, 0.0, 0.0]), + Some(1), + )?); + let algorithm = with_target_prompts( + algorithm, + TargetPrompts::default() + .with("weak", "legacy weak prompt") + .with("plain", "plain prompt"), + ); + let algorithm = with_target_prompts( + algorithm, + TargetPrompts::default() + .with("weak", "weak prompt") + .with("strong", "strong prompt"), + ); + + let mut request = request(); + request.llm_request.preservation.requests.insert( + "openai_chat".into(), + json!({"model": "auto", "messages": [{"role": "user", "content": "hi"}]}), + ); + let outcome = drive(algorithm, request, |_call| async { + Err(test_error("route-only algorithm emitted a model call")) + }) + .await?; + + assert_eq!(instruction_text(&outcome.request), ["weak prompt"]); + let strong = outcome.request_for(&ModelId::from("strong"))?; + assert_eq!(instruction_text(&strong), ["strong prompt"]); + assert!(strong.llm_request.preservation.requests.is_empty()); + assert_eq!( + instruction_text(&outcome.request_for(&ModelId::from("plain"))?), + ["plain prompt"] + ); + assert!(instruction_text(&outcome.request_for(&ModelId::from("bare"))?).is_empty()); + assert!(matches!( + outcome.request_for(&ModelId::from("missing")), + Err(LibsyError::TargetNotFound { target }) if target == "missing" + )); + Ok(()) + } + #[tokio::test] async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> { tokio::time::timeout(std::time::Duration::from_secs(1), async { diff --git a/crates/libsy/src/core/target_prompts.rs b/crates/libsy/src/core/target_prompts.rs new file mode 100644 index 000000000..bcdeb5010 --- /dev/null +++ b/crates/libsy/src/core/target_prompts.rs @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Target-specific system-prompt policy shared by routing algorithms and hosts. + +use std::collections::BTreeMap; + +use switchyard_protocol::ModelId; + +/// System prompts keyed by routing target. A target left unset is routed untouched. +#[derive(Clone, Debug, Default)] +pub struct TargetPrompts { + by_target: BTreeMap, +} + +impl TargetPrompts { + /// Hand `target` this prompt on every turn it serves. + pub fn with(mut self, target: impl Into, prompt: impl Into) -> Self { + self.by_target.insert(target.into(), prompt.into()); + self + } + + /// The prompt configured for `target`, if any. + pub fn get(&self, target: &ModelId) -> Option<&str> { + self.by_target.get(target).map(String::as_str) + } + + /// Whether any target has a prompt, so a caller can skip empty policy layers. + pub fn is_empty(&self) -> bool { + self.by_target.is_empty() + } +} diff --git a/crates/libsy/src/core/testing.rs b/crates/libsy/src/core/testing.rs index 3fdb8c38b..b26627fcc 100644 --- a/crates/libsy/src/core/testing.rs +++ b/crates/libsy/src/core/testing.rs @@ -49,17 +49,20 @@ pub(crate) async fn test_drive( ) -> Result<(ModelId, Response)> { let serve = Arc::new(serve); let routing_serve = Arc::clone(&serve); - let outcome = crate::drive(algorithm, request, move |call| { + let mut outcome = crate::drive(algorithm, request, move |call| { fulfill(Arc::clone(&routing_serve), call) }) .await?; let selected_model = outcome.selected_model_id.clone(); - let response = match outcome.response { + let response = match outcome.response.take() { Some(response) => response, - None => serve - .serve(selected_model.clone(), outcome.request) - .await - .map_err(|source| LibsyError::client_call(selected_model.clone(), source))?, + None => { + let request = outcome.request_for(&selected_model)?; + serve + .serve(selected_model.clone(), request) + .await + .map_err(|source| LibsyError::client_call(selected_model.clone(), source))? + } }; Ok((selected_model, response)) } diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 1bb46db73..adb1068b5 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -5,7 +5,10 @@ #![doc = include_str!("../README.md")] mod core; -pub use core::algorithm::{Algorithm, CallModel, Driver, RoutingOutcome, Step, StepStream, drive}; +pub use core::TargetPrompts; +pub use core::algorithm::{ + Algorithm, CallModel, Driver, RoutingOutcome, Step, StepStream, drive, with_target_prompts, +}; pub use core::classifier::{Classification, Classifier, Score}; pub use core::processor::{Event, Processor}; pub use core::state::{State, StateValue}; @@ -28,7 +31,7 @@ pub use algorithms::util::classifier_contract::{ ClassifierContractConfig, ClassifierResponseFormat, }; pub use algorithms::util::escalation::EscalationJudgeConfig; -pub use algorithms::util::prompts::{SystemPromptProcessor, TargetPrompts, append_note}; +pub use algorithms::util::prompts::{SystemPromptProcessor, append_note}; pub use algorithms::util::subagent::SubagentOverride; pub use algorithms::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignals}; diff --git a/crates/switchyard-py/Cargo.toml b/crates/switchyard-py/Cargo.toml index 6fe4fdbb4..4b3278bbb 100644 --- a/crates/switchyard-py/Cargo.toml +++ b/crates/switchyard-py/Cargo.toml @@ -18,6 +18,7 @@ crate-type = ["cdylib", "rlib"] [dependencies] futures.workspace = true http.workspace = true +parking_lot.workspace = true switchyard-libsy.workspace = true pyo3 = { version = "0.28.3", features = ["abi3-py310", "extension-module"] } pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime"] } diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index 4546c549d..470f7be94 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -8,15 +8,17 @@ use std::sync::Arc; use futures::StreamExt; use http::header::{HeaderName, HeaderValue}; +use parking_lot::Mutex as SyncMutex; use pyo3::exceptions::{PyBaseException, PyStopAsyncIteration, PyTypeError, PyValueError}; use pyo3::prelude::*; +use pyo3::types::{PyMapping, PyMappingMethods}; use serde_json::Value; use switchyard_libsy::{ Algorithm, CallModel, ClassifierContractConfig, ClassifierResponseFormat, CustomClassifierConfig, CustomClassifierPolicy, EscalationJudgeConfig, HandoffNoteConfig, LibsyError as RustLibsyError, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, PickerMode, Random, RoutingOutcome, StageRouter, StageRouterConfig, Step as RustStep, - StepStream, TaskClassifierConfig, + StepStream, TargetPrompts, TaskClassifierConfig, with_target_prompts, }; use switchyard_protocol::{ AggLlmResponse, LlmClientError, LlmResponse, Metadata, ModelId, Request, Response, @@ -382,6 +384,18 @@ impl PyModelCall { self.models.clone() } + /// Prepare the normalized request for one routing-time candidate. + fn request_for(&self, py: Python<'_>, model: String) -> PyResult> { + let call = self + .inner + .as_ref() + .ok_or_else(|| py_libsy_error("model call has already been completed"))?; + let request = call + .request_for(&ModelId::new(model)) + .map_err(py_libsy_error)?; + to_python(py, &request.llm_request) + } + /// Fulfill this call with an aggregate normalized response dictionary. fn respond(&mut self, response: &Bound<'_, PyAny>) -> PyResult<()> { let aggregate = from_python::(response)?; @@ -424,9 +438,7 @@ impl PyModelCall { /// The terminal routing selection, rewritten request, and optional existing response. #[pyclass(name = "RoutingOutcome", module = "switchyard.libsy", frozen)] struct PyRoutingOutcome { - selected_model_id: String, - fallback_models: Vec, - request: Py, + inner: SyncMutex, response: Option>, } @@ -434,20 +446,36 @@ struct PyRoutingOutcome { impl PyRoutingOutcome { /// The model selected by the algorithm and tried first by the host. #[getter] - fn selected_model_id(&self) -> &str { - &self.selected_model_id + fn selected_model_id(&self) -> String { + self.inner.lock().selected_model_id.to_string() } /// Additional models the host may try in order after an eligible failure. #[getter] fn fallback_models(&self) -> Vec { - self.fallback_models.clone() + self.inner + .lock() + .fallback_models + .iter() + .map(ToString::to_string) + .collect() } /// The normalized request after routing-time rewrites. #[getter] - fn request(&self, py: Python<'_>) -> Py { - self.request.clone_ref(py) + fn request(&self, py: Python<'_>) -> PyResult> { + let request = self.inner.lock().request.llm_request.clone(); + to_python(py, &request) + } + + /// Prepare the normalized answer request for the selected model or a fallback. + fn request_for(&self, py: Python<'_>, model: String) -> PyResult> { + let request = { + let outcome = self.inner.lock(); + outcome.request_for(&ModelId::new(model)) + } + .map_err(py_libsy_error)?; + to_python(py, &request.llm_request) } /// An answer produced while routing, when one already exists. @@ -501,6 +529,20 @@ struct PyAlgorithm { #[pymethods] impl PyAlgorithm { + /// Return an algorithm that applies system prompts by answer target. + fn with_target_prompts(&self, prompts: &Bound<'_, PyMapping>) -> PyResult { + let prompts = prompts + .items()? + .extract::>()? + .into_iter() + .fold(TargetPrompts::default(), |prompts, (target, prompt)| { + prompts.with(target, prompt) + }); + Ok(Self { + inner: with_target_prompts(Arc::clone(&self.inner), prompts), + }) + } + /// Run the algorithm as routing-time model calls followed by one terminal outcome. /// /// `headers`, when given, is normalized into the request's correlation @@ -541,13 +583,8 @@ async fn step_to_python(step: RustStep) -> PyResult { }) }), RustStep::Done(outcome) => { - let RoutingOutcome { - selected_model_id, - fallback_models, - request, - response, - } = *outcome; - let response = match response { + let mut outcome = *outcome; + let response = match outcome.response.take() { Some(response) => Some( response .llm_response @@ -562,12 +599,7 @@ async fn step_to_python(step: RustStep) -> PyResult { outcome: Py::new( py, PyRoutingOutcome { - selected_model_id: selected_model_id.to_string(), - fallback_models: fallback_models - .iter() - .map(ToString::to_string) - .collect(), - request: to_python(py, &request.llm_request)?, + inner: SyncMutex::new(outcome), response: response .as_ref() .map(|response| to_python(py, response)) diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index c7037717b..3ffa031a1 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -89,6 +89,8 @@ def request(self) -> dict[str, object]: ... @property def models(self) -> list[str]: ... + def request_for(self, model: str) -> dict[str, object]: ... + def respond(self, response: Mapping[str, object]) -> None: ... def fail(self, error: BaseException) -> None: ... @@ -104,6 +106,8 @@ def fallback_models(self) -> list[str]: ... @property def request(self) -> dict[str, object]: ... + def request_for(self, model: str) -> dict[str, object]: ... + @property def response(self) -> dict[str, object] | None: ... @@ -190,6 +194,8 @@ def __init__( @final class Algorithm: + def with_target_prompts(self, prompts: Mapping[str, str]) -> Algorithm: ... + def run_stream( self, request: Mapping[str, object], diff --git a/tests/test_libsy_minimal_bindings.py b/tests/test_libsy_minimal_bindings.py index ae2a9c7b4..836234c52 100644 --- a/tests/test_libsy_minimal_bindings.py +++ b/tests/test_libsy_minimal_bindings.py @@ -61,7 +61,7 @@ async def run_algorithm( match step: case Step.CallModel(call): for index, target in enumerate(call.models): - candidate_request = {**call.request, "model": target} + candidate_request = call.request_for(target) client = (clients or {})[target] try: response = await client.call(candidate_request) @@ -79,7 +79,7 @@ async def run_algorithm( return outcome.selected_model_id, outcome.response candidates = [outcome.selected_model_id, *outcome.fallback_models] for index, target in enumerate(candidates): - candidate_request = {**outcome.request, "model": target} + candidate_request = outcome.request_for(target) client = (clients or {})[target] try: response = await client.call(candidate_request) @@ -93,7 +93,7 @@ async def run_algorithm( async def test_random_streams_complex_steps_and_accepts_a_dictionary_response() -> None: client = EchoClient("fast") - algorithm = algorithms.random(["fast"]) + algorithm = algorithms.random(["fast"]).with_target_prompts({"fast": "fast prompt"}) outcome: RoutingOutcome | None = None variants: list[str] = [] @@ -110,6 +110,7 @@ async def test_random_streams_complex_steps_and_accepts_a_dictionary_response() assert outcome.response is None response = await client.call(outcome.request) assert client.calls[0]["model"] == "fast" + assert client.calls[0]["instructions"][0]["content"][0]["text"] == "fast prompt" assert client.calls[0]["messages"][0]["content"] == [ {"type": "text", "text": "hello"} ] @@ -361,8 +362,9 @@ def test_invalid_request_is_rejected_at_the_boundary() -> None: async def test_context_window_failure_falls_back_to_the_next_model() -> None: - class OverflowClient: + class OverflowClient(EchoClient): async def call(self, request: dict[str, Any]) -> dict[str, Any]: + self.calls.append(request) raise ContextWindowExceededError("request exceeds context window") algorithm = algorithms.stage_router( @@ -370,11 +372,18 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: "fast", picker="efficient_first", confidence_threshold=0.5, + capable_system_prompt="strong prompt", + efficient_system_prompt="fast prompt", ) + fast = OverflowClient("fast") + strong = EchoClient("strong") selected_model, response = await run_algorithm( algorithm, - {"fast": OverflowClient(), "strong": EchoClient("strong")}, + {"fast": fast, "strong": strong}, ) assert selected_model == "fast" + assert fast.calls[0]["instructions"][0]["content"][0]["text"] == "fast prompt" + assert strong.calls[0]["instructions"][0]["content"][0]["text"] == "strong prompt" + assert len(strong.calls[0]["instructions"]) == 1 assert response["model"] == "strong" From 6a06b79c5ec53bd92dde228a494d9455a739f6b7 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 17 Aug 2026 12:01:26 -0700 Subject: [PATCH 3/3] feat(server): configure system prompts by target Signed-off-by: Alex Fournier --- crates/switchyard-server/CONFIGURATION.md | 4 + crates/switchyard-server/src/config.rs | 112 ++++++++++++++---- crates/switchyard-server/src/lib.rs | 8 +- crates/switchyard-server/tests/server.rs | 103 ++++++++++++---- docs/reference/toml_schema.md | 1 + .../stage_router_routing.md | 14 ++- 6 files changed, 190 insertions(+), 52 deletions(-) diff --git a/crates/switchyard-server/CONFIGURATION.md b/crates/switchyard-server/CONFIGURATION.md index efd438109..f0113337e 100644 --- a/crates/switchyard-server/CONFIGURATION.md +++ b/crates/switchyard-server/CONFIGURATION.md @@ -15,10 +15,14 @@ max_retries = 2 id = "provider/model" llm_client = "provider" extra_body = { chat_template_kwargs = { enable_thinking = false } } +system_prompt = "instructions for this model" ``` `extra_body` is target-specific. It shallow-merges top-level provider options into the outbound request, while explicit request fields win on conflicts. +`system_prompt` is also target-specific. It is prepended only when that target +serves an answer call, including a fallback after another target exceeds its +context window; classifier and judge calls are unchanged. The `chat_template_kwargs.enable_thinking` example is a provider/model-specific vLLM option. It is not a portable Switchyard reasoning switch. Use it on a judge diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index a78695cd5..78f804515 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -13,6 +13,7 @@ use libsy::{ CustomClassifierConfig, CustomClassifierPolicy, EscalationJudgeConfig, GateTrigger, HandoffNoteConfig, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, Passthrough, PickerMode, Random, StageRouter, StageRouterConfig, TargetPrompts, TaskClassifierConfig, + with_target_prompts, }; use serde::Deserialize; use serde_json::Value; @@ -101,8 +102,11 @@ impl ServerConfig { ))); } let algorithm = build_algorithm(route_name, config, &targets)?; + let target_prompts = self.build_route_target_prompts(route_name, config)?; let (client, caller_auth) = self.build_route_clients(route_name, config, &clients)?; - let count_tokens_target = self.build_count_tokens_target(config, &clients); + let count_tokens_target = + self.build_count_tokens_target(config, &clients, &target_prompts); + let algorithm = with_target_prompts(algorithm, target_prompts); routes.push(( config.id().clone(), algorithm, @@ -207,10 +211,40 @@ impl ServerConfig { Ok((ClientRouter::new(by_model), caller_auth)) } + /// Builds answer-target prompt policy and rejects aliases that lose prompt identity. + fn build_route_target_prompts( + &self, + route_name: &str, + route: &RouteConfig, + ) -> ServerResult { + let mut prompts = TargetPrompts::default(); + let mut by_model = HashMap::<&ModelId, (&str, Option<&str>)>::new(); + for (name, legacy_prompt) in route.routing_targets_with_legacy_prompts() { + let target = self.targets.get(name).ok_or_else(|| { + ServerError::new(format!("route references unknown target {name}")) + })?; + let effective = target.system_prompt.as_deref().or(legacy_prompt); + if let Some((previous, previous_prompt)) = + by_model.insert(&target.id, (name, effective)) + && previous_prompt != effective + { + return Err(ServerError::new(format!( + "route {route_name} maps answer targets {previous} and {name} to model {} with different system prompts", + target.id + ))); + } + if let Some(prompt) = effective { + prompts = prompts.with(target.id.clone(), prompt); + } + } + Ok(prompts) + } + fn build_count_tokens_target( &self, route_config: &RouteConfig, clients: &BTreeMap>, + target_prompts: &TargetPrompts, ) -> Option { route_config .routing_target_names() @@ -230,6 +264,7 @@ impl ServerConfig { .map(|(_, _, target, client)| CountTokensTarget { model: target.id.clone(), client: client.clone(), + system_prompt: target_prompts.get(&target.id).map(str::to_owned), }) } } @@ -265,6 +300,7 @@ struct TargetConfig { llm_client: String, #[serde(default)] extra_body: BTreeMap, + system_prompt: Option, } #[derive(Clone, Copy, Debug, Deserialize)] @@ -599,6 +635,27 @@ impl RouteConfig { } } + /// Completion targets paired with the legacy Stage prompt for that role, when present. + fn routing_targets_with_legacy_prompts(&self) -> Vec<(&str, Option<&str>)> { + match self { + Self::StageRouter { + capable_target, + efficient_target, + capable_system_prompt, + efficient_system_prompt, + .. + } => vec![ + (capable_target, capable_system_prompt.as_deref()), + (efficient_target, efficient_system_prompt.as_deref()), + ], + _ => self + .routing_target_names() + .into_iter() + .map(|name| (name, None)) + .collect(), + } + } + /// Every target the algorithm may call, including judge-only targets. /// /// [`routing_target_names`](Self::routing_target_names) covers completion destinations; @@ -1025,8 +1082,6 @@ fn build_algorithm( confidence_threshold, recent_turn_window, handoff_notes, - capable_system_prompt, - efficient_system_prompt, classifier, .. } => { @@ -1040,12 +1095,6 @@ fn build_algorithm( let mut config = StageRouterConfig::new(*picker, *confidence_threshold); config.recent_window = *recent_turn_window; config.handoff_notes = handoff_notes.clone(); - config.tier_prompts = tier_prompts( - &capable, - capable_system_prompt.as_deref(), - &efficient, - efficient_system_prompt.as_deref(), - ); // The judge is called through its own target, so it is not a routing // destination and stays out of the tier pair. config.llm_fallback = classifier @@ -1144,23 +1193,6 @@ fn default_classifier_max_output_tokens() -> u64 { TaskClassifierConfig::default().max_output_tokens } -/// Keys each configured system prompt by the target it belongs to. -fn tier_prompts( - capable: &str, - capable_prompt: Option<&str>, - efficient: &str, - efficient_prompt: Option<&str>, -) -> TargetPrompts { - let mut prompts = TargetPrompts::default(); - if let Some(prompt) = capable_prompt { - prompts = prompts.with(capable, prompt); - } - if let Some(prompt) = efficient_prompt { - prompts = prompts.with(efficient, prompt); - } - prompts -} - fn resolve_targets<'a>( route_name: &str, names: impl IntoIterator, @@ -1500,6 +1532,34 @@ classifier_magic = true } } + #[test] + fn rejects_conflicting_prompts_for_answer_aliases_of_one_model() { + let target_conflict = VALID_CONFIG + .replace("id = \"weak/model\"", "id = \"strong/model\"") + .replace( + "id = \"strong/model\"\nllm_client = \"responses\"", + "id = \"strong/model\"\nllm_client = \"responses\"\nsystem_prompt = \"strong prompt\"", + ); + let stage_conflict = VALID_CONFIG.replace( + "[routes.noop]", + r#"[routes.stage] +id = "switchyard/stage" +type = "stage_router" +capable_target = "strong" +efficient_target = "strong" +picker = "efficient_first" +confidence_threshold = 0.5 +capable_system_prompt = "capable prompt" +efficient_system_prompt = "efficient prompt" + +[routes.noop]"#, + ); + for config in [target_conflict, stage_conflict] { + let error = error_message(&config); + assert!(error.contains("with different system prompts"), "{error}"); + } + } + #[test] fn accepts_duplicate_target_model_ids_on_one_client() -> ServerResult<()> { // Two targets share one model id on one client. The client keeps one and drops the diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 24f1026cc..b6b86d40d 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -146,10 +146,16 @@ impl CallerAuthKind { struct CountTokensTarget { model: ModelId, client: Arc, + system_prompt: Option, } impl CountTokensTarget { - async fn count_tokens(&self, request: Request) -> Result { + async fn count_tokens(&self, mut request: Request) -> Result { + switchyard_translation::prepare_request_for_target( + &mut request.llm_request, + &self.model, + self.system_prompt.as_deref(), + ); self.client.count_tokens(&self.model, request).await } } diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 3992837bf..7470882a2 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -96,14 +96,25 @@ async fn upstream_chat( Json(body): Json, ) -> HttpResponse { calls.lock().await.push(body.clone()); - if body["messages"][0]["content"] == "fail" { + let user_prompt = body["messages"] + .as_array() + .and_then(|messages| { + messages + .iter() + .find_map(|message| match message["role"].as_str() { + Some("user") => message["content"].as_str(), + _ => None, + }) + }) + .unwrap_or(""); + if user_prompt == "fail" { return ( StatusCode::IM_A_TEAPOT, Json(json!({"error": {"message": "upstream rejected request"}})), ) .into_response(); } - if body["messages"][0]["content"] == "auth-fail" { + if user_prompt == "auth-fail" { return ( StatusCode::UNAUTHORIZED, Json(json!({"error": {"message": "upstream authentication failed"}})), @@ -112,15 +123,14 @@ async fn upstream_chat( } let model = body["model"].as_str().unwrap_or("unknown").to_string(); - let prompt = body["messages"][0]["content"].as_str().unwrap_or(""); - if (model == "model/weak" && prompt == "unavailable") || prompt == "all-unavailable" { + if (model == "model/weak" && user_prompt == "unavailable") || user_prompt == "all-unavailable" { return ( StatusCode::SERVICE_UNAVAILABLE, Json(json!({"error": {"message": "upstream is unavailable"}})), ) .into_response(); } - if model == "model/weak" && body["messages"][0]["content"] == "overflow" { + if model == "model/weak" && user_prompt == "overflow" { return ( StatusCode::BAD_REQUEST, Json(json!({ @@ -133,7 +143,7 @@ async fn upstream_chat( .into_response(); } if body["stream"].as_bool() == Some(true) { - if body["messages"][0]["content"] == "stream-error" { + if user_prompt == "stream-error" { let events = [ json!({"id": "chatcmpl-stream-error", "model": model, "choices": [{"index": 0, "delta": {"role": "assistant"}}]}).to_string(), json!({"id": "chatcmpl-stream-error", "model": model, "choices": [{"index": 0, "delta": {"content": "before"}}]}).to_string(), @@ -727,6 +737,7 @@ llm_client = "mock" [targets.second] id = "{second}" llm_client = "mock" +system_prompt = "strong target prompt" [routes.random] id = "{ROUTE_MODEL}" @@ -915,6 +926,7 @@ base_url = "{base_url}" [targets.strong] id = "model/stats-strong" llm_client = "upstream" +system_prompt = "strong target prompt" [targets.weak] id = "model/stats-weak" @@ -927,6 +939,7 @@ capable_target = "strong" efficient_target = "weak" picker = "efficient_first" confidence_threshold = 0.5 +capable_system_prompt = "strong stage prompt" "#, base_url = upstream.base_url ))?; @@ -960,6 +973,13 @@ confidence_threshold = 0.5 Some("model/stats-strong"), "a critical error should escalate on the signals alone" ); + let calls = upstream.calls.lock().await; + assert_eq!(calls.len(), 1); + assert_eq!( + calls[0]["messages"][0], + json!({"role": "system", "content": "strong target prompt"}) + ); + drop(calls); let stats = send(&app, "GET", "/v1/stats", None).await?.json()?; assert_eq!( stats["algorithm_stats"]["stage_router"]["routing_decisions"]["override"]["targets"]["model/stats-strong"], @@ -1091,6 +1111,7 @@ base_url = "{base_url}" [targets.classifier] id = "model/classifier" llm_client = "upstream" +system_prompt = "must not reach judge calls" [targets.strong] id = "model/strong" @@ -1491,6 +1512,11 @@ base_url = "{base_url}" [targets.strong] id = "real/opus" llm_client = "claude" +system_prompt = "count target instructions" + +[targets.legacy] +id = "real/opus-legacy" +llm_client = "claude" [targets.other] id = "real/sonnet" @@ -1500,28 +1526,43 @@ llm_client = "claude" id = "switchyard/random" type = "random" targets = ["other", "strong"] + +[routes.stage] +id = "switchyard/stage" +type = "stage_router" +capable_target = "legacy" +efficient_target = "other" +picker = "efficient_first" +confidence_threshold = 0.5 +capable_system_prompt = "count legacy instructions" "#, base_url = upstream.base_url ))?; let app = build_switchyard_router(state); - let response = send( - &app, - "POST", - "/v1/messages/count_tokens", - Some(json!({ - "model": "switchyard/random", - "messages": [{"role": "user", "content": "hi"}] - })), - ) - .await?; - assert_eq!(response.status, StatusCode::OK); - assert_eq!(response.json()?["input_tokens"], 7); + for route in ["switchyard/random", "switchyard/stage"] { + let response = send( + &app, + "POST", + "/v1/messages/count_tokens", + Some(json!({ + "model": route, + "messages": [{"role": "user", "content": "hi"}] + })), + ) + .await?; + assert_eq!(response.status, StatusCode::OK); + assert_eq!(response.json()?["input_tokens"], 7); + } let calls = upstream.calls.lock().await; - assert_eq!(calls.len(), 1); - // The inbound route name is rewritten to the real upstream model. + assert_eq!(calls.len(), 2); + // The target-level prompt follows the preferred Opus target in the random route. assert_eq!(calls[0]["model"], "real/opus"); + assert_eq!(calls[0]["system"], "count target instructions"); + // Stage's legacy capable prompt uses the same effective policy on this bypass path. + assert_eq!(calls[1]["model"], "real/opus-legacy"); + assert_eq!(calls[1]["system"], "count legacy instructions"); Ok(()) } @@ -2146,13 +2187,18 @@ async fn unavailable_target_fails_over_across_endpoints_and_stops_when_exhausted ); assert_eq!(response.json()?["model"], "model/strong"); let calls = upstream.calls.lock().await; + let fallback_calls = &calls[previous_call_count..]; assert_eq!( - calls[previous_call_count..] + fallback_calls .iter() - .map(|call| call["model"].as_str().unwrap_or("")) + .filter_map(|call| call["model"].as_str()) .collect::>(), ["model/weak", "model/strong"] ); + assert_eq!( + fallback_calls[1]["messages"][0]["content"], + "strong target prompt" + ); } let stats = send(&app, "GET", "/v1/stats", None).await?.json()?; @@ -2620,10 +2666,12 @@ base_url = "{base_url}" [targets.executor] id = "model/executor" llm_client = "upstream" +system_prompt = "executor target prompt" [targets.advisor] id = "model/advisor" llm_client = "upstream" +system_prompt = "must not reach advisor calls" [routes.gated] id = "switchyard/advisor" @@ -2664,6 +2712,17 @@ async fn advisor_route_approve_flow_and_stats() -> TestResult { ); // Executor turn first, then the review consult. assert_eq!(upstream.models().await, ["model/executor", "model/advisor"]); + let calls = upstream.calls.lock().await; + assert_eq!( + calls[0]["messages"][0], + json!({"role": "system", "content": "executor target prompt"}) + ); + assert!( + !calls[1] + .to_string() + .contains("must not reach advisor calls") + ); + drop(calls); let stats = send(&app, "GET", "/v1/stats", None).await?.json()?; assert_eq!(stats["models"]["model/executor"]["calls"], 1); diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index 68d01bc54..b92ff7bcb 100644 --- a/docs/reference/toml_schema.md +++ b/docs/reference/toml_schema.md @@ -83,6 +83,7 @@ calls an upstream. | `id` | Yes | — | Exact model ID sent upstream. | | `llm_client` | Yes | — | Key under `[llm_clients]`. | | `extra_body` | No | `{}` | Values merged into the upstream request when the request does not already set that key. | +| `system_prompt` | No | unset | System prompt prepended when this target serves an answer call. Classifier and judge calls are unchanged. | ## `[routes.]` diff --git a/docs/routing_algorithms/stage_router_routing.md b/docs/routing_algorithms/stage_router_routing.md index f4e4c451e..c0b0ac9a8 100644 --- a/docs/routing_algorithms/stage_router_routing.md +++ b/docs/routing_algorithms/stage_router_routing.md @@ -225,12 +225,20 @@ escalation_note = "the previous model was stalling; pick up the diagnosis" ### Optional: per-tier system prompts ```toml -[routes.stage] +[targets.strong] +# ... +system_prompt = "diagnose before you edit" + +[targets.weak] # ... -capable_system_prompt = "diagnose before you edit" -efficient_system_prompt = "follow the settled plan" +system_prompt = "follow the settled plan" ``` +The existing `capable_system_prompt` and `efficient_system_prompt` route fields +remain supported for stage routes. Target-level prompts also work with the other +routing algorithms and follow the selected target when a call falls back. A +target-level prompt takes precedence over the matching legacy route field. + ### Optional: LLM classifier fallback By default the router uses tool signals only. To break ties on low-confidence