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/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] 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"