diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs index e10a17dd..1d2444d9 100644 --- a/crates/libsy-llm-client/src/lib.rs +++ b/crates/libsy-llm-client/src/lib.rs @@ -15,6 +15,8 @@ //! [`switchyard_libsy::Algorithm::run_stream`] and serves every model call the algorithm //! offloads, so a host that just wants the answer does not have to drive the step stream //! itself. +//! A host that drives the stream itself can use [`ClientRouter::resolve_call`] to resolve the +//! candidate client and apply target-specific prompts before making each call. pub mod backend; pub mod client; @@ -30,5 +32,5 @@ pub use client::{ModelConfig, TranslatingLlmClient}; pub use error::{LlmClientError, Result}; pub use observation::{LlmCallObservation, RunObservation, RunObserver}; pub use raw::RawResponse; -pub use run::{ClientRouter, run}; +pub use run::{ClientRouter, ResolvedLlmCall, run}; pub use switchyard_translation::RawEventStream; diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index 9ef13355..5242e490 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use parking_lot::Mutex; -use switchyard_libsy::{Algorithm, CallModel, LibsyError, Result, drive}; +use switchyard_libsy::{Algorithm, CallModel, LibsyError, Result, TargetPrompts, drive}; use switchyard_protocol::{ Decision, LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason, }; @@ -207,7 +207,6 @@ async fn call_one( count: usize, ) -> Result { let span = tracing::Span::current(); - observability::record_gen_ai_request(&span, &request.llm_request); if let Some(session_id) = call .request .metadata @@ -217,9 +216,13 @@ async fn call_one( span.record("gen_ai.conversation.id", session_id); } let is_answer_call = call.is_answer_call; - // Resolved before the clock starts: picking the client is Switchyard's work, not - // the provider's, so it belongs in the routing overhead. + // Client and prompt selection are Switchyard work, so they remain outside provider latency. let client = clients.route(model_id); + let mut request = request; + if client.is_ok() { + clients.apply_target_prompt(model_id, is_answer_call, &mut request); + } + observability::record_gen_ai_request(&span, &request.llm_request); let started = Instant::now(); let result = match client { Ok(client) => client.call(request).await, @@ -290,7 +293,12 @@ fn request_for(request: &Request, target: &ModelId) -> Request { /// Cloning is cheap — the mapping is shared, so one router can serve every request. #[derive(Clone)] pub struct ClientRouter { - routing: Arc, + inner: Arc, +} + +struct ClientRouterInner { + routing: Routing, + prompts: TargetPrompts, } enum Routing { @@ -300,11 +308,40 @@ enum Routing { ByModel(HashMap>), } +/// One normalized model call resolved to its target client and target-specific prompt. +pub struct ResolvedLlmCall { + client: Arc, + request: Request, +} + +impl ResolvedLlmCall { + /// The request that will be sent, including any selected target's system prompt. + pub fn request(&self) -> &Request { + &self.request + } + + /// Perform the resolved call through its selected client. + pub async fn call(self) -> std::result::Result { + self.client.call(self.request).await + } +} + impl ClientRouter { /// Build a router over `model name -> client`, for targets spread across providers. pub fn new(by_model: HashMap>) -> Self { + Self::new_with_target_prompts(by_model, TargetPrompts::default()) + } + + /// Build a router with system prompts applied to answer calls by selected target. + pub fn new_with_target_prompts( + by_model: HashMap>, + prompts: TargetPrompts, + ) -> Self { Self { - routing: Arc::new(Routing::ByModel(by_model)), + inner: Arc::new(ClientRouterInner { + routing: Routing::ByModel(by_model), + prompts, + }), } } @@ -315,7 +352,32 @@ impl ClientRouter { /// only duplicate that. pub fn single(client: Arc) -> Self { Self { - routing: Arc::new(Routing::Single(client)), + inner: Arc::new(ClientRouterInner { + routing: Routing::Single(client), + prompts: TargetPrompts::default(), + }), + } + } + + /// Resolve a candidate model and request into the exact target call a host should perform. + /// + /// This is the prompt-aware host boundary. It stamps the selected model and prepends + /// that target's configured prompt only for answer calls. + pub fn resolve_call( + &self, + target: &ModelId, + is_answer_call: bool, + mut request: Request, + ) -> std::result::Result { + let client = Arc::clone(self.route(target)?); + request.llm_request.model = Some(target.to_string()); + self.apply_target_prompt(target, is_answer_call, &mut request); + Ok(ResolvedLlmCall { client, request }) + } + + fn apply_target_prompt(&self, target: &ModelId, is_answer_call: bool, request: &mut Request) { + if is_answer_call && let Some(prompt) = self.inner.prompts.get(target) { + switchyard_translation::prepend_system_prompt(&mut request.llm_request, prompt); } } @@ -327,7 +389,7 @@ impl ClientRouter { &self, model: &ModelId, ) -> std::result::Result<&Arc, LlmClientError> { - match self.routing.as_ref() { + match &self.inner.routing { Routing::Single(client) => Ok(client), Routing::ByModel(by_model) => { by_model diff --git a/crates/libsy-llm-client/tests/target_prompts.rs b/crates/libsy-llm-client/tests/target_prompts.rs new file mode 100644 index 00000000..bfc627e8 --- /dev/null +++ b/crates/libsy-llm-client/tests/target_prompts.rs @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::{collections::HashMap, sync::Arc}; + +use async_trait::async_trait; +use parking_lot::Mutex; +use switchyard_libsy::{Random, TargetPrompts}; +use switchyard_llm_client::ClientRouter; +use switchyard_protocol::{ + ContentBlock, LlmClientError, LlmResponse, ModelId, Request, Response, RoutedLlmClient, + text_response, +}; + +const WEAK: &str = "weak/model"; +const STRONG: &str = "strong/model"; + +#[derive(Default)] +struct RecordingClient { + calls: Mutex>, + overflow: Option, +} + +#[async_trait] +impl RoutedLlmClient for RecordingClient { + async fn call(&self, request: Request) -> Result { + let model = request.model_id().unwrap_or_default(); + self.calls.lock().push(request); + if self.overflow.as_ref() == Some(&model) { + return Err(LlmClientError::ContextWindowExceeded { + model, + message: "too long".to_string(), + }); + } + Ok(Response { + llm_response: LlmResponse::Agg(text_response(Some(model.to_string()), "ok")), + metadata: None, + }) + } +} + +fn instruction_text(request: &Request) -> Vec<&str> { + request + .llm_request + .instructions + .iter() + .flat_map(|instruction| instruction.content.iter()) + .filter_map(|block| match block { + ContentBlock::Text { text } => Some(text.as_str()), + _ => None, + }) + .collect() +} + +#[tokio::test] +async fn fallback_call_receives_the_new_targets_prompt() -> switchyard_libsy::Result<()> { + let client = Arc::new(RecordingClient { + calls: Mutex::new(Vec::new()), + overflow: Some(ModelId::from(WEAK)), + }); + let routed_client: Arc = client.clone(); + let clients = HashMap::from([ + (ModelId::from(WEAK), Arc::clone(&routed_client)), + (ModelId::from(STRONG), routed_client), + ]); + let prompts = TargetPrompts::default() + .with(WEAK, "weak prompt") + .with(STRONG, "strong prompt"); + let algorithm = Random::new( + vec![ModelId::from(WEAK), ModelId::from(STRONG)], + Some(vec![1.0, 0.0]), + Some(1), + )?; + + switchyard_llm_client::run( + Arc::new(algorithm), + ClientRouter::new_with_target_prompts(clients, prompts), + Request::default(), + None, + ) + .await?; + + let calls = client.calls.lock(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].model_id().as_deref(), Some(WEAK)); + assert_eq!(instruction_text(&calls[0]), ["weak prompt"]); + assert_eq!(calls[1].model_id().as_deref(), Some(STRONG)); + assert_eq!(instruction_text(&calls[1]), ["strong prompt"]); + Ok(()) +} diff --git a/crates/switchyard-server/CONFIGURATION.md b/crates/switchyard-server/CONFIGURATION.md index efd43810..f0113337 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/README.md b/crates/switchyard-server/README.md index b853929d..f7b89a6e 100644 --- a/crates/switchyard-server/README.md +++ b/crates/switchyard-server/README.md @@ -83,6 +83,8 @@ client's `base_url` should receive the caller's login. A forwarding route must be called through the matching provider API. Target-level `extra_body` values are shallow-merged into the upstream request when the request does not already contain that key. +An optional target-level `system_prompt` is prepended when that target serves an +answer call; it is not added to classifier or judge calls. `max_retries` defaults to `2` and applies to transport failures, timeouts, HTTP 408/429, and 5xx responses. diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index 56139023..888fa672 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -79,6 +79,15 @@ impl ServerConfig { for (target_name, target) in &self.targets { validate_value("target name", target_name)?; validate_value(&format!("target {target_name} id"), &target.id)?; + if target + .system_prompt + .as_deref() + .is_some_and(|prompt| prompt.trim().is_empty()) + { + return Err(ServerError::new(format!( + "target {target_name} system_prompt must not be empty" + ))); + } if !seen_client_model_ids.insert((target.llm_client.as_str(), target.id.as_str())) { tracing::warn!( "target {target_name} reuses model id {} on llm client {}; only one target per id is kept and the other is dropped. Give each target a unique model id, or point both routes at one target.", @@ -100,9 +109,12 @@ impl ServerConfig { "route {route_name} context_window must be greater than zero" ))); } + let target_prompts = self.build_target_prompts(route_name, config)?; let algorithm = build_algorithm(route_name, config, &targets)?; - 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 (client, caller_auth) = + self.build_route_clients(route_name, config, &clients, target_prompts)?; routes.push(( config.id().clone(), algorithm, @@ -176,6 +188,7 @@ impl ServerConfig { route_name: &str, route: &RouteConfig, clients: &BTreeMap>, + target_prompts: TargetPrompts, ) -> ServerResult<(ClientRouter, Option)> { let mut by_model = HashMap::new(); let mut caller_auth = None; @@ -204,13 +217,66 @@ impl ServerConfig { let client: Arc = client.clone(); by_model.insert(target.id.clone(), client); } - Ok((ClientRouter::new(by_model), caller_auth)) + Ok(( + ClientRouter::new_with_target_prompts(by_model, target_prompts), + caller_auth, + )) + } + + /// Resolves target fields and legacy Stage aliases into one prompt per routed model. + /// Judge-only targets are excluded; aliases sharing a model must agree, including unset values. + fn build_target_prompts( + &self, + route_name: &str, + route: &RouteConfig, + ) -> ServerResult { + let mut by_model: BTreeMap> = BTreeMap::new(); + for name in route.routing_target_names() { + let target = self.targets.get(name).ok_or_else(|| { + ServerError::new(format!( + "route {route_name} references unknown target {name}" + )) + })?; + let legacy_prompt = route.legacy_system_prompt(name); + if let (Some(prompt), Some(legacy)) = (target.system_prompt.as_deref(), legacy_prompt) + && prompt != legacy + { + return Err(ServerError::new(format!( + "route {route_name} configures different system prompts for {}", + target.id + ))); + } + let prompt = legacy_prompt.or(target.system_prompt.as_deref()); + // Decisions carry model IDs, so aliases in one route must agree on the prompt. + match by_model.entry(target.id.clone()) { + std::collections::btree_map::Entry::Vacant(entry) => { + entry.insert(prompt.map(str::to_string)); + } + std::collections::btree_map::Entry::Occupied(entry) + if entry.get().as_deref() != prompt => + { + return Err(ServerError::new(format!( + "route {route_name} configures different system prompts for {}", + target.id + ))); + } + _ => {} + } + } + + Ok(by_model + .into_iter() + .filter_map(|(target, prompt)| prompt.map(|prompt| (target, prompt))) + .fold(TargetPrompts::default(), |prompts, (target, prompt)| { + prompts.with(target, prompt) + })) } fn build_count_tokens_target( &self, route_config: &RouteConfig, clients: &BTreeMap>, + target_prompts: &TargetPrompts, ) -> Option { route_config .routing_target_names() @@ -230,6 +296,7 @@ impl ServerConfig { .map(|(_, _, target, client)| CountTokensTarget { model: target.id.clone(), client: client.clone(), + system_prompt: target_prompts.get(&target.id).map(str::to_string), }) } } @@ -265,6 +332,7 @@ struct TargetConfig { llm_client: String, #[serde(default)] extra_body: BTreeMap, + system_prompt: Option, } #[derive(Clone, Copy, Debug, Deserialize)] @@ -546,6 +614,23 @@ impl RouteConfig { } } + // Stage's original route-level fields remain aliases for target prompts. + fn legacy_system_prompt(&self, target_name: &str) -> Option<&str> { + match self { + Self::StageRouter { + capable_target, + capable_system_prompt, + .. + } if target_name == capable_target => capable_system_prompt.as_deref(), + Self::StageRouter { + efficient_target, + efficient_system_prompt, + .. + } if target_name == efficient_target => efficient_system_prompt.as_deref(), + _ => None, + } + } + /// Every target the algorithm may call, including judge-only targets. /// /// [`routing_target_names`](Self::routing_target_names) covers completion destinations; @@ -965,8 +1050,6 @@ fn build_algorithm( confidence_threshold, recent_turn_window, handoff_notes, - capable_system_prompt, - efficient_system_prompt, classifier, .. } => { @@ -980,12 +1063,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 @@ -1017,23 +1094,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, @@ -1144,6 +1204,47 @@ target = "weak" Ok(()) } + #[test] + fn target_system_prompts_must_be_unambiguous() { + let set_and_unset = VALID_CONFIG + .replace( + "[targets.strong]\nid = \"strong/model\"\nllm_client = \"responses\"", + "[targets.strong]\nid = \"shared/model\"\nllm_client = \"responses\"\nsystem_prompt = \"strong prompt\"", + ) + .replace( + "[targets.weak]\nid = \"weak/model\"\nllm_client = \"anthropic\"", + "[targets.weak]\nid = \"shared/model\"\nllm_client = \"anthropic\"", + ); + let error = error_message(&set_and_unset); + assert!( + error.contains("route classifier configures different system prompts for shared/model"), + "{error}" + ); + + let legacy_stage = format!( + r#"{VALID_CONFIG} +[routes.stage] +id = "switchyard/stage" +type = "stage_router" +capable_target = "strong" +efficient_target = "weak" +picker = "efficient_first" +confidence_threshold = 0.5 +capable_system_prompt = "legacy prompt" +"# + ) + .replacen( + "[targets.strong]", + "[targets.strong]\nsystem_prompt = \"target prompt\"", + 1, + ); + let error = error_message(&legacy_stage); + assert!( + error.contains("route stage configures different system prompts for strong/model"), + "{error}" + ); + } + #[test] fn an_escalation_table_switches_the_classifier_route_to_escalation() -> ServerResult<()> { // Present: the classifier target judges the weak tier's reply each turn instead of @@ -1356,6 +1457,14 @@ classifier_magic = true VALID_CONFIG.replace("[targets.strong]", "[targets.\" strong \"]"), "target name must be non-empty and have no surrounding whitespace", ), + ( + VALID_CONFIG.replacen( + "[targets.strong]", + "[targets.strong]\nsystem_prompt = \" \"", + 1, + ), + "target strong system_prompt must not be empty", + ), ( VALID_CONFIG.replace( "targets = [\"strong\", \"weak\"]", diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index fd02cd01..1d5f408e 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -146,10 +146,14 @@ 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 { + if let Some(prompt) = self.system_prompt.as_deref() { + switchyard_translation::prepend_system_prompt(&mut request.llm_request, prompt); + } 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 1cd87dca..f542f92c 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -96,14 +96,24 @@ 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| { + (message["role"] == "user") + .then(|| message["content"].as_str()) + .flatten() + }) + }) + .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 +122,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 +142,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(), @@ -694,10 +703,12 @@ max_retries = 0 [targets.first] id = "{first}" llm_client = "mock" +system_prompt = "weak target prompt" [targets.second] id = "{second}" llm_client = "mock" +system_prompt = "strong target prompt" [routes.random] id = "{ROUTE_MODEL}" @@ -886,10 +897,12 @@ base_url = "{base_url}" [targets.strong] id = "model/stats-strong" llm_client = "upstream" +system_prompt = "strong stage prompt" [targets.weak] id = "model/stats-weak" llm_client = "upstream" +system_prompt = "weak stage prompt" [routes.stage] id = "switchyard/stage" @@ -931,6 +944,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 stage 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"], @@ -1062,6 +1082,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" @@ -1385,6 +1406,7 @@ base_url = "{base_url}" [targets.strong] id = "real/opus" llm_client = "claude" +system_prompt = "count these instructions" [targets.other] id = "real/sonnet" @@ -1416,6 +1438,7 @@ targets = ["other", "strong"] assert_eq!(calls.len(), 1); // The inbound route name is rewritten to the real upstream model. assert_eq!(calls[0]["model"], "real/opus"); + assert_eq!(calls[0]["system"], "count these instructions"); Ok(()) } @@ -2040,12 +2063,26 @@ 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("")) + .map(|call| ( + call["model"].as_str().unwrap_or(""), + call["messages"][0]["content"].as_str().unwrap_or(""), + )) .collect::>(), - ["model/weak", "model/strong"] + [ + ("model/weak", "weak target prompt"), + ("model/strong", "strong target prompt"), + ] + ); + assert!( + fallback_calls[1]["messages"] + .as_array() + .is_some_and(|messages| messages + .iter() + .all(|message| { message["content"] != "weak target prompt" })) ); } diff --git a/crates/switchyard-translation/src/lib.rs b/crates/switchyard-translation/src/lib.rs index d07a3497..0fda8a03 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, prepend_system_prompt, + sanitize_anthropic_tool_use_id, }; diff --git a/crates/switchyard-translation/src/util.rs b/crates/switchyard-translation/src/util.rs index 9fcf769f..484085c6 100644 --- a/crates/switchyard-translation/src/util.rs +++ b/crates/switchyard-translation/src/util.rs @@ -9,8 +9,8 @@ use serde_json::{Map, Value, json}; use crate::diagnostic::TranslationDiagnostic; use crate::error::{Result, TranslationError}; -use crate::format::FormatId; -use crate::llm::{ContentBlock, LlmRequest, Message, PreservationMetadata}; +use crate::format::{FormatId, WireFormat}; +use crate::llm::{ContentBlock, InstructionBlock, LlmRequest, Message, PreservationMetadata, Role}; use crate::policy::{ LossyConversionPolicy, PreservationPolicy, TranslationPolicy, UnknownFieldPolicy, }; @@ -271,6 +271,173 @@ pub fn exact_preserved_response( .flatten() } +/// Prepends a system prompt while preserving unrelated fields in built-in request formats. +/// +/// Exact preserved bodies are patched in place so unrelated provider fields survive. A custom +/// or malformed preserved body is discarded and will be rebuilt from the normalized request. +pub fn prepend_system_prompt(request: &mut LlmRequest, prompt: &str) { + let is_already_first = request.instructions.first().is_some_and(|instruction| { + instruction.role == Role::System + && matches!( + instruction.content.as_slice(), + [ContentBlock::Text { text }] if text_starts_with_prompt(text, prompt) + ) + }); + if !is_already_first { + request.instructions.insert( + 0, + InstructionBlock { + role: Role::System, + content: vec![ContentBlock::Text { + text: prompt.to_string(), + }], + }, + ); + } + let mut embedded_formats = Vec::new(); + request.preservation.requests.retain(|format, body| { + let patched = match format.as_str() { + format if format == WireFormat::OpenAiChat.as_str() => { + prepend_openai_chat_system_prompt(body, prompt) + } + format if format == WireFormat::OpenAiResponses.as_str() => { + prepend_text_system_prompt(body, "instructions", prompt) + } + format if format == WireFormat::AnthropicMessages.as_str() => { + prepend_anthropic_system_prompt(body, prompt) + } + _ => false, + }; + if patched && take_embedded_preservation(body) { + embedded_formats.push(format.clone()); + } + patched + }); + // Refresh envelopes only on bodies that arrived with one. The snapshot is serialized while + // every body is envelope-free, so it cannot recursively nest stale preservation metadata. + if !embedded_formats.is_empty() + && let Ok(envelope) = serde_json::to_value(&request.preservation) + { + for format in embedded_formats { + if let Some(body) = request.preservation.requests.get_mut(&format) { + restore_embedded_preservation(body, &envelope); + } + } + } +} + +// Inserts a system message ahead of an exact Chat Completions message list. +fn prepend_openai_chat_system_prompt(body: &mut Value, prompt: &str) -> bool { + let Some(body) = body.as_object_mut() else { + return false; + }; + let messages = body + .entry("messages".to_string()) + .or_insert_with(|| Value::Array(Vec::new())); + let Some(messages) = messages.as_array_mut() else { + return false; + }; + if messages.first().is_some_and(|message| { + message.get("role").and_then(Value::as_str) == Some("system") + && message + .get("content") + .is_some_and(|content| content_starts_with_prompt(content, prompt)) + }) { + return true; + } + messages.insert(0, json!({"role": "system", "content": prompt})); + true +} + +// Recognizes the scalar and block-array forms accepted for system content. +fn content_starts_with_prompt(content: &Value, prompt: &str) -> bool { + match content { + Value::String(text) => text_starts_with_prompt(text, prompt), + Value::Array(blocks) => blocks + .first() + .is_some_and(|block| text_block_matches_prompt(block, prompt)), + _ => false, + } +} + +fn text_block_matches_prompt(block: &Value, prompt: &str) -> bool { + block.get("type").and_then(Value::as_str) == Some("text") + && block + .get("text") + .and_then(Value::as_str) + .is_some_and(|text| text_starts_with_prompt(text, prompt)) +} + +fn text_starts_with_prompt(text: &str, prompt: &str) -> bool { + text == prompt + || text + .strip_prefix(prompt) + .is_some_and(|suffix| suffix.starts_with("\n\n")) +} + +// Prepends a prompt to a provider's scalar instruction field. +fn prepend_text_system_prompt(body: &mut Value, field: &str, prompt: &str) -> bool { + let Some(body) = body.as_object_mut() else { + return false; + }; + prepend_text_field(body, field, prompt) +} + +fn prepend_text_field(body: &mut Map, field: &str, prompt: &str) -> bool { + match body.get_mut(field) { + Some(Value::String(existing)) if text_starts_with_prompt(existing, prompt) => { + return true; + } + Some(Value::String(existing)) if existing.is_empty() => { + *existing = prompt.to_string(); + } + Some(Value::String(existing)) => { + *existing = format!("{prompt}\n\n{existing}"); + } + Some(value @ Value::Null) => { + *value = Value::String(prompt.to_string()); + } + None => { + body.insert(field.to_string(), Value::String(prompt.to_string())); + } + Some(_) => return false, + } + true +} + +// Preserves Anthropic's structured system blocks, including cache-control metadata. +fn prepend_anthropic_system_prompt(body: &mut Value, prompt: &str) -> bool { + let Some(body) = body.as_object_mut() else { + return false; + }; + if let Some(Value::Array(blocks)) = body.get_mut("system") { + if blocks + .first() + .is_some_and(|block| text_block_matches_prompt(block, prompt)) + { + return true; + } + blocks.insert(0, json!({"type": "text", "text": prompt})); + return true; + } + prepend_text_field(body, "system", prompt) +} + +// Removes a stale embedded snapshot and reports whether the caller must refresh it. +fn take_embedded_preservation(body: &mut Value) -> bool { + if let Some(metadata) = body.get_mut("metadata").and_then(Value::as_object_mut) { + return metadata.remove(SWITCHYARD_METADATA_KEY).is_some(); + } + false +} + +// Attaches the refreshed snapshot to a body that originally carried one. +fn restore_embedded_preservation(body: &mut Value, envelope: &Value) { + if let Some(metadata) = body.get_mut("metadata").and_then(Value::as_object_mut) { + metadata.insert(SWITCHYARD_METADATA_KEY.to_string(), envelope.clone()); + } +} + /// 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 e14d38e9..fd16f33e 100644 --- a/crates/switchyard-translation/tests/request_translation.rs +++ b/crates/switchyard-translation/tests/request_translation.rs @@ -6,11 +6,118 @@ use pretty_assertions::assert_eq; use serde_json::{Value, json}; use switchyard_translation::{ - LossyConversionPolicy, TranslationEngine, TranslationPolicy, WireFormat, + LossyConversionPolicy, TranslationEngine, TranslationPolicy, WireFormat, prepend_system_prompt, }; type TestResult = std::result::Result<(), Box>; +// Provider-only fields and embedded cross-format snapshots survive prompt insertion. +#[test] +fn prepending_a_system_prompt_preserves_exact_provider_fields() -> TestResult { + let engine = TranslationEngine::default(); + let policy = TranslationPolicy::default(); + let cases = [ + ( + WireFormat::OpenAiChat, + json!({ + "model": "gpt", + "messages": [ + {"role": "system", "name": "caller", "content": "client prompt"}, + {"role": "user", "content": "hi"} + ], + "prompt_cache_key": "session-1", + "stream": true, + "stream_options": {"include_usage": true} + }), + "/messages/0", + json!({"role": "system", "content": "target prompt"}), + "/stream_options/include_usage", + json!(true), + ), + ( + WireFormat::OpenAiResponses, + json!({ + "model": "gpt", + "instructions": "client prompt", + "input": "hi", + "include": ["reasoning.encrypted_content"], + "parallel_tool_calls": false, + "prompt_cache_key": "session-1", + "reasoning": {"effort": "high", "summary": "auto"}, + "metadata": { + "tenant": "example", + "_switchyard_translation": { + "requests": { + "anthropic_messages": { + "model": "claude", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}] + } + } + } + }, + "store": false, + "truncation": "auto" + }), + "/instructions", + json!("target prompt\n\nclient prompt"), + "/metadata/tenant", + json!("example"), + ), + ( + WireFormat::AnthropicMessages, + json!({ + "model": "claude", + "max_tokens": 64, + "system": [{ + "type": "text", + "text": "client prompt", + "cache_control": {"type": "ephemeral"} + }], + "messages": [{"role": "user", "content": "hi"}] + }), + "/system/0", + json!({"type": "text", "text": "target prompt"}), + "/system/1/cache_control/type", + json!("ephemeral"), + ), + ]; + + for (format, body, prompt_path, expected_prompt, preserved_path, expected_preserved) in cases { + let mut request = engine.decode_request(format, &body, &policy)?.request; + prepend_system_prompt(&mut request, "target prompt"); + let once = engine.encode_request(format, &request, &policy)?.body; + assert_eq!( + once.pointer(prompt_path), + Some(&expected_prompt), + "{format}" + ); + assert_eq!( + once.pointer(preserved_path), + Some(&expected_preserved), + "{format}" + ); + if format == WireFormat::OpenAiResponses { + assert_eq!( + once.pointer( + "/metadata/_switchyard_translation/requests/anthropic_messages/system" + ), + Some(&json!("target prompt")), + "embedded preservation must reflect the prompt mutation" + ); + let relayed = engine + .translate_request(format, WireFormat::AnthropicMessages, &once, &policy)? + .body; + assert_eq!(relayed["system"], "target prompt"); + } + + prepend_system_prompt(&mut request, "target prompt"); + let twice = engine.encode_request(format, &request, &policy)?.body; + assert_eq!(twice, once, "{format}: prompt insertion must be idempotent"); + } + Ok(()) +} + // Verifies Anthropic-only request fields are dropped or mapped for OpenAI Chat. #[test] fn anthropic_request_translates_to_openai_chat_without_anthropic_only_fields() -> TestResult { diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index 68d01bc5..b92ff7bc 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 f4e4c451..78af3fe6 100644 --- a/docs/routing_algorithms/stage_router_routing.md +++ b/docs/routing_algorithms/stage_router_routing.md @@ -225,12 +225,19 @@ 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. + ### Optional: LLM classifier fallback By default the router uses tool signals only. To break ties on low-confidence