From c529f29268942b04972769f36c63579156e85c79 Mon Sep 17 00:00:00 2001 From: Graham King Date: Wed, 9 Sep 2026 18:38:34 -0400 Subject: [PATCH] feat(libsy): Pass models at runtime not construction time- #630 Instead of giving the available models to the algorithm in `new` we pass them alongside the request in `run_stream` where they go in the `Driver`. See https://github.com/NVIDIA-NeMo/Switchyard/issues/588 Assisted-by: Codex:GPT 5.6 Sol high Signed-off-by: Graham King --- README.md | 8 +- ...2-telecom-custom-opus-qwen-aggressive.toml | 13 +- ...au2-telecom-custom-opus-qwen-balanced.toml | 17 +- crates/libsy-llm-client/src/run.rs | 71 ++- .../libsy-llm-client/tests/observability.rs | 139 ++++-- crates/libsy/src/algorithms/composite.rs | 115 +++-- crates/libsy/src/algorithms/escalation.rs | 93 ++-- crates/libsy/src/algorithms/fall_through.rs | 209 ++------ crates/libsy/src/algorithms/llm_class.rs | 457 +++++++++--------- crates/libsy/src/algorithms/passthrough.rs | 39 +- crates/libsy/src/algorithms/rand.rs | 216 ++++----- crates/libsy/src/algorithms/stage.rs | 271 +++++++---- crates/libsy/src/algorithms/subagent.rs | 81 ++-- .../src/algorithms/subagent_affinity_tests.rs | 10 +- crates/libsy/src/algorithms/util/affinity.rs | 16 +- .../libsy/src/algorithms/util/escalation.rs | 36 +- crates/libsy/src/algorithms/util/llm_judge.rs | 57 ++- crates/libsy/src/algorithms/util/prompts.rs | 229 +-------- crates/libsy/src/algorithms/util/stage.rs | 134 ++--- .../src/algorithms/util/target_selector.rs | 75 +-- crates/libsy/src/core/algorithm.rs | 58 ++- crates/libsy/src/core/testing.rs | 17 +- crates/libsy/src/lib.rs | 6 +- crates/prefill-router/tests/unit/algorithm.rs | 3 +- crates/protocol/src/category.rs | 48 ++ crates/protocol/src/lib.rs | 2 + .../src/runtime.rs | 3 +- crates/switchyard-py/src/libsy_bindings.rs | 117 ++--- crates/switchyard-runner/src/algorithm.rs | 392 +++++++++------ crates/switchyard-runner/src/config.rs | 58 ++- crates/switchyard-runner/src/lib.rs | 2 +- crates/switchyard-runner/src/route.rs | 32 +- crates/switchyard-runner/src/runner.rs | 4 + crates/switchyard-runner/tests/route.rs | 31 +- crates/switchyard-server/src/lib.rs | 9 +- crates/switchyard-server/tests/server.rs | 50 +- docs/reference/toml_schema.md | 15 +- .../llm_classifier_routing.md | 23 +- docs/routing_algorithms/subagent_routing.md | 11 +- examples/libsy.py | 3 +- .../plugins/stage_routing_plugin.py | 9 +- .../plugins/switchyard_routing_plugin.py | 23 +- .../unit/test_switchyard_routing_plugin.py | 37 +- switchyard_rust/libsy.py | 21 +- tests/test_libsy_minimal_bindings.py | 113 +++-- 45 files changed, 1745 insertions(+), 1628 deletions(-) create mode 100644 crates/protocol/src/category.rs diff --git a/README.md b/README.md index fa972f4cc..f003decc0 100644 --- a/README.md +++ b/README.md @@ -171,17 +171,15 @@ switchyard-protocol = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", b tokio = { version = "1", features = ["macros", "rt"] } ``` -**2. Construct an algorithm.** Target names are whatever your harness calls its -models. This is the stage router from the benchmark; `random`, -`llm_task_classifier`, and `llm_classifier` are built the same way. +**2. Construct an algorithm.** Models are supplied when each request runs. This +is the stage router from the benchmark; `random`, `llm_task_classifier`, and +`llm_classifier` are built the same way. ```python from switchyard.libsy import LlmResponse, Step from switchyard.libsy.algorithms import stage_router algorithm = stage_router( - "capable", - "efficient", picker="efficient_first", confidence_threshold=0.5, ) diff --git a/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-aggressive.toml b/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-aggressive.toml index c44f9737f..01d0577da 100644 --- a/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-aggressive.toml +++ b/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-aggressive.toml @@ -50,19 +50,18 @@ llm_client = "openrouter" [routes.switchyard] id = "switchyard" type = "llm_classifier" -classifier_target = "classifier" mode = "custom" recent_turn_window = 6 # Re-classify when the user speaks again, and hold that target across the tool calls # in between, so a tool chain never switches tier mid-task. classify_trigger = "user_turn" -targets = ["weak", "strong"] -default_target = "strong" +models = { judge = ["classifier"], capable = ["strong"], efficient = ["weak"], any = ["strong", "weak"] } +default_target = "capable" response_schema = ''' { "type": "object", "properties": { - "route": { "type": "string", "enum": ["weak", "strong"] }, + "route": { "type": "string", "enum": ["efficient", "capable"] }, "confidence": { "type": "number" }, "abstain": { "type": "boolean" } }, @@ -74,10 +73,10 @@ prompt = ''' You are a routing classifier inside a customer-service agent. Return exactly one JSON object: -{"route": "weak" or "strong", "confidence": number 0..1, "abstain": boolean} +{"route": "efficient" or "capable", "confidence": number 0..1, "abstain": boolean} -State the route DIRECTLY: "weak" = the on-device assistant handles this turn; -"strong" = escalate this turn to the frontier model. +State the route DIRECTLY: "efficient" = the on-device assistant handles this turn; +"capable" = escalate this turn to the frontier model. ROUTING BIAS: the WEAK tier is the DEFAULT — it handles nearly all support work end-to-end: lookups, standard actions, troubleshooting with known steps, diff --git a/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-balanced.toml b/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-balanced.toml index 9f8ed0a54..c7858df45 100644 --- a/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-balanced.toml +++ b/benchmark/routing-profiles/tau2-telecom-custom-opus-qwen-balanced.toml @@ -50,19 +50,18 @@ llm_client = "openrouter" [routes.switchyard] id = "switchyard" type = "llm_classifier" -classifier_target = "classifier" mode = "custom" recent_turn_window = 6 # Re-classify when the user speaks again, and hold that target across the tool calls # in between, so a tool chain never switches tier mid-task. classify_trigger = "user_turn" -targets = ["weak", "strong"] -default_target = "strong" +models = { judge = ["classifier"], capable = ["strong"], efficient = ["weak"], any = ["strong", "weak"] } +default_target = "capable" response_schema = ''' { "type": "object", "properties": { - "route": { "type": "string", "enum": ["weak", "strong"] }, + "route": { "type": "string", "enum": ["efficient", "capable"] }, "confidence": { "type": "number" }, "abstain": { "type": "boolean" } }, @@ -76,20 +75,20 @@ You see a condensed view of the conversation: the original request, recent turns (including tool results), and the customer's newest message. Return exactly one JSON object: -{"route": "weak" or "strong", "confidence": number 0..1, "abstain": boolean} +{"route": "efficient" or "capable", "confidence": number 0..1, "abstain": boolean} -State the route DIRECTLY: "weak" = the on-device assistant handles this turn; -"strong" = escalate this turn to the frontier model. Decide for the customer's +State the route DIRECTLY: "efficient" = the on-device assistant handles this turn; +"capable" = escalate this turn to the frontier model. Decide for the customer's NEWEST request, using the recent turns as context. -Route "weak" when the newest request is ROUTINE — the procedure is +Route "efficient" when the newest request is ROUTINE — the procedure is clear and it's about executing it: account/order/status lookups, reading or relaying tool results, standard single-step actions (toggle a setting, resend a code, restart a service), collecting information from the customer, confirmations, pleasantries, straightforward troubleshooting with an obvious next step. -Route "strong" when the newest request needs NON-OBVIOUS +Route "capable" when the newest request needs NON-OBVIOUS JUDGMENT the routine tier may get wrong: applying or reconciling POLICY with multiple conditions (eligibility, refunds, exceptions, proration), conflicts between what the customer wants and what diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index 8561816e7..75f103043 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -21,7 +21,7 @@ use http::StatusCode; use parking_lot::Mutex; use switchyard_libsy::{Algorithm, CallModel, LibsyError, Result, RoutingOutcome, drive}; use switchyard_protocol::{ - LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason, + Category, LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason, }; use switchyard_translation::prepare_request_for_target; @@ -44,6 +44,7 @@ pub async fn run( algorithm: Arc, clients: ClientRouter, request: Request, + models: HashMap>, observer: Option, ) -> Result<(ModelId, Response)> { let algorithm_name = algorithm.name().to_string(); @@ -52,7 +53,7 @@ pub async fn run( // This says if we have an observer, put Some(..) in routing_observations. // No observer means we don't want any routing_observations. let routing_observations = observer.as_ref().map(|_| Arc::new(Mutex::new(Vec::new()))); - let outcome = drive(algorithm, request, { + let outcome = drive(algorithm, request, models, { let routing_observations = routing_observations.clone(); move |call| serve(routing_clients.clone(), call, routing_observations.clone()) }) @@ -110,9 +111,10 @@ pub async fn decide( algorithm: Arc, clients: ClientRouter, request: Request, + models: HashMap>, ) -> Result { let routing_clients = clients.clone(); - let mut outcome = drive(algorithm, request, move |call| { + let mut outcome = drive(algorithm, request, models, move |call| { serve(routing_clients.clone(), call, None) }) .await?; @@ -453,9 +455,7 @@ mod tests { use crate::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; - struct CandidateAlgorithm { - models: Vec, - } + struct CandidateAlgorithm {} struct AnsweredAlgorithm { model: ModelId, @@ -469,13 +469,14 @@ mod tests { async fn route( self: Arc, - _driver: Driver, + driver: Driver, request: Request, ) -> Result { - let selected_model = self.models.first().cloned().ok_or(LibsyError::NoTargets)?; + let models = driver.models_for(Category::Any); + let selected_model = models.first().cloned().ok_or(LibsyError::NoTargets)?; Ok(RoutingOutcome::route_to( selected_model, - self.models.iter().skip(1).cloned().collect(), + models.iter().skip(1).cloned().collect(), request, )) } @@ -603,19 +604,27 @@ mod tests { requests: Mutex::new(Vec::new()), first, }); - let algorithm = Arc::new(CandidateAlgorithm { - models: vec!["weak".into(), "strong".into()], - }); + let algorithm = Arc::new(CandidateAlgorithm {}); + let models = to_category_map(&["weak", "strong"]); let result = run( algorithm, ClientRouter::single(client.clone()), request(), + models, None, ) .await; (client, result) } + fn to_category_map(names: &[&str]) -> HashMap> { + [( + Category::Any, + names.iter().map(|name| ModelId::from(*name)).collect(), + )] + .into() + } + #[tokio::test] async fn answered_outcome_does_not_make_a_second_model_call() -> Result<()> { let client = Arc::new(CandidateClient { @@ -633,6 +642,7 @@ mod tests { }), ClientRouter::single(client.clone()), request(), + HashMap::new(), Some(observer), ) .await?; @@ -697,11 +707,10 @@ mod tests { ); run( - Arc::new(CandidateAlgorithm { - models: vec!["weak".into(), "strong".into()], - }), + Arc::new(CandidateAlgorithm {}), clients, request(), + to_category_map(&["weak", "strong"]), None, ) .await?; @@ -733,6 +742,7 @@ mod tests { }), clients, request(), + HashMap::new(), ) .await?; @@ -763,11 +773,10 @@ mod tests { ); let outcome = decide( - Arc::new(CandidateAlgorithm { - models: vec!["weak".into(), "strong".into()], - }), + Arc::new(CandidateAlgorithm {}), clients, request(), + to_category_map(&["weak", "strong"]), ) .await?; @@ -904,10 +913,15 @@ mod tests { ]) .map_err(|error| LibsyError::external("building test client", error))?, ); - let algorithm = Arc::new(CandidateAlgorithm { - models: vec!["weak".into(), "strong".into()], - }); - run(algorithm, ClientRouter::single(client), request(), None).await?; + let algorithm = Arc::new(CandidateAlgorithm {}); + run( + algorithm, + ClientRouter::single(client), + request(), + to_category_map(&["weak", "strong"]), + None, + ) + .await?; assert_eq!(&*calls.lock(), &["weak", "weak", "weak", "strong"]); Ok(()) @@ -995,9 +1009,7 @@ mod tests { ]) .expect("building test client"), ); - let algorithm = Arc::new(CandidateAlgorithm { - models: vec!["weak".into(), "strong".into()], - }); + let algorithm = Arc::new(CandidateAlgorithm {}); let mut llm_request = text_request(Some("auto".to_string()), "hello".to_string()); llm_request.stream = true; let request = Request { @@ -1005,7 +1017,14 @@ mod tests { raw_request: None, metadata: None, }; - let result = run(algorithm, ClientRouter::single(client), request, None).await; + let result = run( + algorithm, + ClientRouter::single(client), + request, + to_category_map(&["weak", "strong"]), + None, + ) + .await; (server, calls, result) } diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 8ae0a551e..1225fdb84 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -10,7 +10,7 @@ //! and model name. Counters are cumulative across flushes; the helpers take the //! latest (max) matching data point. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::fmt; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, OnceLock}; @@ -39,7 +39,7 @@ use switchyard_libsy::{ StageRouterConfig, Step, TaskClassifierConfig, }; use switchyard_llm_client::{ClientRouter, RunObservation, RunObserver}; -use switchyard_protocol::ModelId; +use switchyard_protocol::{Category, ModelId}; use switchyard_protocol::{ ContentBlock, LlmRequest, LlmResponse, Message, Metadata, Request, Response, Role, RoutedLlmClient, ToolCall, ToolResult, Usage, WireFormat, @@ -585,19 +585,19 @@ async fn run( client: Arc, request: Request, ) -> switchyard_libsy::Result<(ModelId, Response)> { - switchyard_llm_client::run(algorithm, ClientRouter::single(client), request, None).await + switchyard_llm_client::run( + algorithm, + ClientRouter::single(client), + request, + HashMap::new(), + None, + ) + .await } -fn classifier_router( - judge_model: &str, - efficient_model: &str, - capable_model: &str, -) -> switchyard_libsy::Result> { +fn classifier_router() -> switchyard_libsy::Result> { Ok(Arc::new(LlmTaskClassifier::new( LlmClassifierConfig::Capability { - judge_target: ModelId::from(judge_model), - efficient_target: ModelId::from(efficient_model), - capable_target: ModelId::from(capable_model), config: TaskClassifierConfig { base_threshold: 0.5, ..TaskClassifierConfig::default() @@ -606,6 +606,40 @@ fn classifier_router( )?)) } +fn classifier_models( + judge_model: &str, + efficient_model: &str, + capable_model: &str, +) -> HashMap> { + [ + (Category::Judge, vec![judge_model.into()]), + (Category::Efficient, vec![efficient_model.into()]), + (Category::Capable, vec![capable_model.into()]), + ( + Category::Any, + vec![efficient_model.into(), capable_model.into()], + ), + ] + .into() +} + +async fn run_classifier( + judge_model: &str, + efficient_model: &str, + capable_model: &str, + client: Arc, + request: Request, +) -> switchyard_libsy::Result<(ModelId, Response)> { + switchyard_llm_client::run( + classifier_router()?, + ClientRouter::single(client), + request, + classifier_models(judge_model, efficient_model, capable_model), + None, + ) + .await +} + fn classifier_request() -> Request { Request { llm_request: text_request(Some("auto".to_string()), "classify this"), @@ -716,9 +750,6 @@ async fn affinity_keeps_the_algorithm_selection_after_client_fallback() efficient_available: AtomicBool::new(false), }); let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: "affinity-fallback-judge".into(), - efficient_target: "affinity-fallback-weak".into(), - capable_target: "affinity-fallback-strong".into(), config: TaskClassifierConfig { base_threshold: 0.5, classify_trigger: ClassifyTrigger::NewSession, @@ -731,6 +762,11 @@ async fn affinity_keeps_the_algorithm_selection_after_client_fallback() Arc::clone(&router), ClientRouter::single(client.clone()), request.clone(), + classifier_models( + "affinity-fallback-judge", + "affinity-fallback-weak", + "affinity-fallback-strong", + ), None, ) .await?; @@ -741,9 +777,18 @@ async fn affinity_keeps_the_algorithm_selection_after_client_fallback() ); client.efficient_available.store(true, Ordering::Relaxed); - let (selected, second_response) = - switchyard_llm_client::run(router, ClientRouter::single(client.clone()), request, None) - .await?; + let (selected, second_response) = switchyard_llm_client::run( + router, + ClientRouter::single(client.clone()), + request, + classifier_models( + "affinity-fallback-judge", + "affinity-fallback-weak", + "affinity-fallback-strong", + ), + None, + ) + .await?; assert_eq!(selected, "affinity-fallback-weak"); assert_eq!( second_response.served_model().map(ModelId::as_str), @@ -1042,12 +1087,10 @@ async fn stage_router_records_algorithm_owned_metrics() -> switchyard_libsy::Res let (_, exporter, provider, _, _) = telemetry(); const STRONG: &str = "obs-stage-strong"; const WEAK: &str = "obs-stage-weak"; - let target = |name: &str| name.to_string(); - let algorithm = Arc::new(StageRouter::new( - target(STRONG).into(), - target(WEAK).into(), - StageRouterConfig::new(PickerMode::EfficientFirst, 0.5), - )?) as Arc; + let algorithm = Arc::new(StageRouter::new(StageRouterConfig::new( + PickerMode::EfficientFirst, + 0.5, + ))?) as Arc; let request = Request { llm_request: LlmRequest { model: Some("auto".to_string()), @@ -1085,7 +1128,22 @@ async fn stage_router_records_algorithm_owned_metrics() -> switchyard_libsy::Res usage: Usage::default(), }) as Arc; - let (selected_model, _) = run(algorithm, client, request).await?; + let (selected_model, _) = switchyard_llm_client::run( + algorithm, + ClientRouter::single(client), + request, + [ + (Category::Capable, vec![ModelId::from(STRONG)]), + (Category::Efficient, vec![ModelId::from(WEAK)]), + ( + Category::Any, + vec![ModelId::from(STRONG), ModelId::from(WEAK)], + ), + ] + .into(), + None, + ) + .await?; assert_eq!(selected_model, STRONG); let snapshots = flushed_metrics(exporter, provider); @@ -1126,6 +1184,7 @@ async fn observed_run_reports_one_successful_routed_call() -> switchyard_libsy:: algo(ALGO, MODEL), ClientRouter::single(client), request_with_metadata("observed-session", "observed-correlation"), + HashMap::new(), Some(observer), ) .await?; @@ -1299,8 +1358,10 @@ async fn upstream_body_is_redacted_from_the_client_call_span() -> switchyard_lib judge_model: JUDGE.into(), outcome: JudgeOutcome::CallFailure, }) as Arc; - run( - classifier_router(JUDGE, "redaction-weak", "redaction-strong")?, + run_classifier( + JUDGE, + "redaction-weak", + "redaction-strong", client, classifier_request(), ) @@ -1333,7 +1394,10 @@ async fn failed_call_records_metrics_without_error_details() -> switchyard_libsy name: ALGO.to_string(), target: MODEL.into(), }); - let stream = algorithm.run_stream(request_with_metadata("obs-session-2", "obs-corr-2")); + let stream = algorithm.run_stream( + request_with_metadata("obs-session-2", "obs-corr-2"), + HashMap::new(), + ); tokio::pin!(stream); let mut saw_error_step = false; @@ -1420,9 +1484,8 @@ async fn classifier_metrics_count_routing_and_answer_calls_once() -> switchyard_ classifier_delay: Duration::from_millis(60), routed_delay: Duration::from_millis(200), }) as Arc; - let router = classifier_router("classifier", "weak", "strong")?; - - let (selected_model, _response) = run(router, client, classifier_request()).await?; + let (selected_model, _response) = + run_classifier("classifier", "weak", "strong", client, classifier_request()).await?; assert_eq!(selected_model, "weak"); @@ -1521,8 +1584,10 @@ async fn classifier_fail_open_records_each_failure_stage() -> switchyard_libsy:: judge_model: judge_model.into(), outcome, }) as Arc; - run( - classifier_router(judge_model, "fo-weak", "fo-strong")?, + run_classifier( + judge_model, + "fo-weak", + "fo-strong", client, classifier_request(), ) @@ -1564,7 +1629,10 @@ async fn in_flight_gauge_reads_a_run_parked_on_an_unanswered_routing_call() name: ALGO.to_string(), target: MODEL.into(), }); - let stream = algorithm.run_stream(request_with_metadata("obs-session-if", "obs-corr-if")); + let stream = algorithm.run_stream( + request_with_metadata("obs-session-if", "obs-corr-if"), + HashMap::new(), + ); tokio::pin!(stream); let attributes = [("algorithm", ALGO)]; @@ -1619,7 +1687,10 @@ async fn in_flight_gauge_clears_when_a_run_is_abandoned() -> switchyard_libsy::R // disconnected client abandons a run. The run task is aborted mid-await and never // reaches the code that follows it, so only a drop can return the count. { - let stream = algorithm.run_stream(request_with_metadata("obs-session-ab", "obs-corr-ab")); + let stream = algorithm.run_stream( + request_with_metadata("obs-session-ab", "obs-corr-ab"), + HashMap::new(), + ); tokio::pin!(stream); let Some(Ok(Step::CallModel(_call))) = stream.next().await else { return Err(test_error("expected an offloaded routing call")); diff --git a/crates/libsy/src/algorithms/composite.rs b/crates/libsy/src/algorithms/composite.rs index af75f7b64..3e71e813c 100644 --- a/crates/libsy/src/algorithms/composite.rs +++ b/crates/libsy/src/algorithms/composite.rs @@ -16,13 +16,13 @@ use super::fall_through::FallThrough; use super::llm_class::{LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig}; use super::stage::{StageRouterConfig, build_stage_route}; use super::util::affinity::{ClassifyTrigger, evict_if_full, has_new_user_turn, retention_key}; -use super::util::stage::{StageTargets, Tier, set_fall_open}; +use super::util::stage::{Tier, set_fall_open}; use crate::core::algorithm::{Algorithm, Driver, RoutingIdentity}; use crate::core::classifier::Classifier; use crate::core::processor::{Event, Processor}; use crate::core::state::State; use crate::{LibsyError, Result}; -use switchyard_protocol::{ModelId, Request}; +use switchyard_protocol::{Category, ModelId, Request}; const COMPOSITE: &str = "composite"; @@ -33,13 +33,23 @@ const COMPOSITE: &str = "composite"; /// into state on every request so the cascade below reads it. struct TierSetter { judge: Arc>, - targets: StageTargets, trigger: ClassifyTrigger, message_hash_fallback: bool, tiers: Mutex>, } impl TierSetter { + fn tier_for(driver: Option<&Driver>, target: &ModelId) -> Option { + let driver = driver?; + if driver.models_for(Category::Capable).first() == Some(target) { + Some(Tier::Capable) + } else if driver.models_for(Category::Efficient).first() == Some(target) { + Some(Tier::Efficient) + } else { + None + } + } + /// Two requests for one identity can both pass this and both judge, since a /// judge call sits between here and [`retain`](Self::retain). The later wins. fn is_due(&self, identity: Option<&RoutingIdentity>, request: &Request) -> bool { @@ -76,7 +86,7 @@ impl Processor for TierSetter { if self.is_due(identity.as_ref(), request) { let (classification, _) = self.judge.score(state, request, driver).await?; if let Some(winner) = classification.argmax(false)? - && let Some(tier) = self.targets.tier_for(&winner.target) + && let Some(tier) = Self::tier_for(driver, &winner.target) { set_fall_open(state, tier); if let Some(identity) = identity { @@ -100,8 +110,6 @@ impl Processor for TierSetter { /// A judge stacked over a stage router. pub struct CompositeRouterConfig { - /// Target the judge is called through. Not a routing destination. - pub judge_target: ModelId, /// Judge settings, including how often `classify_trigger` runs it. pub judge: TaskClassifierConfig, /// Serves the turns, with the tier the judge picked as its fall-open default. @@ -120,11 +128,7 @@ impl CompositeRouter { /// /// A stage router carrying its own judge is allowed, but that judge sits ahead /// of the fall-open tier and so answers most of the turns this one set a tier for. - pub fn new( - capable: ModelId, - efficient: ModelId, - config: CompositeRouterConfig, - ) -> Result { + pub fn new(config: CompositeRouterConfig) -> Result { if config.judge.classify_trigger == ClassifyTrigger::EveryRequest { return Err(LibsyError::AlgorithmError { message: "composite: classify_trigger must be user_turn or new_session".to_string(), @@ -141,19 +145,15 @@ impl CompositeRouter { ..config.judge }; let judge = LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: config.judge_target, - efficient_target: efficient.clone(), - capable_target: capable.clone(), config: judge_config, })?; let setter = TierSetter { judge: Arc::new(judge), - targets: StageTargets::new(capable.clone(), efficient.clone()), trigger, message_hash_fallback, tiers: Mutex::new(HashMap::new()), }; - let route = build_stage_route(capable, efficient, config.stage)? + let route = build_stage_route(config.stage)? .with_name(COMPOSITE) .with_processor(Arc::new(setter)); Ok(Self { route }) @@ -177,14 +177,28 @@ impl Algorithm for CompositeRouter { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; - use switchyard_protocol::{Message, Role}; + use switchyard_protocol::{Category, Message, Role}; use super::*; use crate::algorithms::util::stage::PickerMode; use crate::algorithms::util::tier_fixtures::{JUDGE, Recorder, turn_request}; - use crate::core::testing::test_drive; + use crate::core::testing::test_drive_with_models; + + fn runtime_models() -> HashMap> { + [ + (Category::Judge, vec![ModelId::from(JUDGE)]), + (Category::Efficient, vec![ModelId::from("weak")]), + (Category::Capable, vec![ModelId::from("strong")]), + ( + Category::Any, + vec![ModelId::from("strong"), ModelId::from("weak")], + ), + ] + .into() + } fn user_turn_request() -> Request { let mut request = turn_request(false); @@ -204,47 +218,36 @@ mod tests { } fn hash_keyed_router() -> Result> { - Ok(Arc::new(CompositeRouter::new( - ModelId::from("strong"), - ModelId::from("weak"), - CompositeRouterConfig { - judge_target: ModelId::from(JUDGE), - judge: TaskClassifierConfig { - base_threshold: 0.5, - classify_trigger: ClassifyTrigger::UserTurn, - message_hash_fallback: true, - ..Default::default() - }, - stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5), + Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig { + judge: TaskClassifierConfig { + base_threshold: 0.5, + classify_trigger: ClassifyTrigger::UserTurn, + message_hash_fallback: true, + ..Default::default() }, - )?)) + stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5), + })?)) } fn router() -> Result> { - Ok(Arc::new(CompositeRouter::new( - ModelId::from("strong"), - ModelId::from("weak"), - CompositeRouterConfig { - judge_target: ModelId::from(JUDGE), - judge: TaskClassifierConfig { - base_threshold: 0.5, - classify_trigger: ClassifyTrigger::UserTurn, - ..Default::default() - }, - stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5), + Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig { + judge: TaskClassifierConfig { + base_threshold: 0.5, + classify_trigger: ClassifyTrigger::UserTurn, + ..Default::default() }, - )?)) + stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5), + })?)) } #[test] fn rejects_every_request_as_a_trigger() { let config = CompositeRouterConfig { - judge_target: ModelId::from(JUDGE), judge: TaskClassifierConfig::default(), stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5), }; assert!(matches!( - CompositeRouter::new(ModelId::from("strong"), ModelId::from("weak"), config), + CompositeRouter::new(config), Err(LibsyError::AlgorithmError { .. }) )); } @@ -255,15 +258,17 @@ mod tests { *recorder.judge_p_solve.lock() = 0.1; let router = hash_keyed_router()?; - test_drive( + test_drive_with_models( router.clone(), unkeyed(user_turn_request()), + runtime_models(), recorder.serve(), ) .await?; - test_drive( + test_drive_with_models( router.clone(), unkeyed(turn_request(false)), + runtime_models(), recorder.serve(), ) .await?; @@ -287,8 +292,20 @@ mod tests { *recorder.judge_p_solve.lock() = 0.1; let router = router()?; - test_drive(router.clone(), user_turn_request(), recorder.serve()).await?; - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + user_turn_request(), + runtime_models(), + recorder.serve(), + ) + .await?; + test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; let routed = recorder.routed(); assert_eq!( diff --git a/crates/libsy/src/algorithms/escalation.rs b/crates/libsy/src/algorithms/escalation.rs index 47523dc13..747dcf418 100644 --- a/crates/libsy/src/algorithms/escalation.rs +++ b/crates/libsy/src/algorithms/escalation.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use async_trait::async_trait; use switchyard_protocol::{ - AggLlmResponse, LlmClientError, LlmResponse, Message, ModelId, Request, Response, Role, + AggLlmResponse, Category, LlmClientError, LlmResponse, Message, Request, Response, Role, }; use super::util::classifier_contract::ClassifierContractConfig; @@ -44,33 +44,19 @@ fn assistant_message(response: &AggLlmResponse) -> Message { /// not pay for a second model call. struct EscalationClassifier { judge: JudgeClassifier, - capable: ModelId, - efficient: ModelId, /// Consecutive escalate verdicts required to latch. confirmations: u32, } /// Builds the escalation classifier used by the shared LLM classifier route shell. pub(super) fn build_classifier( - judge_target: ModelId, - efficient_target: &ModelId, - capable_target: &ModelId, contract_config: ClassifierContractConfig, config: EscalationJudgeConfig, max_output_tokens: u64, ) -> Result>> { let confirmations = config.confirmations; let classifier: Arc> = Arc::new(EscalationClassifier { - judge: escalation::build_judge( - judge_target, - capable_target.clone(), - efficient_target.clone(), - &contract_config, - config, - max_output_tokens, - )?, - capable: capable_target.clone(), - efficient: efficient_target.clone(), + judge: escalation::build_judge(&contract_config, config, max_output_tokens)?, confirmations, }); Ok(classifier) @@ -89,6 +75,8 @@ impl Classifier for EscalationClassifier { message: "escalation classifier requires a driver".into(), }); }; + let capable = driver.first_model_for(Category::Capable)?.clone(); + let efficient = driver.first_model_for(Category::Efficient)?.clone(); // A confirmed session stays capable without a judge call. if streak(state) >= self.confirmations { @@ -96,7 +84,7 @@ impl Classifier for EscalationClassifier { "source": "escalation", "verdict": "latched", })); - return Ok((decisive(&self.capable), None)); + return Ok((decisive(&capable), None)); } // Call efficient model and buffer the response so the judge can read it. @@ -104,11 +92,11 @@ impl Classifier for EscalationClassifier { // If the efficient model exceeds its context window, fall through to capable. This call // deliberately has one candidate so the classifier sees the efficient model's error. tracing::info!( - target = %self.efficient, + target = %efficient, "escalation classifier selected efficient tier" ); let efficient_response = match driver - .call_model(request.clone(), vec![self.efficient.clone()]) + .call_model(request.clone(), vec![efficient.clone()]) .await { Ok(r) => r, @@ -120,7 +108,7 @@ impl Classifier for EscalationClassifier { "source": "fallback", "reason_code": "context_window", })); - return Ok((decisive(&self.capable), None)); + return Ok((decisive(&capable), None)); } Err(e) => return Err(e), }; @@ -133,10 +121,10 @@ impl Classifier for EscalationClassifier { "source": "fallback", "reason_code": "transport", })); - return Ok((decisive(&self.capable), None)); + return Ok((decisive(&capable), None)); } Err(source) => { - return Err(LibsyError::client_call(self.efficient.clone(), source)); + return Err(LibsyError::client_call(efficient.clone(), source)); } }; // Append the efficient reply so the judge reads this turn's completed trajectory. @@ -162,7 +150,7 @@ impl Classifier for EscalationClassifier { let held = streak(state); let best = classification.argmax(false)?; let (escalate, pending) = match &best { - Some(score) if score.target == self.capable => (true, held + 1), + Some(score) if score.target == capable => (true, held + 1), Some(_) => (false, 0), None => (false, held), }; @@ -176,7 +164,7 @@ impl Classifier for EscalationClassifier { "source": "escalation", "verdict": "escalate", })); - return Ok((decisive(&self.capable), None)); + return Ok((decisive(&capable), None)); } if escalate { @@ -186,25 +174,25 @@ impl Classifier for EscalationClassifier { })); } - Ok((decisive(&self.efficient), Some(efficient_response))) + Ok((decisive(&efficient), Some(efficient_response))) } } #[cfg(test)] mod tests { - use std::collections::VecDeque; + use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use parking_lot::Mutex; use switchyard_protocol::{ - ContentBlock, LlmClientError, LlmResponse, LlmResponseChunk, Metadata, Request, Response, - completion_text, text_request, text_response, + ContentBlock, LlmClientError, LlmResponse, LlmResponseChunk, Metadata, ModelId, Request, + Response, completion_text, text_request, text_response, }; use super::*; use crate::algorithms::llm_class::{LlmClassifierConfig, LlmTaskClassifier}; use crate::algorithms::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS; - use crate::core::testing::{Serve, reply, test_drive}; + use crate::core::testing::{Serve, reply, test_drive_with_models}; /// A queue of replies, drained in order. struct Queue(Mutex>); @@ -259,6 +247,19 @@ mod tests { } } + fn runtime_models() -> HashMap> { + [ + (Category::Judge, vec![ModelId::from("judge")]), + (Category::Efficient, vec![ModelId::from("efficient")]), + (Category::Capable, vec![ModelId::from("capable")]), + ( + Category::Any, + vec![ModelId::from("capable"), ModelId::from("efficient")], + ), + ] + .into() + } + /// Returns a stream that emits partial content before failing during aggregation. fn streamed_then_error(error: LlmClientError) -> Response { Response { @@ -278,9 +279,6 @@ mod tests { fn escalation_router() -> Result> { Ok(Arc::new(LlmTaskClassifier::new( LlmClassifierConfig::Escalation { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), contract: ClassifierContractConfig::default(), config: EscalationJudgeConfig { confirmations: 1, @@ -296,9 +294,10 @@ mod tests { let judge = Queue::new([r#"{"escalate":false,"reason":"progressing"}"#]); let model = Queue::new(["efficient answer"]); - let (selected_model, response) = test_drive( + let (selected_model, response) = test_drive_with_models( escalation_router()?, classify_request(), + runtime_models(), queued(model, judge), ) .await?; @@ -334,9 +333,6 @@ mod tests { } }; let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Escalation { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), contract: ClassifierContractConfig::default().with_prompt("Custom trajectory rubric."), config: EscalationJudgeConfig { confirmations: 1, @@ -345,7 +341,7 @@ mod tests { max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS, })?); - test_drive(router, classify_request(), serve).await?; + test_drive_with_models(router, classify_request(), runtime_models(), serve).await?; assert_eq!(&*prompts.lock(), &["Custom trajectory rubric."]); Ok(()) @@ -356,9 +352,10 @@ mod tests { let judge = Queue::new([r#"{"escalate":true,"reason":"stuck in a loop"}"#]); let model = Queue::new(["efficient draft", "capable answer"]); - let (selected_model, response) = test_drive( + let (selected_model, response) = test_drive_with_models( escalation_router()?, classify_request(), + runtime_models(), queued(model, judge), ) .await?; @@ -378,13 +375,15 @@ mod tests { let router = escalation_router()?; let request = classify_session_request(); - test_drive( + test_drive_with_models( router.clone(), request.clone(), + runtime_models(), queued(Arc::clone(&model), Arc::clone(&judge)), ) .await?; - let (selected_model, _) = test_drive(router, request, queued(model, judge)).await?; + let (selected_model, _) = + test_drive_with_models(router, request, runtime_models(), queued(model, judge)).await?; assert_eq!(selected_model, "capable"); Ok(()) @@ -403,8 +402,13 @@ mod tests { } }; - let (selected_model, response) = - test_drive(escalation_router()?, classify_request(), serve).await?; + let (selected_model, response) = test_drive_with_models( + escalation_router()?, + classify_request(), + runtime_models(), + serve, + ) + .await?; assert_eq!(selected_model, "capable"); assert_eq!( @@ -440,7 +444,8 @@ mod tests { let mut request = classify_request(); request.llm_request.stream = true; - let result = test_drive(escalation_router()?, request, serve).await; + let result = + test_drive_with_models(escalation_router()?, request, runtime_models(), serve).await; assert_eq!(&*calls.lock(), &["efficient", "capable"]); let (_, response) = result?; @@ -465,7 +470,7 @@ mod tests { let mut request = classify_request(); request.llm_request.stream = true; - match test_drive(escalation_router()?, request, serve).await { + match test_drive_with_models(escalation_router()?, request, runtime_models(), serve).await { Err(LibsyError::ClientCall { target, source: LlmClientError::InvalidResponse { .. }, diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs index c2cebe68a..0b181268a 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -26,11 +26,11 @@ use async_trait::async_trait; use parking_lot::Mutex; use tokio::sync::Mutex as AsyncMutex; -use crate::core::algorithm::{self, Algorithm, Driver}; +use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::processor::{Event, Processor}; use crate::{LibsyError, Result, RoutingOutcome}; -use switchyard_protocol::{ModelId, Request, Response}; +use switchyard_protocol::{Category, ModelId, Request, Response}; struct SessionState { state: Arc>, @@ -96,19 +96,17 @@ pub struct FallThrough { name: String, processors: Vec>>, classifiers: Vec>>, - targets: Vec, session_states: Option>>, cleanup_started: Once, } impl FallThrough<()> { /// Creates an empty stateless router. - pub fn new(targets: Vec) -> Self { + pub fn new(_targets: Vec) -> Self { Self { name: "fall_through".to_string(), processors: Vec::new(), classifiers: Vec::new(), - targets, session_states: None, cleanup_started: Once::new(), } @@ -120,12 +118,11 @@ where S: Default + Send + 'static, { /// Creates a router that retains one private `S` per session. - pub fn new_with_state(targets: Vec) -> Self { + pub fn new_with_state(_targets: Vec) -> Self { Self { name: "fall_through".to_string(), processors: Vec::new(), classifiers: Vec::new(), - targets, session_states: Some(Arc::new(Mutex::new(HashMap::new()))), cleanup_started: Once::new(), } @@ -199,7 +196,13 @@ where match served { Some(response) => Ok(RoutingOutcome::answered(target, request, response)), None => { - let fallback_models = self.fallbacks(&target); + // Every configured target other than the selection, in fallback order. + let fallback_models: Vec = driver + .models_for(Category::Any) + .iter() + .filter(|candidate| **candidate != target) + .cloned() + .collect(); Ok(RoutingOutcome::route_to(target, fallback_models, request)) } } @@ -212,15 +215,6 @@ where } } - /// Every configured target other than the selection, in fallback order. - fn fallbacks(&self, target: &ModelId) -> Vec { - self.targets - .iter() - .filter(|candidate| *candidate != target) - .cloned() - .collect() - } - /// Returns this request's retained state without holding the registry lock. fn session_state(&self, request: &Request) -> Option>> { let states = self.session_states.as_ref()?; @@ -270,7 +264,6 @@ where }; // 3. Resolve the target and log the choice. - algorithm::ensure_model_is_target(&self.targets, &score.target)?; let target = score.target.clone(); tracing::info!(algorithm=self.name, target=%score.target, confidence=score.confidence, "Model selected"); @@ -341,8 +334,6 @@ mod tests { use super::*; use crate::algorithms::util::prompts; use crate::core::classifier::Classification; - use crate::{SystemPromptProcessor, TargetPrompts}; - use crate::core::testing::{Serve, echo, reply, test_drive}; use switchyard_protocol::{LlmRequest, Message, Metadata, Role, completion_text, text_request}; @@ -368,95 +359,12 @@ mod tests { } } - const CAPABLE_PROMPT: &str = "diagnose before you edit"; - const EFFICIENT_PROMPT: &str = "follow the settled plan"; const NOTE: &str = "the previous model was stalling"; - /// One model call as the prompt and note tests observe it. - #[derive(Clone, Debug, Default)] - struct RecordedCall { - target: String, - messages: Vec, - instructions: Vec, - } - - /// Captures the prompt-bearing request that reached the selected target. - #[derive(Default)] - struct PromptRecorder(Mutex>); - - impl PromptRecorder { - fn serve(self: &Arc) -> impl Serve { - let recorder = Arc::clone(self); - move |target: ModelId, request: Request| { - let recorder = Arc::clone(&recorder); - async move { - *recorder.0.lock() = Some(RecordedCall { - target: target.to_string(), - messages: request - .llm_request - .messages - .iter() - .filter_map(|message| message.text_content("|")) - .collect(), - instructions: request - .llm_request - .instructions - .iter() - .filter_map(|block| block.content.iter().find_map(text_of)) - .collect(), - }); - Ok(reply(target)) - } - } - } - } - - fn text_of(block: &switchyard_protocol::ContentBlock) -> Option { - match block { - switchyard_protocol::ContentBlock::Text { text } => Some(text.clone()), - _ => None, - } - } - fn target_set(names: &[&str]) -> Vec { names.iter().map(|name| ModelId::from(*name)).collect() } - fn target_prompts() -> TargetPrompts { - TargetPrompts::default() - .with("capable", CAPABLE_PROMPT) - .with("efficient", EFFICIENT_PROMPT) - } - - /// Routes one turn on a prompt test cascade and returns the recorded model call. - async fn routed_prompt_call( - recorder: &Arc, - router: FallThrough, - ) -> Result { - test_drive( - Arc::new(router), - Request { - llm_request: text_request(Some("auto".to_string()), "fix the build"), - raw_request: None, - metadata: None, - }, - recorder.serve(), - ) - .await?; - let call = recorder.0.lock().take(); - match call { - Some(call) => Ok(call), - None => panic!("the model was never called"), - } - } - - /// A prompt cascade that always routes to `target`. - fn prompt_router(target: &str, prompts: TargetPrompts) -> FallThrough { - FallThrough::new(target_set(&["capable", "efficient"])) - .with_processor(Arc::new(SystemPromptProcessor::new(prompts))) - .with_classifier(Arc::new(DefaultTarget::new(target))) - } - /// A classifier that emits fixed scores (empty = abstain). struct FixedClassifier(Vec); @@ -546,64 +454,6 @@ mod tests { // --- tests ------------------------------------------------------------------------- - #[tokio::test] - async fn each_target_gets_its_own_prompt() -> Result<()> { - for (target, expected) in [("capable", CAPABLE_PROMPT), ("efficient", EFFICIENT_PROMPT)] { - let recorder = Arc::new(PromptRecorder::default()); - let call = - routed_prompt_call(&recorder, prompt_router(target, target_prompts())).await?; - assert_eq!(call.target, target); - assert_eq!(call.instructions, vec![expected.to_string()]); - } - Ok(()) - } - - #[tokio::test] - async fn a_target_with_no_prompt_is_left_untouched() -> Result<()> { - let recorder = Arc::new(PromptRecorder::default()); - let only_capable = TargetPrompts::default().with("capable", CAPABLE_PROMPT); - - let call = routed_prompt_call(&recorder, prompt_router("efficient", only_capable)).await?; - - assert!( - call.instructions.is_empty(), - "one target's prompt must not leak onto another: {:?}", - call.instructions - ); - Ok(()) - } - - #[tokio::test] - async fn the_prompt_follows_the_target_whichever_classifier_picked_it() -> Result<()> { - // The first classifier abstains, so the second decides; the prompt follows the - // target the cascade settled on rather than the classifier that named it. - struct Abstains; - - #[async_trait] - impl Classifier for Abstains { - async fn score( - &self, - _state: &mut (), - _request: &mut Request, - _driver: Option<&Driver>, - ) -> Result<(Classification, Option)> { - Ok((Classification::Ambiguous(Vec::new()), None)) - } - } - - let recorder = Arc::new(PromptRecorder::default()); - let router = FallThrough::new(target_set(&["capable", "efficient"])) - .with_processor(Arc::new(SystemPromptProcessor::new(target_prompts()))) - .with_classifier(Arc::new(Abstains)) - .with_classifier(Arc::new(DefaultTarget::new("capable"))); - - let call = routed_prompt_call(&recorder, router).await?; - - assert_eq!(call.target, "capable"); - assert_eq!(call.instructions, vec![CAPABLE_PROMPT.to_string()]); - Ok(()) - } - #[tokio::test] async fn a_note_reaches_the_model_in_the_conversation() -> Result<()> { // Appends a note to every outbound request, the way a router would on a turn it @@ -620,15 +470,37 @@ mod tests { } } - let recorder = Arc::new(PromptRecorder::default()); + let captured = Arc::new(Mutex::new(None)); let router = FallThrough::new(target_set(&["capable", "efficient"])) .with_processor(Arc::new(Noting)) .with_classifier(Arc::new(DefaultTarget::new("capable"))); - let call = routed_prompt_call(&recorder, router).await?; + test_drive( + Arc::new(router), + Request { + llm_request: text_request(Some("auto".to_string()), "fix the build"), + raw_request: None, + metadata: None, + }, + capturing(Arc::clone(&captured)), + ) + .await?; + let request = captured + .lock() + .take() + .ok_or_else(|| LibsyError::external("test", TestError("the model was never called")))?; + let messages: Vec = request + .llm_request + .messages + .iter() + .filter_map(|message| message.text_content("|")) + .collect(); - assert_eq!(call.messages, vec![format!("fix the build|{NOTE}")]); - assert!(call.instructions.is_empty(), "a note is not an instruction"); + assert_eq!(messages, vec![format!("fix the build|{NOTE}")]); + assert!( + request.llm_request.instructions.is_empty(), + "a note is not an instruction" + ); Ok(()) } @@ -637,10 +509,11 @@ mod tests { use futures::StreamExt; let router = Arc::new( - FallThrough::<()>::new(target_set(&["weak", "mid", "strong"])) - .with_classifier(fixed(vec![score("mid", 0.9)])), + FallThrough::<()>::new(vec![]).with_classifier(fixed(vec![score("mid", 0.9)])), ); - let stream = router.run_stream(request()); + + let models = Category::to_map(Category::Any, &["weak", "mid", "strong"]); + let stream = router.run_stream(request(), models); tokio::pin!(stream); while let Some(step) = stream.next().await { if let crate::Step::Done(outcome) = step? { diff --git a/crates/libsy/src/algorithms/llm_class.rs b/crates/libsy/src/algorithms/llm_class.rs index 652d647cf..19be971c0 100644 --- a/crates/libsy/src/algorithms/llm_class.rs +++ b/crates/libsy/src/algorithms/llm_class.rs @@ -3,16 +3,16 @@ //! Judge-backed capability, escalation, and custom-policy routing. -use std::collections::{BTreeMap, BTreeSet, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use async_trait::async_trait; use serde::{Deserialize, Deserializer}; use serde_json::Value; -use switchyard_protocol::{ContentBlock, Message, ModelId, Role}; +use switchyard_protocol::{Category, ContentBlock, Message, Role}; use super::escalation; -use super::fall_through::{DefaultTarget, FallThrough}; +use super::fall_through::FallThrough; use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS; use super::util::affinity::{AffinityRouter, ClassifyTrigger}; use super::util::classifier_contract::{ @@ -24,7 +24,7 @@ use super::util::llm_judge::{ SerdeDecoder, StructuredJudge, }; use super::util::target_selector::TargetSelectorPolicy; -use crate::core::algorithm::{self, Algorithm, Driver}; +use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::state::State; use crate::{LibsyError, Result}; @@ -186,24 +186,17 @@ impl ClassifierInput for TaskInput { } } +#[cfg(test)] type CapabilityJudge = StructuredJudge>; struct TaskClassifierPolicy { - efficient_target: ModelId, - capable_target: ModelId, base_threshold: f64, threshold_step: f64, } impl TaskClassifierPolicy { - fn new( - efficient_target: impl Into, - capable_target: impl Into, - config: &TaskClassifierConfig, - ) -> Self { + fn new(config: &TaskClassifierConfig) -> Self { Self { - efficient_target: efficient_target.into(), - capable_target: capable_target.into(), base_threshold: config.base_threshold, threshold_step: config.threshold_step, } @@ -218,28 +211,32 @@ impl TaskClassifierPolicy { impl JudgePolicy for TaskClassifierPolicy { type Verdict = TaskClassifierVerdict; - fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification { + fn to_classification( + &self, + verdict: Option<&Self::Verdict>, + driver: &Driver, + ) -> Result { // Judge output is untrusted. An absent, invalid, or inconsistent verdict is // ambiguous so the surrounding router applies its configured fallback. let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else { - return Classification::Ambiguous(vec![]); + return Ok(Classification::Ambiguous(vec![])); }; // A usable verdict below the capability threshold is still a decision: the judge // does not trust the efficient tier with this task. let Some(threshold) = self.threshold(verdict) else { - return Classification::Ambiguous(vec![]); + return Ok(Classification::Ambiguous(vec![])); }; - let target = if verdict.p_solve >= threshold + let category = if verdict.p_solve >= threshold || (threshold - verdict.p_solve).abs() <= f64::EPSILON { - &self.efficient_target + Category::Efficient } else { - &self.capable_target + Category::Capable }; - Classification::Scores(vec![Score { - target: target.clone(), + Ok(Classification::Scores(vec![Score { + target: driver.first_model_for(category)?.clone(), confidence: 1.0, - }]) + }])) } } @@ -400,7 +397,7 @@ impl TaskClassifierConfig { /// Policy that maps a custom classifier verdict to a routing target. #[derive(Clone, Debug)] pub enum CustomClassifierPolicy { - /// Resolves a JSON Pointer and treats its string value as a configured target label. + /// Resolves a JSON Pointer and treats its string value as a model category. TargetSelector { /// JSON Pointer evaluated against each schema-validated verdict. selector: String, @@ -408,7 +405,7 @@ pub enum CustomClassifierPolicy { } impl CustomClassifierPolicy { - /// Creates a policy that selects a target label through a JSON Pointer. + /// Creates a policy that selects a model category through a JSON Pointer. pub fn target_selector(selector: impl Into) -> Self { Self::TargetSelector { selector: selector.into(), @@ -476,18 +473,17 @@ enum CustomPolicyRuntime { impl JudgePolicy for CustomPolicyRuntime { type Verdict = Value; - fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification { + fn to_classification( + &self, + verdict: Option<&Self::Verdict>, + driver: &Driver, + ) -> Result { match self { - Self::TargetSelector(policy) => policy.to_classification(verdict), + Self::TargetSelector(policy) => policy.to_classification(verdict, driver), } } } -struct TaskClassifier { - classifier: JudgeClassifier, - capable_target: ModelId, -} - /// Builds the affinity router a trigger calls for, if any. fn affinity_router( trigger: ClassifyTrigger, @@ -514,34 +510,47 @@ pub struct LlmTaskClassifier { } struct ClassifierRouteConfig { - default_target: ModelId, + default_target: Category, classify_trigger: ClassifyTrigger, message_hash_fallback: bool, } +struct DefaultCategoryClassifier(Category); + +#[async_trait] +impl Classifier for DefaultCategoryClassifier { + async fn score( + &self, + _state: &mut State, + _request: &mut Request, + driver: Option<&Driver>, + ) -> Result<(Classification, Option)> { + let Some(driver) = driver else { + return Err(LibsyError::AlgorithmError { + message: "default category requires a driver".to_string(), + }); + }; + Ok(( + Classification::Scores(vec![Score { + target: driver.first_model_for(self.0)?.clone(), + confidence: 0.0, + }]), + None, + )) + } +} + /// Complete construction settings for one LLM classifier mode. #[derive(Clone)] #[non_exhaustive] pub enum LlmClassifierConfig { /// Routes between efficient and capable targets from a task-level verdict. Capability { - /// Target that produces classifier verdicts. - judge_target: ModelId, - /// Target used when the efficient tier can handle the task. - efficient_target: ModelId, - /// Target used when the task needs the capable tier. - capable_target: ModelId, /// Capability classifier settings. config: TaskClassifierConfig, }, /// Judges efficient responses and escalates after a confirmed streak. Escalation { - /// Target that produces escalation verdicts. - judge_target: ModelId, - /// Target called before each escalation decision. - efficient_target: ModelId, - /// Target used after escalation is confirmed. - capable_target: ModelId, /// Prompt and verdict contract settings for the escalation judge. contract: ClassifierContractConfig, /// Escalation policy settings. @@ -549,14 +558,10 @@ pub enum LlmClassifierConfig { /// Maximum completion tokens available to the escalation verdict. max_output_tokens: u64, }, - /// Routes among named targets using a user-supplied schema and policy. + /// Routes among model categories using a user-supplied schema and policy. Custom { - /// Target that produces classifier verdicts. - judge_target: ModelId, - /// User-facing labels paired with their resolved routing targets. - targets: Vec<(String, ModelId)>, - /// Label selected when the judge does not produce a usable verdict. - default_target: String, + /// Category selected when the judge does not produce a usable verdict. + default_target: Category, /// Custom classifier settings. config: CustomClassifierConfig, }, @@ -571,49 +576,26 @@ impl LlmTaskClassifier { /// settings are invalid. pub fn new(config: LlmClassifierConfig) -> Result { match config { - LlmClassifierConfig::Capability { - judge_target, - efficient_target, - capable_target, - config, - } => Self::build_capability(judge_target, efficient_target, capable_target, config), + LlmClassifierConfig::Capability { config } => Self::build_capability(config), LlmClassifierConfig::Escalation { - judge_target, - efficient_target, - capable_target, - contract, - config, - max_output_tokens, - } => Self::build_escalation( - judge_target, - efficient_target, - capable_target, contract, config, max_output_tokens, - ), + } => Self::build_escalation(contract, config, max_output_tokens), LlmClassifierConfig::Custom { - judge_target, - targets, default_target, config, - } => Self::build_custom(judge_target, targets, default_target, config), + } => Self::build_custom(default_target, config), } } - fn build_capability( - judge_target: ModelId, - efficient_target: ModelId, - capable_target: ModelId, - config: TaskClassifierConfig, - ) -> Result { + fn build_capability(config: TaskClassifierConfig) -> Result { config.validate()?; let contract = Self::load_capability_contract(&config.contract)?; - let targets = vec![efficient_target.clone(), capable_target.clone()]; let classify_trigger = config.classify_trigger; let message_hash_fallback = config.message_hash_fallback; - let classifier = Arc::new(TaskClassifier { - classifier: JudgeClassifier::new( + let classifier: Arc> = Arc::new( + JudgeClassifier::new( StructuredJudge::new( TaskInput { recent_turn_window: config.recent_turn_window, @@ -622,75 +604,22 @@ impl LlmTaskClassifier { SerdeDecoder::new(), JudgeRuntimeConfig::new(config.max_output_tokens)?, ), - judge_target.clone(), - TaskClassifierPolicy::new( - efficient_target.clone(), - capable_target.clone(), - &config, - ), + TaskClassifierPolicy::new(&config), ) .with_evidence(capability_evidence), - capable_target: capable_target.clone(), - }); - let inner: Arc> = classifier.clone(); + ); Self::from_classifier( - targets, - inner, + classifier, ClassifierRouteConfig { - default_target: classifier.capable_target.clone(), + default_target: Category::Capable, classify_trigger, message_hash_fallback, }, ) } - fn build_custom( - judge_target: ModelId, - targets: Vec<(String, ModelId)>, - default_target: String, - config: CustomClassifierConfig, - ) -> Result { + fn build_custom(default_target: Category, config: CustomClassifierConfig) -> Result { config.validate()?; - if targets.len() < 2 { - return Err(LibsyError::AlgorithmError { - message: "custom classifier requires at least two targets".to_string(), - }); - } - - let mut labels = BTreeSet::new(); - let mut resolved_names = BTreeSet::new(); - let mut target_map = BTreeMap::new(); - let mut resolved_targets = Vec::with_capacity(targets.len()); - for (label, target) in targets { - if label.trim().is_empty() || label.trim() != label { - return Err(LibsyError::AlgorithmError { - message: "custom classifier target labels must be non-empty and have no surrounding whitespace" - .to_string(), - }); - } - if !labels.insert(label.clone()) { - return Err(LibsyError::AlgorithmError { - message: format!("custom classifier target label {label:?} is duplicated"), - }); - } - if !resolved_names.insert(target.clone()) { - return Err(LibsyError::AlgorithmError { - message: format!("custom classifier resolved target {target:?} is duplicated"), - }); - } - target_map.insert(label, target.clone()); - resolved_targets.push(target); - } - let default_name = - target_map - .get(&default_target) - .cloned() - .ok_or_else(|| LibsyError::AlgorithmError { - message: format!( - "default_target {default_target:?} must be one of the configured targets" - ), - })?; - let CustomClassifierConfig { prompt, response_schema, @@ -703,9 +632,7 @@ impl LlmTaskClassifier { let contract = ClassifierContract::from_inner_schema(&prompt, response_schema)?; let policy = match policy { CustomClassifierPolicy::TargetSelector { selector } => { - CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new( - selector, target_map, - )?) + CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(selector)?) } }; let classifier: Arc> = Arc::new(JudgeClassifier::new( @@ -715,15 +642,13 @@ impl LlmTaskClassifier { JsonSchemaDecoder::new(), JudgeRuntimeConfig::new(max_output_tokens)?, ), - judge_target, policy, )); Self::from_classifier( - resolved_targets, classifier, ClassifierRouteConfig { - default_target: default_name, + default_target, classify_trigger, message_hash_fallback, }, @@ -731,24 +656,13 @@ impl LlmTaskClassifier { } fn build_escalation( - judge_target: ModelId, - efficient_target: ModelId, - capable_target: ModelId, contract_config: ClassifierContractConfig, config: EscalationJudgeConfig, max_output_tokens: u64, ) -> Result { - let inner = escalation::build_classifier( - judge_target, - &efficient_target, - &capable_target, - contract_config, - config, - max_output_tokens, - )?; - let targets = vec![capable_target, efficient_target]; + let inner = escalation::build_classifier(contract_config, config, max_output_tokens)?; Ok(Self { - route: FallThrough::::new_with_state(targets) + route: FallThrough::::new_with_state(vec![]) .with_name(ALGORITHM_NAME) .with_classifier(Arc::clone(&inner)), inner, @@ -762,11 +676,9 @@ impl LlmTaskClassifier { /// Keeps affinity and fallback ordering identical across judge-backed modes. fn from_classifier( - targets: Vec, inner: Arc>, config: ClassifierRouteConfig, ) -> Result { - algorithm::ensure_model_is_target(&targets, &config.default_target)?; if config.message_hash_fallback && config.classify_trigger != ClassifyTrigger::NewSession { return Err(LibsyError::AlgorithmError { message: "message_hash_fallback requires classify_trigger = new_session" @@ -774,7 +686,7 @@ impl LlmTaskClassifier { }); } // Affinity comes first so a retained assignment short-circuits the judge call. - let mut route = FallThrough::::new_with_state(targets).with_name(ALGORITHM_NAME); + let mut route = FallThrough::::new_with_state(vec![]).with_name(ALGORITHM_NAME); if let Some(affinity) = affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref() { @@ -783,7 +695,7 @@ impl LlmTaskClassifier { .with_processor(affinity.clone()) .with_classifier(affinity.clone()); } - let fallback = DefaultTarget::new(config.default_target); + let fallback = DefaultCategoryClassifier(config.default_target); Ok(Self { route: route .with_classifier(inner.clone()) @@ -793,18 +705,6 @@ impl LlmTaskClassifier { } } -#[async_trait] -impl Classifier for TaskClassifier { - async fn score( - &self, - state: &mut State, - request: &mut Request, - driver: Option<&Driver>, - ) -> Result<(Classification, Option)> { - self.classifier.score(state, request, driver).await - } -} - #[async_trait] impl Classifier for LlmTaskClassifier { async fn score( @@ -834,6 +734,7 @@ impl Algorithm for LlmTaskClassifier { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; use parking_lot::Mutex; @@ -841,12 +742,12 @@ mod tests { use super::*; use switchyard_protocol::{ - ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ToolCall, ToolResult, - completion_text, text_request, text_response, + ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ModelId, ToolCall, + ToolResult, completion_text, text_request, text_response, }; use crate::algorithms::util::llm_judge::Judge; - use crate::core::testing::{Serve, test_drive}; + use crate::core::testing::{Serve, test_drive_with_models}; use switchyard_protocol::{LlmResponse, Response}; const TEST_THRESHOLD: f64 = 0.5; @@ -859,7 +760,24 @@ mod tests { } fn policy() -> TaskClassifierPolicy { - TaskClassifierPolicy::new("efficient", "capable", &test_config(TEST_THRESHOLD)) + TaskClassifierPolicy::new(&test_config(TEST_THRESHOLD)) + } + + fn runtime_models() -> HashMap> { + [ + (Category::Judge, vec![ModelId::from("judge")]), + (Category::Efficient, vec![ModelId::from("efficient")]), + (Category::Capable, vec![ModelId::from("capable")]), + ( + Category::Any, + vec![ModelId::from("efficient"), ModelId::from("capable")], + ), + ] + .into() + } + + fn policy_driver() -> Driver { + Driver::new("test", runtime_models()).0 } fn verdict( @@ -880,7 +798,7 @@ mod tests { verdict: Option<&TaskClassifierVerdict>, ) -> Result { policy - .to_classification(verdict) + .to_classification(verdict, &policy_driver())? .argmax(false)? .map(|score| score.target) .ok_or_else(|| LibsyError::AlgorithmError { @@ -978,9 +896,6 @@ mod tests { fn router() -> Result> { Ok(Arc::new(LlmTaskClassifier::new( LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), config: test_config(TEST_THRESHOLD), }, )?)) @@ -1021,8 +936,13 @@ mod tests { async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> { let router = router()?; - let (selected_model, response) = - test_drive(router, classify_request(), unreachable_judge()).await?; + let (selected_model, response) = test_drive_with_models( + router, + classify_request(), + runtime_models(), + unreachable_judge(), + ) + .await?; assert_eq!(selected_model, "capable"); assert_eq!( @@ -1036,10 +956,17 @@ mod tests { async fn classifier_judges_each_request_without_affinity() -> Result<()> { let recorder = Arc::new(Recorder::default()); let router = router()?; - let request = classify_request; + let request = classify_request(); + let models = runtime_models(); - test_drive(router.clone(), request(), recorder.serve()).await?; - test_drive(router.clone(), request(), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + request.clone(), + models.clone(), + recorder.serve(), + ) + .await?; + test_drive_with_models(router, request, models, recorder.serve()).await?; assert_eq!( recorder.calls(), @@ -1061,16 +988,19 @@ mod tests { async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> { let recorder = Arc::new(Recorder::default()); let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), config: TaskClassifierConfig { max_output_tokens: 512, ..test_config(TEST_THRESHOLD) }, })?); - test_drive(router, classify_request(), recorder.serve()).await?; + test_drive_with_models( + router, + classify_request(), + runtime_models(), + recorder.serve(), + ) + .await?; assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]); Ok(()) @@ -1080,9 +1010,6 @@ mod tests { async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> { let recorder = Arc::new(Recorder::default()); let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), config: TaskClassifierConfig { contract: ClassifierContractConfig::default() .with_prompt("Custom capability rubric."), @@ -1090,7 +1017,13 @@ mod tests { }, })?); - test_drive(router, classify_request(), recorder.serve()).await?; + test_drive_with_models( + router, + classify_request(), + runtime_models(), + recorder.serve(), + ) + .await?; let prompts = recorder.judge_system_prompts(); assert_eq!(prompts.len(), 1); @@ -1102,18 +1035,22 @@ mod tests { async fn classifier_config_enables_new_session_trigger() -> Result<()> { let recorder = Arc::new(Recorder::default()); let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), config: TaskClassifierConfig { classify_trigger: ClassifyTrigger::NewSession, ..test_config(TEST_THRESHOLD) }, })?); - let session_request = classify_session_request; - test_drive(router.clone(), session_request(), recorder.serve()).await?; - test_drive(router.clone(), session_request(), recorder.serve()).await?; + let request = classify_session_request(); + let models = runtime_models(); + test_drive_with_models( + router.clone(), + request.clone(), + models.clone(), + recorder.serve(), + ) + .await?; + test_drive_with_models(router, request, models, recorder.serve()).await?; assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]); Ok(()) @@ -1123,9 +1060,6 @@ mod tests { async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> { let recorder = Arc::new(Recorder::default()); let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("efficient"), - capable_target: ModelId::from("capable"), config: TaskClassifierConfig { classify_trigger: ClassifyTrigger::NewSession, message_hash_fallback: true, @@ -1134,10 +1068,18 @@ mod tests { }, })?); - test_drive(router.clone(), classify_request(), recorder.serve()).await?; - test_drive( + let models = runtime_models(); + test_drive_with_models( router.clone(), + classify_request(), + models.clone(), + recorder.serve(), + ) + .await?; + test_drive_with_models( + router, classify_follow_up_request(), + models, recorder.serve(), ) .await?; @@ -1146,6 +1088,77 @@ mod tests { Ok(()) } + #[tokio::test] + async fn one_classifier_uses_each_requests_runtime_models() -> Result<()> { + let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { + config: TaskClassifierConfig { + classify_trigger: ClassifyTrigger::NewSession, + ..test_config(TEST_THRESHOLD) + }, + })?); + let calls = Arc::new(Mutex::new(Vec::new())); + let serve = |calls: Arc>>| { + move |model: ModelId, _request: Request| { + let calls = Arc::clone(&calls); + async move { + calls.lock().push(model.to_string()); + let text = if model.as_str().starts_with("judge-") { + r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string() + } else { + model.to_string() + }; + Ok(Response { + llm_response: LlmResponse::Agg(text_response(None, text)), + metadata: None, + }) + } + } + }; + let models = |suffix: &str| { + [ + ( + Category::Judge, + vec![ModelId::from(format!("judge-{suffix}"))], + ), + ( + Category::Efficient, + vec![ModelId::from(format!("efficient-{suffix}"))], + ), + ( + Category::Capable, + vec![ModelId::from(format!("capable-{suffix}"))], + ), + ( + Category::Any, + vec![ + ModelId::from(format!("efficient-{suffix}")), + ModelId::from(format!("capable-{suffix}")), + ], + ), + ] + .into() + }; + let request = classify_session_request(); + + let (first, _) = test_drive_with_models( + router.clone(), + request.clone(), + models("a"), + serve(Arc::clone(&calls)), + ) + .await?; + let (second, _) = + test_drive_with_models(router, request, models("b"), serve(Arc::clone(&calls))).await?; + + assert_eq!(first, "efficient-a"); + assert_eq!(second, "efficient-b"); + assert_eq!( + &*calls.lock(), + &["judge-a", "efficient-a", "judge-b", "efficient-b"] + ); + Ok(()) + } + #[test] fn the_threshold_boundary_is_inclusive() -> Result<()> { let policy = policy(); @@ -1159,8 +1172,8 @@ mod tests { #[test] fn the_threshold_moves_the_routing_boundary() -> Result<()> { let borderline = verdict(0.5, "supported", "SUP-1"); - let strict = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.9)); - let lenient = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.1)); + let strict = TaskClassifierPolicy::new(&test_config(0.9)); + let lenient = TaskClassifierPolicy::new(&test_config(0.1)); assert_eq!(selected(&strict, Some(&borderline))?, "capable"); assert_eq!(selected(&lenient, Some(&borderline))?, "efficient"); Ok(()) @@ -1187,9 +1200,6 @@ mod tests { for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] { assert!( LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("e"), - capable_target: ModelId::from("c"), config: test_config(bad), }) .is_err(), @@ -1218,21 +1228,10 @@ mod tests { ..TaskClassifierConfig::default() }, ] { - assert!( - LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("e"), - capable_target: ModelId::from("c"), - config, - }) - .is_err() - ); + assert!(LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).is_err()); } for base_threshold in [0.0, 1.0] { LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: ModelId::from("judge"), - efficient_target: ModelId::from("e"), - capable_target: ModelId::from("c"), config: test_config(base_threshold), })?; } @@ -1257,7 +1256,7 @@ mod tests { None, ]; for verdict in unusable { - let classification = policy.to_classification(verdict.as_ref()); + let classification = policy.to_classification(verdict.as_ref(), &policy_driver())?; assert!(matches!(classification, Classification::Ambiguous(_))); assert!(classification.argmax(false)?.is_none()); assert!(classification.argmax(true)?.is_none()); @@ -1267,14 +1266,10 @@ mod tests { #[test] fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> { - let policy = TaskClassifierPolicy::new( - "efficient", - "capable", - &TaskClassifierConfig { - threshold_step: 0.1, - ..test_config(0.4) - }, - ); + let policy = TaskClassifierPolicy::new(&TaskClassifierConfig { + threshold_step: 0.1, + ..test_config(0.4) + }); assert_eq!( selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?, diff --git a/crates/libsy/src/algorithms/passthrough.rs b/crates/libsy/src/algorithms/passthrough.rs index bd2fa5994..bd4d71daa 100644 --- a/crates/libsy/src/algorithms/passthrough.rs +++ b/crates/libsy/src/algorithms/passthrough.rs @@ -5,24 +5,14 @@ use std::sync::Arc; -use switchyard_protocol::{ModelId, Request}; +use switchyard_protocol::{Category, Request}; use crate::core::algorithm::{Algorithm, Driver}; -use crate::{Result, RoutingOutcome}; +use crate::{LibsyError, Result, RoutingOutcome}; /// Routing algorithm that always selects one configured target. -pub struct Passthrough { - target: ModelId, -} - -impl Passthrough { - /// Creates an algorithm that always selects `target`. - pub fn new(target: impl Into) -> Self { - Self { - target: target.into(), - } - } -} +#[derive(Default)] +pub struct Passthrough {} #[async_trait::async_trait] impl Algorithm for Passthrough { @@ -30,10 +20,13 @@ impl Algorithm for Passthrough { "passthrough" } - async fn route(self: Arc, _driver: Driver, request: Request) -> Result { - tracing::info!(target = %self.target, "passthrough selected target"); + async fn route(self: Arc, driver: Driver, request: Request) -> Result { + let Some(target) = driver.models_for(Category::Any).first() else { + return Err(LibsyError::NoTargets); + }; + tracing::info!(target = %target, "passthrough selected target"); Ok(RoutingOutcome::route_to( - self.target.clone(), + target.clone(), Vec::new(), request, )) @@ -45,9 +38,9 @@ mod tests { use std::sync::Arc; use super::Passthrough; - use crate::core::algorithm::Algorithm; - use crate::core::testing::{echo, test_drive}; - use switchyard_protocol::{Request, completion_text, text_request}; + use crate::core::testing::echo; + use crate::core::{algorithm::Algorithm, testing::test_drive_with_models}; + use switchyard_protocol::{Category, Request, completion_text, text_request}; #[tokio::test] async fn test_passthrough() -> crate::Result<()> { @@ -57,8 +50,10 @@ mod tests { raw_request: None, metadata: None, }; - let algorithm: Arc = Arc::new(Passthrough::new(MODEL_ID)); - let (selected_model, response) = test_drive(algorithm, request, echo()).await?; + let algorithm: Arc = Arc::new(Passthrough::default()); + let models = Category::to_map(Category::Any, &[MODEL_ID]); + let (selected_model, response) = + test_drive_with_models(algorithm, request, models, echo()).await?; assert_eq!( response diff --git a/crates/libsy/src/algorithms/rand.rs b/crates/libsy/src/algorithms/rand.rs index 87f89ca61..8f0389dc0 100644 --- a/crates/libsy/src/algorithms/rand.rs +++ b/crates/libsy/src/algorithms/rand.rs @@ -6,11 +6,11 @@ //! [`RandomClassifier`] selects one target; [`FallThrough`] owns the common //! processor/classifier/target-call orchestration. -use std::collections::BTreeSet; use std::sync::Arc; use async_trait::async_trait; use parking_lot::Mutex; +use rand::RngExt as _; use rand::SeedableRng; use rand::distr::{Distribution, weighted::WeightedIndex}; use rand::rngs::StdRng; @@ -19,12 +19,11 @@ use crate::algorithms::fall_through::FallThrough; use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::{Classification, Classifier, Score}; use crate::{LibsyError, Result}; -use switchyard_protocol::{ModelId, Request, Response}; +use switchyard_protocol::{Category, Request, Response}; /// Stateless weighted classifier used by random fall-through routing. pub struct RandomClassifier { - targets: Vec, - distribution: WeightedIndex, + distribution: Option>, rng: Mutex, } @@ -37,63 +36,36 @@ impl RandomClassifier { /// /// # Errors /// - /// Returns an error when targets are empty or duplicated, or when explicit - /// weights have the wrong length, are negative or non-finite, or contain no + /// Returns an error explicit weights are negative or non-finite, or contain no /// positive value. - pub fn new( - targets: Vec, - weights: Option>, - seed: Option, - ) -> Result { - let target_count = targets.len(); - if target_count == 0 { - return Err(LibsyError::NoTargets); - } - let unique_targets = targets.iter().map(ModelId::as_str).collect::>(); - if unique_targets.len() != target_count { - return Err(LibsyError::AlgorithmError { - message: "random targets must be unique".to_string(), - }); - } - - let weights = weights.unwrap_or_else(|| vec![1.0; target_count]); - if weights.len() != target_count { - return Err(invalid_weights(format!( - "expected {target_count} weights, got {}", - weights.len() - ))); - } - if weights - .iter() - .any(|weight| !weight.is_finite() || *weight < 0.0) - { - return Err(invalid_weights( - "weights must be finite and nonnegative".to_string(), - )); - } - if !weights.iter().any(|weight| *weight > 0.0) { - return Err(invalid_weights( - "at least one weight must be positive".to_string(), - )); - } - let distribution = - WeightedIndex::new(weights).map_err(|error| invalid_weights(error.to_string()))?; + pub fn new(weights: Option>, seed: Option) -> Result { + let distribution = if let Some(weights) = weights { + if weights + .iter() + .any(|weight| !weight.is_finite() || *weight < 0.0) + { + return Err(invalid_weights( + "weights must be finite and nonnegative".to_string(), + )); + } + if !weights.iter().any(|weight| *weight > 0.0) { + return Err(invalid_weights( + "at least one weight must be positive".to_string(), + )); + } + Some(WeightedIndex::new(weights).map_err(|error| invalid_weights(error.to_string()))?) + } else { + None + }; let rng = match seed { Some(seed) => StdRng::seed_from_u64(seed), None => rand::make_rng(), }; Ok(Self { - targets, distribution, rng: Mutex::new(rng), }) } - - fn select_target(&self) -> ModelId { - let mut rng = self.rng.lock(); - let index = self.distribution.sample(&mut *rng); - self.targets[index].clone() - } } fn invalid_weights(message: String) -> LibsyError { @@ -111,12 +83,30 @@ where &self, _state: &mut S, _request: &mut Request, - _driver: Option<&Driver>, + driver: Option<&Driver>, ) -> Result<(Classification, Option)> { + let Some(driver) = driver else { + // Temp until can remove Option from driver + return Err(LibsyError::NoTargets); + }; + // All the available models + let options = driver.models_for(Category::Any); + if options.is_empty() { + return Err(LibsyError::NoTargets); + } + let mut rng = self.rng.lock(); + let index = if let Some(distribution) = self.distribution.as_ref() { + // The user gave us weights + distribution.sample(&mut *rng) + } else { + // No weights, assume equal probability + rng.random_range(..options.len()) + }; + let target = options[index].clone(); Ok(( Classification::Scores(vec![Score { confidence: 1.0, - target: self.select_target(), + target, }]), None, )) @@ -129,18 +119,10 @@ pub struct Random { } impl Random { - /// Creates a router over `targets`. - /// - /// # Errors - /// - /// Returns an error when targets or weights are invalid for [`RandomClassifier`]. - pub fn new( - targets: Vec, - weights: Option>, - seed: Option, - ) -> Result { - let classifier = Arc::new(RandomClassifier::new(targets.clone(), weights, seed)?); - let inner = FallThrough::<()>::new(targets) + /// Creates a random router. The models themselves will be passed at runtime. + pub fn new(weights: Option>, seed: Option) -> Result { + let classifier = Arc::new(RandomClassifier::new(weights, seed)?); + let inner = FallThrough::<()>::new(vec![]) .with_name("random") .with_classifier(classifier); Ok(Self { inner }) @@ -165,12 +147,12 @@ impl Algorithm for Random { #[cfg(test)] mod tests { use super::*; - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; - use switchyard_protocol::{Metadata, completion_text, text_request}; + use switchyard_protocol::{Metadata, ModelId, completion_text, text_request}; use crate::algorithms::util::affinity::AffinityRouter; - use crate::core::testing::{echo, test_drive}; + use crate::core::testing::{echo, test_drive_with_models}; use switchyard_protocol::Request; fn request() -> Request { @@ -191,22 +173,30 @@ mod tests { } } + /* fn target_set(names: &[&str]) -> Vec { names.iter().map(|name| ModelId::from(*name)).collect() } + */ - fn algorithm(names: &[&str], weights: Option>, seed: Option) -> Result { - Random::new(target_set(names), weights, seed) + fn algorithm(weights: Option>, seed: Option) -> Result { + Random::new(weights, seed) } - fn shared_algorithm(names: &[&str]) -> Result> { - Ok(Arc::new(algorithm(names, None, None)?)) + fn shared_algorithm() -> Result> { + Ok(Arc::new(algorithm(None, None)?)) } - async fn selected_models(algorithm: Arc, count: usize) -> Result> { + async fn selected_models( + algorithm: Arc, + count: usize, + models: HashMap>, + ) -> Result> { let mut selected = Vec::with_capacity(count); for _ in 0..count { - let (_, response) = test_drive(algorithm.clone(), request(), echo()).await?; + let (_, response) = + test_drive_with_models(algorithm.clone(), request(), models.clone(), echo()) + .await?; selected.push( response .llm_response @@ -218,10 +208,16 @@ mod tests { Ok(selected) } + fn to_category_map(names: &[&str]) -> HashMap> { + Category::to_map(Category::Any, names) + } + #[tokio::test] async fn single_target_is_always_selected_and_called() -> Result<()> { - let algorithm = shared_algorithm(&["only/model"])?; - let (selected_model, response) = test_drive(algorithm, request(), echo()).await?; + let algorithm = shared_algorithm()?; + let models = to_category_map(&["only/model"]); + let (selected_model, response) = + test_drive_with_models(algorithm, request(), models, echo()).await?; assert_eq!( response @@ -237,12 +233,14 @@ mod tests { #[tokio::test] async fn selection_covers_all_targets_over_many_runs() -> Result<()> { - let algorithm = shared_algorithm(&["a/model", "b/model"])?; + let algorithm = shared_algorithm()?; + let models = to_category_map(&["a/model", "b/model"]); let mut seen = HashSet::new(); for _ in 0..100 { let (selected_model, response) = - test_drive(algorithm.clone(), request(), echo()).await?; + test_drive_with_models(algorithm.clone(), request(), models.clone(), echo()) + .await?; let served_model = response .llm_response .as_agg() @@ -263,19 +261,12 @@ mod tests { #[tokio::test] async fn weighted_seeded_selection_is_reproducible() -> Result<()> { - let first: Arc = Arc::new(algorithm( - &["a/model", "b/model"], - Some(vec![1.0, 3.0]), - Some(42), - )?); - let second: Arc = Arc::new(algorithm( - &["a/model", "b/model"], - Some(vec![1.0, 3.0]), - Some(42), - )?); - - let first_selections = selected_models(first, 1_000).await?; - let second_selections = selected_models(second, 1_000).await?; + let models = to_category_map(&["a/model", "b/model"]); + let first: Arc = Arc::new(algorithm(Some(vec![1.0, 3.0]), Some(42))?); + let second: Arc = Arc::new(algorithm(Some(vec![1.0, 3.0]), Some(42))?); + + let first_selections = selected_models(first, 1_000, models.clone()).await?; + let second_selections = selected_models(second, 1_000, models).await?; assert_eq!(first_selections, second_selections); let second_count = first_selections @@ -292,22 +283,24 @@ mod tests { #[tokio::test] async fn affinity_reuses_the_initial_random_selection() -> Result<()> { let names = ["a/model", "b/model"]; + let models = to_category_map(&names); let affinity = Arc::new(AffinityRouter::new()); - let random = Arc::new(RandomClassifier::new( - names.iter().map(|name| ModelId::from(*name)).collect(), - None, - Some(42), - )?); + let random = Arc::new(RandomClassifier::new(None, Some(42))?); let algorithm: Arc = Arc::new( - FallThrough::<()>::new(target_set(&names)) + FallThrough::<()>::new(vec![]) .with_name("affinity_random") .with_processor(affinity.clone()) .with_classifier(affinity.clone()) .with_classifier(random), ); - let (_, first) = - test_drive(algorithm.clone(), request_for_session("session-1"), echo()).await?; + let (_, first) = test_drive_with_models( + algorithm.clone(), + request_for_session("session-1"), + models.clone(), + echo(), + ) + .await?; let selected = first .llm_response .as_agg() @@ -326,7 +319,9 @@ mod tests { Some(ModelId::from(selected.clone())) ); - let (_, second) = test_drive(algorithm, request_for_session("session-1"), echo()).await?; + let (_, second) = + test_drive_with_models(algorithm, request_for_session("session-1"), models, echo()) + .await?; assert_eq!( second .llm_response @@ -341,14 +336,13 @@ mod tests { #[test] fn rejects_invalid_weights() { let cases = [ - (vec![1.0], "expected 2 weights"), (vec![1.0, -1.0], "finite and nonnegative"), (vec![0.0, 0.0], "at least one weight must be positive"), (vec![1.0, f64::INFINITY], "finite and nonnegative"), ]; for (weights, expected) in cases { - let error = algorithm(&["a/model", "b/model"], Some(weights), None) + let error = algorithm(Some(weights), None) .err() .map(|error| error.to_string()) .unwrap_or_default(); @@ -356,22 +350,12 @@ mod tests { } } - #[test] - fn rejects_invalid_targets() { - let error = algorithm(&[], None, None).err(); - assert!(matches!(error, Some(LibsyError::NoTargets))); - - let error = algorithm(&["same/model", "same/model"], None, None) - .err() - .map(|error| error.to_string()) - .unwrap_or_default(); - assert!(error.contains("random targets must be unique")); - } - #[tokio::test] async fn decision_is_inspectable() -> Result<()> { - let algorithm = shared_algorithm(&["only/model"])?; - let (selected_model, _) = test_drive(algorithm, request(), echo()).await?; + let algorithm = shared_algorithm()?; + let models = to_category_map(&["only/model"]); + let (selected_model, _) = + test_drive_with_models(algorithm, request(), models, echo()).await?; assert_eq!(selected_model, "only/model"); Ok(()) } diff --git a/crates/libsy/src/algorithms/stage.rs b/crates/libsy/src/algorithms/stage.rs index 3fd4acc1e..11cac75a9 100644 --- a/crates/libsy/src/algorithms/stage.rs +++ b/crates/libsy/src/algorithms/stage.rs @@ -19,20 +19,79 @@ use async_trait::async_trait; use super::fall_through::FallThrough; use super::llm_class::{LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig}; -use super::util::prompts::{SystemPromptProcessor, TargetPrompts}; +use super::util::prompts::prepend_system_prompt; use super::util::stage::{ - DecisionSource, HandoffNoteConfig, PickerMode, StageClassifier, StageTargets, Tier, - fall_open_tier, record_decision_source, record_routing_decision, + DecisionSource, HandoffNoteConfig, PickerMode, StageClassifier, Tier, fall_open_tier, + record_decision_source, record_routing_decision, }; use super::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignalProcessor}; use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::{Classification, Classifier, Score}; -use crate::core::state::State; +use crate::core::processor::{Event, Processor}; +use crate::core::state::{State, StateValue}; use crate::{LibsyError, Result}; -use switchyard_protocol::{ModelId, Request, Response}; +use switchyard_protocol::{Category, ModelId, Request, Response}; /// Telemetry name for a router this module assembles. const STAGE_ROUTER: &str = "stage_router"; +const CAPABLE_MODEL_KEY: &str = "stage.capable_model"; +const EFFICIENT_MODEL_KEY: &str = "stage.efficient_model"; + +struct TierPromptProcessor { + capable: Option, + efficient: Option, +} + +impl TierPromptProcessor { + fn remember_model(state: &mut State, key: &str, driver: &Driver, category: Category) { + if let Some(model) = driver.models_for(category).first() { + state + .extra + .insert(key.to_string(), StateValue::String(model.to_string())); + } else { + state.extra.remove(key); + } + } + + fn is_model(state: &State, key: &str, selected: &ModelId) -> bool { + matches!( + state.extra.get(key), + Some(StateValue::String(model)) if model == selected.as_str() + ) + } +} + +#[async_trait] +impl Processor for TierPromptProcessor { + async fn process(&self, state: &mut State, event: Event<'_>) -> Result<()> { + match event { + Event::Request { + driver: Some(driver), + .. + } => { + Self::remember_model(state, CAPABLE_MODEL_KEY, driver, Category::Capable); + Self::remember_model(state, EFFICIENT_MODEL_KEY, driver, Category::Efficient); + } + Event::Decision { + request, + selected_model_id, + } => { + let prompt = if Self::is_model(state, CAPABLE_MODEL_KEY, selected_model_id) { + self.capable.as_deref() + } else if Self::is_model(state, EFFICIENT_MODEL_KEY, selected_model_id) { + self.efficient.as_deref() + } else { + None + }; + if let Some(prompt) = prompt { + prepend_system_prompt(request, prompt); + } + } + _ => {} + } + Ok(()) + } +} /// Attributes a turn to the classifier it wraps, when that classifier decides it. /// @@ -68,7 +127,6 @@ impl Classifier for SourceStamp { /// Closes the cascade at zero confidence: a fallback, not a judgement. struct FallOpen { - targets: StageTargets, default_tier: Tier, } @@ -78,10 +136,17 @@ impl Classifier for FallOpen { &self, state: &mut State, _request: &mut Request, - _driver: Option<&Driver>, + driver: Option<&Driver>, ) -> Result<(Classification, Option)> { let tier = fall_open_tier(state).unwrap_or(self.default_tier); - let target = self.targets.name(tier).clone(); + let category = match tier { + Tier::Capable => Category::Capable, + Tier::Efficient => Category::Efficient, + }; + let target = driver + .ok_or(LibsyError::NoTargets)? + .first_model_for(category)? + .clone(); Ok(( Classification::Scores(vec![Score { target, @@ -94,9 +159,6 @@ impl Classifier for FallOpen { /// The capability judge a stage router falls through to. pub struct LlmFallback { - /// Target the judge model is called through. It is not a routing - /// destination, so it does not belong in the router's target set. - pub judge_target: ModelId, /// Judge configuration. `recent_turn_window` is worth setting to this router's /// `recent_window` so the judge reads the same span the signal scorer scored. /// Note: `classify_trigger = new_session` and `message_hash_fallback` have no effect here — @@ -116,12 +178,12 @@ pub struct StageRouterConfig { /// Note handed to the model on a signal-driven escalation, and on a /// hand-back to the efficient tier when a de-escalation note is configured. pub handoff_notes: Option, - /// System prompts keyed by target, handed over on every turn that target - /// serves. Empty by default. - pub tier_prompts: TargetPrompts, - /// Capability judge consulted on turns the signals leave undecided — the - /// judge's own target, plus the same configuration the standalone capability - /// route takes. + /// System prompt handed to the runtime capable model. + pub capable_system_prompt: Option, + /// System prompt handed to the runtime efficient model. + pub efficient_system_prompt: Option, + /// Capability judge consulted on turns the signals leave undecided. It uses + /// the runtime judge model and the standalone capability route's settings. pub llm_fallback: Option, } @@ -134,7 +196,8 @@ impl StageRouterConfig { confidence_threshold, recent_window: None, handoff_notes: None, - tier_prompts: TargetPrompts::default(), + capable_system_prompt: None, + efficient_system_prompt: None, llm_fallback: None, } } @@ -148,14 +211,13 @@ pub struct StageRouter { } impl StageRouter { - /// Routes between the `capable` and `efficient` targets. The - /// judge, when configured, is called through its own target and is not a - /// routing destination. + /// Routes between the runtime `capable` and `efficient` models. The judge, + /// when configured, is called through the runtime `judge` model. /// /// 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(config: StageRouterConfig) -> Result { Ok(Self { - route: build_stage_route(capable, efficient, config)?, + route: build_stage_route(config)?, }) } } @@ -177,11 +239,7 @@ impl Algorithm for StageRouter { /// Wires the cascade the wrapper drives. Exposed so a composition above can /// stack a prelude onto it. -pub(crate) fn build_stage_route( - capable: ModelId, - efficient: ModelId, - config: StageRouterConfig, -) -> Result> { +pub(crate) fn build_stage_route(config: StageRouterConfig) -> Result> { if !(0.0..=1.0).contains(&config.confidence_threshold) { return Err(LibsyError::AlgorithmError { message: format!( @@ -190,35 +248,23 @@ pub(crate) fn build_stage_route( ), }); } - // The tiers are a fixed pair; their targets are whatever the deployment calls - // them, and the classifier scores onto those names. - let targets = StageTargets::new(capable.clone(), efficient.clone()); let default_tier = config.mode.default_tier(); - let fall_open = FallOpen { - targets: targets.clone(), - default_tier, - }; + let fall_open = FallOpen { default_tier }; - let mut classifier = StageClassifier::new(targets, config.mode, config.confidence_threshold); + let mut classifier = StageClassifier::new(config.mode, config.confidence_threshold); if let Some(notes) = config.handoff_notes { classifier = classifier.with_handoff_notes(notes); } let signals = ToolSignalProcessor { recent_window: config.recent_window.unwrap_or(DEFAULT_RECENT_WINDOW), }; - let target_set = vec![capable.clone(), efficient.clone()]; - let mut router = FallThrough::::new_with_state(target_set) + let mut router = FallThrough::::new_with_state(vec![]) .with_name(STAGE_ROUTER) .with_processor(Arc::new(signals)) .with_classifier(Arc::new(classifier)); if let Some(fallback) = config.llm_fallback { - // The capability judge takes its tiers in the same order the capability - // route passes them: efficient first, capable second. router = router.with_classifier(Arc::new(SourceStamp { inner: Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: fallback.judge_target, - efficient_target: efficient, - capable_target: capable, config: fallback.config, })?), source: DecisionSource::LlmClassifier, @@ -229,15 +275,18 @@ pub(crate) fn build_stage_route( inner: Arc::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))); + if config.capable_system_prompt.is_some() || config.efficient_system_prompt.is_some() { + router = router.with_processor(Arc::new(TierPromptProcessor { + capable: config.capable_system_prompt, + efficient: config.efficient_system_prompt, + })); + } Ok(router) } #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; @@ -246,10 +295,8 @@ mod tests { use super::*; use crate::algorithms::util::stage::{DECISION_SOURCE_KEY, clear_fall_open, set_fall_open}; use crate::algorithms::util::tier_fixtures::{JUDGE, Recorder, turn_request}; - use crate::core::processor::{Event, Processor}; - use crate::core::state::StateValue; - use crate::core::testing::test_drive; - use switchyard_protocol::Response; + use crate::core::testing::test_drive_with_models; + use switchyard_protocol::{Category, ModelId}; /// A classifier that always picks `target`, standing in for a cascade member. struct Fixed(&'static str); @@ -322,12 +369,29 @@ mod tests { StageRouterConfig::new(PickerMode::EfficientFirst, 0.5) } + fn runtime_models() -> HashMap> { + runtime_models_for("strong", "weak") + } + + fn runtime_models_for(capable: &str, efficient: &str) -> HashMap> { + [ + (Category::Judge, vec![ModelId::from(JUDGE)]), + (Category::Efficient, vec![ModelId::from(efficient)]), + (Category::Capable, vec![ModelId::from(capable)]), + ( + Category::Any, + vec![ModelId::from(capable), ModelId::from(efficient)], + ), + ] + .into() + } + #[test] fn rejects_an_out_of_range_confidence_threshold() { let mut config = config(); config.confidence_threshold = 1.5; assert!(matches!( - StageRouter::new(ModelId::from("strong"), ModelId::from("weak"), config), + StageRouter::new(config), Err(LibsyError::AlgorithmError { .. }) )); } @@ -336,21 +400,20 @@ mod tests { fn rejects_an_out_of_range_judge_threshold() { let mut config = config(); config.llm_fallback = Some(LlmFallback { - judge_target: ModelId::from("judge"), config: TaskClassifierConfig { base_threshold: -0.1, ..Default::default() }, }); assert!(matches!( - StageRouter::new(ModelId::from("strong"), ModelId::from("weak"), config), + StageRouter::new(config), Err(LibsyError::AlgorithmError { .. }) )); } #[test] - fn builds_over_both_tiers() -> Result<()> { - let router = StageRouter::new(ModelId::from("strong"), ModelId::from("weak"), config())?; + fn builds() -> Result<()> { + let router = StageRouter::new(config())?; assert_eq!(router.name(), STAGE_ROUTER); Ok(()) } @@ -359,11 +422,7 @@ mod tests { const ESCALATION: &str = "the previous model was stalling; pick up the diagnosis"; fn recording_router(config: StageRouterConfig) -> Result> { - Ok(Arc::new(StageRouter::new( - ModelId::from("strong"), - ModelId::from("weak"), - config, - )?)) + Ok(Arc::new(StageRouter::new(config)?)) } fn config_with_notes() -> StageRouterConfig { @@ -376,7 +435,6 @@ mod tests { *recorder.judge_p_solve.lock() = p_solve; let mut c = config(); c.llm_fallback = Some(LlmFallback { - judge_target: ModelId::from(JUDGE), config: TaskClassifierConfig { base_threshold: 0.5, recent_turn_window: Some(3), @@ -414,15 +472,19 @@ mod tests { let recorder = Arc::new(Recorder::default()); // The picker would fall open to "strong"; the override says "weak". let config = StageRouterConfig::new(PickerMode::CapableFirst, 0.5); - let route: Arc = Arc::new( - build_stage_route(ModelId::from("strong"), ModelId::from("weak"), config)? - .with_processor(Arc::new(TierDecider::default())), - ); - - test_drive(route.clone(), turn_request(false), recorder.serve()).await?; - test_drive(route.clone(), turn_request(false), recorder.serve()).await?; - test_drive(route.clone(), turn_request(true), recorder.serve()).await?; - test_drive(route.clone(), turn_request(false), recorder.serve()).await?; + let route: Arc = + Arc::new(build_stage_route(config)?.with_processor(Arc::new(TierDecider::default()))); + + let models = runtime_models(); + for is_turn_failed in [false, false, true, false] { + test_drive_with_models( + route.clone(), + turn_request(is_turn_failed), + models.clone(), + recorder.serve(), + ) + .await?; + } let routed = recorder.routed(); assert_eq!( @@ -449,12 +511,24 @@ mod tests { let recorder = Arc::new(Recorder::default()); let router = recording_router(config_with_notes())?; - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; - test_drive(router.clone(), turn_request(true), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; + test_drive_with_models( + router.clone(), + turn_request(true), + runtime_models_for("runtime-strong", "runtime-weak"), + recorder.serve(), + ) + .await?; let calls = recorder.routed(); assert_eq!(calls[0].target, "weak"); - assert_eq!(calls[1].target, "strong"); + assert_eq!(calls[1].target, "runtime-strong"); assert!( !calls[0].messages.iter().any(|t| t.contains(ESCALATION)), "steady-state turn should carry no note: {:?}", @@ -476,8 +550,13 @@ mod tests { let recorder = Arc::new(Recorder::default()); let router = recording_router(config_with_judge(&recorder, 0.1))?; - let (selected_model, _) = - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + let (selected_model, _) = test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; let calls = recorder.calls.lock(); assert!( @@ -498,7 +577,13 @@ mod tests { let recorder = Arc::new(Recorder::default()); let router = recording_router(config_with_judge(&recorder, 0.9))?; - test_drive(router.clone(), turn_request(true), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + turn_request(true), + runtime_models(), + recorder.serve(), + ) + .await?; assert!( !recorder.calls.lock().iter().any(|c| c.target == JUDGE), @@ -513,9 +598,21 @@ mod tests { let recorder = Arc::new(Recorder::default()); let router = recording_router(config_with_judge(&recorder, 0.1))?; - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; *recorder.judge_p_solve.lock() = 0.9; - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; let routed = recorder.routed(); assert_eq!(routed[0].target, "strong"); @@ -538,7 +635,13 @@ mod tests { let recorder = Arc::new(Recorder::default()); let router = recording_router(config_with_judge(&recorder, 42.0))?; - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; assert_eq!(recorder.routed()[0].target, "weak"); Ok(()) @@ -549,7 +652,13 @@ mod tests { let recorder = Arc::new(Recorder::default()); let router = recording_router(config_with_judge(&recorder, 0.9))?; - test_drive(router.clone(), turn_request(false), recorder.serve()).await?; + test_drive_with_models( + router.clone(), + turn_request(false), + runtime_models(), + recorder.serve(), + ) + .await?; let judged = recorder .calls diff --git a/crates/libsy/src/algorithms/subagent.rs b/crates/libsy/src/algorithms/subagent.rs index bb918f05e..82db78b90 100644 --- a/crates/libsy/src/algorithms/subagent.rs +++ b/crates/libsy/src/algorithms/subagent.rs @@ -118,15 +118,15 @@ mod tests { use parking_lot::Mutex; use serde_json::json; use switchyard_protocol::{ - ContentBlock, InstructionBlock, Message, Metadata, ModelId, Request, Response, Role, - text_request, + Category, ContentBlock, InstructionBlock, Message, Metadata, ModelId, Request, Response, + Role, text_request, }; use super::{SubagentRouter, SubagentRouterConfig}; use crate::algorithms::passthrough::Passthrough; use crate::core::algorithm::Algorithm; use crate::core::classifier::{Classification, Classifier, Score}; - use crate::core::testing::{echo, reply, test_drive}; + use crate::core::testing::{echo, reply, test_drive_with_models}; use crate::{ ClassifyTrigger, CustomClassifierConfig, CustomClassifierPolicy, Driver, LlmClassifierConfig, LlmTaskClassifier, State, @@ -178,7 +178,7 @@ mod tests { } fn parent() -> Arc { - Arc::new(Passthrough::new("parent")) + Arc::new(Passthrough::default()) } fn configured(classifier: Arc>) -> crate::Result> { @@ -201,11 +201,21 @@ mod tests { }); let router = configured(classifier.clone())?; - let (parent, _) = test_drive(router.clone(), request(None), echo()).await?; - let (first, _) = test_drive(router.clone(), child("child-1"), echo()).await?; - let (same_child, _) = test_drive(router.clone(), child("child-1"), echo()).await?; - let (sibling, _) = test_drive(router.clone(), child("child-2"), echo()).await?; - let (defaulted, _) = test_drive(router.clone(), child("child-3"), echo()).await?; + let models = Category::to_map(Category::Any, &["parent", "worker", "reviewer"]); + let (parent, _) = + test_drive_with_models(router.clone(), request(None), models.clone(), echo()).await?; + let (first, _) = + test_drive_with_models(router.clone(), child("child-1"), models.clone(), echo()) + .await?; + let (same_child, _) = + test_drive_with_models(router.clone(), child("child-1"), models.clone(), echo()) + .await?; + let (sibling, _) = + test_drive_with_models(router.clone(), child("child-2"), models.clone(), echo()) + .await?; + let (defaulted, _) = + test_drive_with_models(router.clone(), child("child-3"), models.clone(), echo()) + .await?; let maintenance = request(Some(Metadata { session_id: Some("session-1".to_string()), agent_id: Some("child-1".to_string()), @@ -213,7 +223,8 @@ mod tests { is_delegated_work: false, ..Metadata::default() })); - let (maintenance, _) = test_drive(router, maintenance, echo()).await?; + let (maintenance, _) = + test_drive_with_models(router, maintenance, models.clone(), echo()).await?; assert_eq!(parent, "parent"); assert_eq!(first, "worker"); @@ -228,18 +239,13 @@ mod tests { #[tokio::test] async fn custom_classifier_receives_only_the_delegated_prompt() -> crate::Result<()> { let classifier = LlmTaskClassifier::new(LlmClassifierConfig::Custom { - judge_target: ModelId::from("judge"), - targets: vec![ - ("worker".to_string(), ModelId::from("worker")), - ("reviewer".to_string(), ModelId::from("reviewer")), - ], - default_target: "worker".to_string(), + default_target: Category::Capable, config: CustomClassifierConfig::new( "classify the delegated task", json!({ "type": "object", "properties": { - "target": {"type": "string", "enum": ["worker", "reviewer"]} + "target": {"type": "string", "enum": ["capable", "efficient"]} }, "required": ["target"], "additionalProperties": false @@ -270,19 +276,34 @@ mod tests { let calls = Arc::new(Mutex::new(Vec::new())); let served_calls = calls.clone(); - let (selected, _) = test_drive(router, request, move |target, request| { - let calls = served_calls.clone(); - async move { - let completion = if target == "judge" { - r#"{"target":"reviewer"}"# - } else { - "child answer" - }; - calls.lock().push((target, request)); - Ok(reply(completion)) - } - }) - .await?; + let models = [ + (Category::Judge, vec![ModelId::from("judge")]), + (Category::Capable, vec![ModelId::from("worker")]), + (Category::Efficient, vec![ModelId::from("reviewer")]), + ( + Category::Any, + vec![ + ModelId::from("parent"), + ModelId::from("worker"), + ModelId::from("reviewer"), + ], + ), + ] + .into(); + let (selected, _) = + test_drive_with_models(router, request, models, move |target, request| { + let calls = served_calls.clone(); + async move { + let completion = if target == "judge" { + r#"{"target":"efficient"}"# + } else { + "child answer" + }; + calls.lock().push((target, request)); + Ok(reply(completion)) + } + }) + .await?; assert_eq!(selected, "reviewer"); let calls = calls.lock(); diff --git a/crates/libsy/src/algorithms/subagent_affinity_tests.rs b/crates/libsy/src/algorithms/subagent_affinity_tests.rs index 4a390351c..31ec44a73 100644 --- a/crates/libsy/src/algorithms/subagent_affinity_tests.rs +++ b/crates/libsy/src/algorithms/subagent_affinity_tests.rs @@ -7,6 +7,7 @@ //! *which* target delegated work belongs on, affinity decides *how long* that decision //! lives. +use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; @@ -17,9 +18,10 @@ use super::util::subagent::SubagentOverride; use crate::Result; use crate::core::algorithm::Driver; use crate::core::classifier::{Classification, Classifier, Score}; -use crate::core::testing::{echo, test_drive}; +use crate::core::testing::{echo, test_drive_with_models}; use switchyard_protocol::{ - Metadata, ModelId, Request, Response, completion_text, slice_to_header_map, text_request, + Category, Metadata, ModelId, Request, Response, completion_text, slice_to_header_map, + text_request, }; /// The cascade's terminal classifier: always picks the orchestrator. @@ -73,7 +75,9 @@ fn router() -> Arc { /// Runs one turn, returning the target that served it. async fn turn(router: &Arc, headers: &[(&str, &str)]) -> Result { - let (_, response) = test_drive(router.clone(), request(headers), echo()).await?; + let models = HashMap::from([(Category::Any, targets())]); + let (_, response) = + test_drive_with_models(router.clone(), request(headers), models, echo()).await?; Ok(response .llm_response .as_agg() diff --git a/crates/libsy/src/algorithms/util/affinity.rs b/crates/libsy/src/algorithms/util/affinity.rs index f557a28d6..92ede7567 100644 --- a/crates/libsy/src/algorithms/util/affinity.rs +++ b/crates/libsy/src/algorithms/util/affinity.rs @@ -24,7 +24,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use async_trait::async_trait; use parking_lot::Mutex; use serde::Deserialize; -use switchyard_protocol::{ContentBlock, Message, ModelId, Request, Role}; +use switchyard_protocol::{Category, ContentBlock, Message, ModelId, Request, Role}; use crate::core::algorithm::{Driver, RoutingIdentity}; use crate::core::classifier::{Classification, Classifier, Score}; @@ -242,17 +242,27 @@ where if self.release_on_user_turn && has_new_user_turn(&request.llm_request.messages) { return Ok((Classification::Scores(Vec::new()), None)); } - let assigned = self.assignments.lock().get(&key).cloned(); + let mut assignments = self.assignments.lock(); + let assigned = assignments.get(&key).cloned(); if assigned.is_some() && let Some(driver) = driver { driver.set_evidence(serde_json::json!({"source": "retained"})); } + let assigned = match (assigned.as_ref(), driver) { + (Some(target), Some(driver)) if !driver.models_for(Category::Any).contains(target) => { + if assigned.as_ref() == Some(target) { + assignments.remove(&key); + } + None + } + (assigned, _) => assigned, + }; Ok(( Classification::Scores(match assigned { Some(target) => vec![Score { confidence: 1.0, - target, + target: target.clone(), }], None => Vec::new(), }), diff --git a/crates/libsy/src/algorithms/util/escalation.rs b/crates/libsy/src/algorithms/util/escalation.rs index 00b47f98e..059a0eb2c 100644 --- a/crates/libsy/src/algorithms/util/escalation.rs +++ b/crates/libsy/src/algorithms/util/escalation.rs @@ -9,13 +9,14 @@ use serde::Deserialize; use serde_json::Value; -use switchyard_protocol::{ContentBlock, Message, ModelId, Role}; +use switchyard_protocol::{Category, ContentBlock, Message, Role}; use super::classifier_contract::{ClassifierContract, ClassifierContractConfig}; use super::llm_judge::{ ClassifierInput, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig, SerdeDecoder, StructuredJudge, }; +use crate::core::algorithm::Driver; use crate::core::classifier::{Classification, Score}; use crate::core::state::State; use crate::{LibsyError, Result}; @@ -117,25 +118,26 @@ pub(crate) type EscalationJudge = StructuredJudge) -> Classification { + fn to_classification( + &self, + verdict: Option<&EscalationVerdict>, + driver: &Driver, + ) -> Result { match verdict { - Some(verdict) if verdict.escalate => Classification::Scores(vec![Score { - target: self.capable.clone(), + Some(verdict) if verdict.escalate => Ok(Classification::Scores(vec![Score { + target: driver.first_model_for(Category::Capable)?.clone(), confidence: 1.0, - }]), - Some(_) => Classification::Scores(vec![Score { - target: self.efficient.clone(), + }])), + Some(_) => Ok(Classification::Scores(vec![Score { + target: driver.first_model_for(Category::Efficient)?.clone(), confidence: 1.0, - }]), - None => Classification::Ambiguous(Vec::new()), + }])), + None => Ok(Classification::Ambiguous(Vec::new())), } } } @@ -153,14 +155,11 @@ fn escalation_evidence( }) } -/// Builds the trajectory judge over `judge_target`, scoring `capable` when it escalates. +/// Builds the trajectory judge, scoring the runtime capable category when it escalates. /// /// Loads the packaged prompt and schema, so an unusable asset or an unusable `config` value /// fails here rather than on the first request. pub(crate) fn build_judge( - judge_target: ModelId, - capable: ModelId, - efficient: ModelId, contract_config: &ClassifierContractConfig, config: EscalationJudgeConfig, max_output_tokens: u64, @@ -175,8 +174,7 @@ pub(crate) fn build_judge( SerdeDecoder::new(), JudgeRuntimeConfig::new(max_output_tokens)?, ), - judge_target, - EscalationPolicy { capable, efficient }, + EscalationPolicy, ) .with_evidence(escalation_evidence)) } diff --git a/crates/libsy/src/algorithms/util/llm_judge.rs b/crates/libsy/src/algorithms/util/llm_judge.rs index a99fb82e5..5012e9f98 100644 --- a/crates/libsy/src/algorithms/util/llm_judge.rs +++ b/crates/libsy/src/algorithms/util/llm_judge.rs @@ -13,7 +13,7 @@ use async_trait::async_trait; use serde::de::DeserializeOwned; use serde_json::Value; use switchyard_protocol::{ - AggLlmResponse, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Role, + AggLlmResponse, Category, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Role, completion_text, }; @@ -199,19 +199,22 @@ pub trait Judge: Send + Sync { pub trait JudgePolicy: Send + Sync { type Verdict: Send + Sync; - fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification; + fn to_classification( + &self, + verdict: Option<&Self::Verdict>, + driver: &Driver, + ) -> Result; } type EvidenceFn = fn(&P, Option<&V>) -> Option; -/// A classifier that calls one judge target and routes through its verdict policy. +/// A classifier that calls the runtime judge models and routes through its verdict policy. pub struct JudgeClassifier where J: Judge, P: JudgePolicy, { judge: J, - target: ModelId, policy: P, evidence: Option>, } @@ -221,11 +224,10 @@ where J: Judge, P: JudgePolicy, { - /// Combines a judge target with a verdict policy. - pub fn new(judge: J, target: ModelId, policy: P) -> Self { + /// Combines a judge with a verdict policy. + pub fn new(judge: J, policy: P) -> Self { Self { judge, - target, policy, evidence: None, } @@ -239,7 +241,11 @@ where /// Adds fail-open evidence only for evidence-enabled judges and preserves an earlier decision. fn report_fail_open(&self, driver: &Driver, error: String, reason: &'static str) { - report_fail_open(self.target.as_str(), error, reason); + let judge_target = driver + .first_model_for(Category::Judge) + .map(|c| c.as_str()) + .unwrap_or("missing"); + report_fail_open(judge_target, error, reason); if self.evidence.is_some() { driver.set_evidence_if_empty(serde_json::json!({ "source": "fail_open", @@ -260,14 +266,15 @@ where state: &mut State, request: &Request, driver: &Driver, + judge_models: &[ModelId], ) -> Option { - let judge_model = self.target.as_str(); + let judge_model = judge_models.first()?.as_str(); tracing::info!(target = judge_model, "consulting llm judge"); let response = driver .call_model( self.judge.build_request(state, request), - vec![self.target.clone()], + judge_models.to_vec(), ) .await .inspect_err(|error| { @@ -343,14 +350,17 @@ where // A missing driver is a broken composition, not an unavailable judge. let Some(driver) = driver else { return Err(LibsyError::AlgorithmError { - message: format!( - "judge classifier for target {:?} requires a driver to call it", - self.target - ), + message: "judge classifier requires a driver to call it".to_string(), }); }; - let verdict = self.verdict(state, request, driver).await; - let classification = self.policy.to_classification(verdict.as_ref()); + let judge_models = driver.models_for(Category::Judge); + if judge_models.is_empty() { + return Err(LibsyError::AlgorithmError { + message: "no models available for category Judge".to_string(), + }); + } + let verdict = self.verdict(state, request, driver, judge_models).await; + let classification = self.policy.to_classification(verdict.as_ref(), driver)?; if let Some(evidence) = self .evidence .and_then(|evidence| evidence(&self.policy, verdict.as_ref())) @@ -428,21 +438,25 @@ mod tests { impl JudgePolicy for TestPolicy { type Verdict = TestVerdict; - fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification { + fn to_classification( + &self, + verdict: Option<&Self::Verdict>, + _driver: &Driver, + ) -> Result { let target = if verdict.is_some() { "verdict" } else { "no-verdict" }; - Classification::Scores(vec![Score { + Ok(Classification::Scores(vec![Score { target: ModelId::from(target), confidence: 1.0, - }]) + }])) } } fn classifier() -> JudgeClassifier { - JudgeClassifier::new(TestJudge, ModelId::from("judge"), TestPolicy) + JudgeClassifier::new(TestJudge, TestPolicy) } fn request() -> Request { @@ -557,7 +571,8 @@ mod tests { /// Serves the single offloaded judge call with `reply` through a standalone step receiver. async fn score_served_with(reply: Result) -> Result { - let (driver, step_rx) = Driver::new("test"); + let models = [(Category::Judge, vec![ModelId::from("judge")])].into(); + let (driver, step_rx) = Driver::new("test", models); let mut steps = tokio_stream::wrappers::ReceiverStream::new(step_rx); let classifier = classifier(); let mut state = State::default(); diff --git a/crates/libsy/src/algorithms/util/prompts.rs b/crates/libsy/src/algorithms/util/prompts.rs index d8bc939fd..1e66c4286 100644 --- a/crates/libsy/src/algorithms/util/prompts.rs +++ b/crates/libsy/src/algorithms/util/prompts.rs @@ -3,33 +3,14 @@ //! Adding text to a request on its way to the model it was routed to. //! -//! 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`] — 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 +//! These helpers mutate the normalized request. 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; - -use async_trait::async_trait; -use switchyard_protocol::{ContentBlock, InstructionBlock, Message, ModelId, Request, Role}; - -use crate::Result; -use crate::core::processor::{Event, Processor}; +use switchyard_protocol::{ContentBlock, InstructionBlock, Message, Request, Role}; /// Appends `note` to the request as conversation text. /// @@ -70,84 +51,26 @@ 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. -pub struct SystemPromptProcessor { - prompts: TargetPrompts, -} - -impl SystemPromptProcessor { - /// Hand each target the prompt configured for it. - pub fn new(prompts: TargetPrompts) -> Self { - Self { prompts } - } -} - -#[async_trait] -impl Processor for SystemPromptProcessor { - async fn process(&self, _state: &mut S, event: Event<'_>) -> Result<()> { - // The decision event carries both the routing outcome and the outbound request, - // so the target is read straight off it — whichever classifier picked it, and - // with nothing kept between turns. - let Event::Decision { - request, - selected_model_id, - } = event - else { - return Ok(()); - }; - 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(), - }], - }, - ); - drop_exact_replay(request); - Ok(()) - } +/// Prepends a system prompt and disables exact replay so the edit reaches the provider. +pub(crate) fn prepend_system_prompt(request: &mut Request, prompt: &str) { + request.llm_request.instructions.insert( + 0, + InstructionBlock { + role: Role::System, + content: vec![ContentBlock::Text { + text: prompt.to_string(), + }], + }, + ); + drop_exact_replay(request); } #[cfg(test)] mod tests { use super::*; - use switchyard_protocol::{LlmRequest, ModelId, ToolResult, text_request}; + use switchyard_protocol::{LlmRequest, ToolResult, text_request}; const NOTE: &str = "recovering from an error"; - const STRONG_PROMPT: &str = "diagnose before you edit"; - const WEAK_PROMPT: &str = "follow the settled plan"; /// Every test request carries the exact inbound body a codec keeps for /// same-format replay, so each assertion below also says what happens to it. @@ -244,126 +167,4 @@ mod tests { "a same-format hop would replay the body captured before the note" ); } - - /// The instruction text the request carries. - fn instructions(request: &Request) -> Vec { - request - .llm_request - .instructions - .iter() - .filter_map(|block| { - block.content.iter().find_map(|content| match content { - ContentBlock::Text { text } => Some(text.clone()), - _ => None, - }) - }) - .collect() - } - - /// Runs one outbound request routed to `target` through `processor`. - async fn run(processor: &SystemPromptProcessor, target: &'static str) -> Result { - let mut request = Request { - llm_request: LlmRequest { - preservation: preserved_body(), - ..LlmRequest::default() - }, - ..Request::default() - }; - let selected_model_id = ModelId::from(target); - processor - .process( - &mut (), - Event::Decision { - request: &mut request, - selected_model_id: &selected_model_id, - }, - ) - .await?; - Ok(request) - } - - fn prompts() -> TargetPrompts { - TargetPrompts::default() - .with("strong", STRONG_PROMPT) - .with("weak", WEAK_PROMPT) - } - - #[tokio::test] - async fn each_target_gets_its_own_prompt() -> Result<()> { - let processor = SystemPromptProcessor::new(prompts()); - for (target, expected) in [("strong", STRONG_PROMPT), ("weak", WEAK_PROMPT)] { - let request = run(&processor, target).await?; - assert_eq!(instructions(&request), vec![expected]); - assert!( - !replays_exactly(&request), - "{target}: a same-format hop would replay the body captured before the prompt" - ); - } - Ok(()) - } - - #[tokio::test] - async fn an_unconfigured_target_is_left_untouched() -> Result<()> { - // One target's prompt must not leak onto another, whatever ran before. - let processor = - SystemPromptProcessor::new(TargetPrompts::default().with("strong", STRONG_PROMPT)); - assert_eq!( - instructions(&run(&processor, "strong").await?), - vec![STRONG_PROMPT] - ); - let untouched = run(&processor, "weak").await?; - assert!(instructions(&untouched).is_empty()); - assert!( - replays_exactly(&untouched), - "an untouched request must keep its lossless same-format replay" - ); - Ok(()) - } - - #[tokio::test] - async fn the_prompt_leads_the_client_instructions() -> Result<()> { - let processor = SystemPromptProcessor::new(prompts()); - let mut request = Request::default(); - request.llm_request.instructions.push(InstructionBlock { - role: Role::System, - content: vec![ContentBlock::Text { - text: "you are a coding agent".to_string(), - }], - }); - let selected_model_id = ModelId::from("strong"); - - processor - .process( - &mut (), - Event::Decision { - request: &mut request, - selected_model_id: &selected_model_id, - }, - ) - .await?; - - assert_eq!( - instructions(&request), - vec![STRONG_PROMPT, "you are a coding agent"] - ); - Ok(()) - } - - #[tokio::test] - async fn the_inbound_request_is_left_alone() -> Result<()> { - // The inbound hook runs before the cascade has picked anything. - let processor = SystemPromptProcessor::new(prompts()); - let mut request = Request::default(); - processor - .process( - &mut (), - Event::Request { - request: &mut request, - driver: None, - }, - ) - .await?; - assert!(instructions(&request).is_empty()); - Ok(()) - } } diff --git a/crates/libsy/src/algorithms/util/stage.rs b/crates/libsy/src/algorithms/util/stage.rs index d5da79596..5edd42c29 100644 --- a/crates/libsy/src/algorithms/util/stage.rs +++ b/crates/libsy/src/algorithms/util/stage.rs @@ -28,8 +28,7 @@ use crate::core::algorithm::Driver; use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::state::{State, StateValue}; use crate::observability::meter; -use switchyard_protocol::ModelId; -use switchyard_protocol::Request; +use switchyard_protocol::{Category, Request}; /// Turn depth below which stall signals stay quiet — early no-write turns are /// normal exploration, not a stall. @@ -98,51 +97,6 @@ impl Tier { } } -/// The targets a stage router's two tiers route to. -/// -/// The tiers are a fixed pair, but their targets are whatever the deployment -/// calls them, so the classifier scores onto those names and the routed call -/// reaches the right model. -#[derive(Clone, Debug)] -pub struct StageTargets { - capable: ModelId, - efficient: ModelId, -} - -impl StageTargets { - /// Name the targets the two tiers route to. - pub fn new(capable: impl Into, efficient: impl Into) -> Self { - Self { - capable: capable.into(), - efficient: efficient.into(), - } - } - - /// The target `tier` routes to. - pub fn name(&self, tier: Tier) -> &ModelId { - match tier { - Tier::Capable => &self.capable, - Tier::Efficient => &self.efficient, - } - } - - /// The tier a routed target belongs to, or `None` for one outside the pair. - pub fn tier_for(&self, target: &ModelId) -> Option { - if *target == self.capable { - Some(Tier::Capable) - } else if *target == self.efficient { - Some(Tier::Efficient) - } else { - None - } - } - - /// The tier label for a routed target, or `None` for one outside the pair. - pub fn label_for(&self, target: &ModelId) -> Option<&'static str> { - self.tier_for(target).map(Tier::label) - } -} - /// Which tier to default to when the scorer is not confident. #[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize)] #[serde(rename_all = "snake_case")] @@ -556,18 +510,16 @@ impl HandoffNoteConfig { /// With [`with_handoff_notes`](Self::with_handoff_notes) it also splices a note /// into the request explaining why the signals sent the turn where they did. pub struct StageClassifier { - targets: StageTargets, mode: PickerMode, confidence_threshold: f64, handoff_notes: Option, } impl StageClassifier { - /// Scores onto `targets`, with the given default tier (`mode`) and - /// `confidence_threshold`. - pub fn new(targets: StageTargets, mode: PickerMode, confidence_threshold: f64) -> Self { + /// Scores onto the runtime capable and efficient models, with the given + /// default tier (`mode`) and `confidence_threshold`. + pub fn new(mode: PickerMode, confidence_threshold: f64) -> Self { Self { - targets, mode, confidence_threshold, handoff_notes: None, @@ -622,24 +574,27 @@ impl Classifier for StageClassifier { probability, confidence, } => { - let target = self.targets.name(tier); + let category = match tier { + Tier::Capable => Category::Capable, + Tier::Efficient => Category::Efficient, + }; + let driver = driver.ok_or(crate::LibsyError::NoTargets)?; + let target = driver.first_model_for(category)?; record_decision_source(state, source); record_routing_decision(source, target); // Only a resolved turn routes on this classifier's target, so it // is the only branch whose tier the signals actually chose — an // ambiguous turn is decided further down the cascade. self.apply_handoff_note(request, tier, source); - if let Some(driver) = driver { - let evidence = match (source, confidence) { - (DecisionSource::Dimensions, Some(confidence)) => serde_json::json!({ - "source": source.as_str(), - "confidence": confidence, - "threshold": self.confidence_threshold, - }), - _ => serde_json::json!({"source": source.as_str()}), - }; - driver.set_evidence(evidence); - } + let evidence = match (source, confidence) { + (DecisionSource::Dimensions, Some(confidence)) => serde_json::json!({ + "source": source.as_str(), + "confidence": confidence, + "threshold": self.confidence_threshold, + }), + _ => serde_json::json!({"source": source.as_str()}), + }; + driver.set_evidence(evidence); // Prefer pick_tier's own confidence (e.g. 1.0 for an Override) // over re-deriving it from the neutral 0.5 placeholder. let conf = confidence.unwrap_or_else(|| 2.0 * (probability - 0.5).abs()); @@ -660,7 +615,19 @@ impl Classifier for StageClassifier { mod tests { use super::*; use serde_json::json; - use switchyard_protocol::{Metadata, Request, WireFormat, text_request}; + use switchyard_protocol::{Metadata, ModelId, Request, WireFormat, text_request}; + + fn driver() -> Driver { + Driver::new( + "stage_test", + [ + (Category::Capable, vec![ModelId::from("strong")]), + (Category::Efficient, vec![ModelId::from("weak")]), + ] + .into(), + ) + .0 + } fn signal_from(messages: serde_json::Value) -> ToolSignals { let raw_request = Some(json!({"model": "m", "messages": messages})); @@ -748,11 +715,6 @@ mod tests { // ─── StageClassifier ───────────────────────────────────────────────── - /// Tiers named the way a deployment would name them. - fn tiers() -> StageTargets { - StageTargets::new("strong", "weak") - } - /// A `State` carrying `signal` as its tool signals. fn state_with(signal: ToolSignals) -> State { State { @@ -766,7 +728,7 @@ mod tests { // No tool activity yet — nothing to score, so the signals have no opinion // and the turn belongs to whatever the cascade has behind them. let mut state = State::default(); - let classification = StageClassifier::new(tiers(), PickerMode::EfficientFirst, 0.5) + let classification = StageClassifier::new(PickerMode::EfficientFirst, 0.5) .score(&mut state, &mut Request::default(), None) .await?; assert!(classification.0.argmax(false)?.is_none()); @@ -785,8 +747,9 @@ mod tests { ..Default::default() }; let mut state = state_with(signal); - let classification = StageClassifier::new(tiers(), PickerMode::EfficientFirst, 0.5) - .score(&mut state, &mut Request::default(), None) + let driver = driver(); + let classification = StageClassifier::new(PickerMode::EfficientFirst, 0.5) + .score(&mut state, &mut Request::default(), Some(&driver)) .await?; match classification.0 { Classification::Scores(scores) => { @@ -816,8 +779,9 @@ mod tests { ..Default::default() }; let mut state = state_with(signal); - let classification = StageClassifier::new(tiers(), PickerMode::EfficientFirst, 0.5) - .score(&mut state, &mut Request::default(), None) + let driver = driver(); + let classification = StageClassifier::new(PickerMode::EfficientFirst, 0.5) + .score(&mut state, &mut Request::default(), Some(&driver)) .await?; match classification.0 { Classification::Scores(scores) => { @@ -834,7 +798,7 @@ mod tests { // A quiet signal corroborates neither axis, so the scorer abstains and // records why. let mut state = state_with(ToolSignals::default()); - let classification = StageClassifier::new(tiers(), PickerMode::EfficientFirst, 0.5) + let classification = StageClassifier::new(PickerMode::EfficientFirst, 0.5) .score(&mut state, &mut Request::default(), None) .await?; assert!(classification.0.argmax(false)?.is_none()); @@ -904,7 +868,7 @@ mod tests { /// A classifier that hands the capable tier an escalation note, gated to /// signal-driven escalations. fn noting_classifier(mode: PickerMode) -> StageClassifier { - StageClassifier::new(tiers(), mode, 0.5).with_handoff_notes(HandoffNoteConfig::new( + StageClassifier::new(mode, 0.5).with_handoff_notes(HandoffNoteConfig::new( ESCALATION, Some(DEESCALATION.to_string()), true, @@ -941,9 +905,10 @@ mod tests { async fn a_signal_driven_escalation_carries_the_note() -> Result<()> { let mut state = state_with(critical()); let mut request = request(); + let driver = driver(); noting_classifier(PickerMode::EfficientFirst) - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, Some(&driver)) .await?; assert_eq!(trailing_text(&request), Some(format!("hi|{ESCALATION}"))); @@ -956,10 +921,13 @@ mod tests { // of escalated turns each carries one. Nothing tracks the previous tier. let classifier = noting_classifier(PickerMode::EfficientFirst); let mut state = state_with(critical()); + let driver = driver(); for _ in 0..3 { let mut request = request(); - classifier.score(&mut state, &mut request, None).await?; + classifier + .score(&mut state, &mut request, Some(&driver)) + .await?; assert_eq!(trailing_text(&request), Some(format!("hi|{ESCALATION}"))); } Ok(()) @@ -976,9 +944,10 @@ mod tests { }; let mut state = state_with(signal); let mut request = request(); + let driver = driver(); noting_classifier(PickerMode::EfficientFirst) - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, Some(&driver)) .await?; assert_eq!(trailing_text(&request), Some(format!("hi|{DEESCALATION}"))); @@ -1005,9 +974,10 @@ mod tests { async fn no_note_when_notes_are_unconfigured() -> Result<()> { let mut state = state_with(critical()); let mut request = request(); + let driver = driver(); - StageClassifier::new(tiers(), PickerMode::EfficientFirst, 0.5) - .score(&mut state, &mut request, None) + StageClassifier::new(PickerMode::EfficientFirst, 0.5) + .score(&mut state, &mut request, Some(&driver)) .await?; assert_eq!(trailing_text(&request), Some("hi".to_string())); diff --git a/crates/libsy/src/algorithms/util/target_selector.rs b/crates/libsy/src/algorithms/util/target_selector.rs index 4e36b648a..68132075e 100644 --- a/crates/libsy/src/algorithms/util/target_selector.rs +++ b/crates/libsy/src/algorithms/util/target_selector.rs @@ -3,28 +3,23 @@ //! Deterministic target selection from a validated JSON classifier verdict. -use std::collections::BTreeMap; - use jsonptr::PointerBuf; use serde_json::Value; use super::llm_judge::JudgePolicy; +use crate::core::algorithm::Driver; use crate::core::classifier::{Classification, Score}; use crate::{LibsyError, Result}; -use switchyard_protocol::ModelId; +use switchyard_protocol::Category; -/// Maps one string field in a validated verdict to a configured routing target. +/// Maps one string field in a validated verdict to a runtime model category. pub(crate) struct TargetSelectorPolicy { selector: PointerBuf, - targets: BTreeMap, } impl TargetSelectorPolicy { /// Parses a JSON Pointer used to read validated verdicts. - pub(crate) fn new( - selector: impl Into, - targets: BTreeMap, - ) -> Result { + pub(crate) fn new(selector: impl Into) -> Result { let selector = PointerBuf::parse(selector.into()).map_err(|error| LibsyError::AlgorithmError { message: format!("policy selector is not a valid JSON Pointer: {error}"), @@ -34,24 +29,28 @@ impl TargetSelectorPolicy { message: "policy selector must identify a response field".to_string(), }); } - Ok(Self { selector, targets }) + Ok(Self { selector }) } } impl JudgePolicy for TargetSelectorPolicy { type Verdict = Value; - fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification { - let target = verdict + fn to_classification( + &self, + verdict: Option<&Self::Verdict>, + driver: &Driver, + ) -> Result { + let category = verdict .and_then(|verdict| self.selector.resolve(verdict).ok()) .and_then(Value::as_str) - .and_then(|label| self.targets.get(label)); - match target { - Some(target) => Classification::Scores(vec![Score { - target: target.clone(), + .and_then(|label| label.parse::().ok()); + match category { + Some(category) => Ok(Classification::Scores(vec![Score { + target: driver.first_model_for(category)?.clone(), confidence: 1.0, - }]), - None => Classification::Ambiguous(vec![]), + }])), + None => Ok(Classification::Ambiguous(vec![])), } } } @@ -59,46 +58,50 @@ impl JudgePolicy for TargetSelectorPolicy { #[cfg(test)] mod tests { use serde_json::json; + use switchyard_protocol::ModelId; use super::*; use crate::Result; + fn driver() -> Driver { + Driver::new( + "test", + [(Category::Capable, vec![ModelId::from("model/opus")])].into(), + ) + .0 + } + #[test] - fn a_verdict_selects_its_mapped_target() -> Result<()> { - let policy = TargetSelectorPolicy::new( - "/decision/target", - BTreeMap::from([ - ("opus".to_string(), ModelId::from("model/opus")), - ("sonnet".to_string(), ModelId::from("model/sonnet")), - ]), + fn a_verdict_selects_its_runtime_category() -> Result<()> { + let policy = TargetSelectorPolicy::new("/decision/target")?; + let classification = policy.to_classification( + Some(&json!({ + "decision": {"target": "capable"} + })), + &driver(), )?; - let classification = policy.to_classification(Some(&json!({ - "decision": {"target": "sonnet"} - }))); assert_eq!( classification.argmax(false)?.map(|score| score.target), - Some(ModelId::from("model/sonnet")) + Some(ModelId::from("model/opus")) ); Ok(()) } #[test] fn a_missing_or_unknown_target_abstains() -> Result<()> { - let policy = TargetSelectorPolicy::new( - "/target", - BTreeMap::from([("sonnet".to_string(), ModelId::from("model/sonnet"))]), - )?; + let policy = TargetSelectorPolicy::new("/target")?; + let driver = driver(); assert_eq!( policy - .to_classification(Some(&json!({"target": "unknown"}))) + .to_classification(Some(&json!({"target": "unknown"})), &driver)? .argmax(false)?, None ); assert_eq!( policy - .to_classification(Some(&json!({"reason": "missing"}))) + .to_classification(Some(&json!({"reason": "missing"})), &driver)? .argmax(false)?, None ); @@ -107,7 +110,7 @@ mod tests { #[test] fn an_invalid_json_pointer_is_rejected() { - let result = TargetSelectorPolicy::new("/target~2name", BTreeMap::new()); + let result = TargetSelectorPolicy::new("/target~2name"); assert!(matches!(result, Err(LibsyError::AlgorithmError { message }) if message.contains("valid JSON Pointer"))); } diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 6e01c20c7..c91df1e79 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -4,7 +4,10 @@ //! The [`Algorithm`] trait and its [`Driver`] — the orchestration contract every //! algorithm implements and the offload channel it uses for routing-time model calls. -use std::{future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc, time::Instant}; +use std::{ + collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc, + time::Instant, +}; use async_trait::async_trait; use futures::{FutureExt, Stream, StreamExt}; @@ -21,7 +24,7 @@ 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::{Category, ModelId, Request, Response}; use crate::{DriverError, LibsyError, Result, observability}; @@ -121,16 +124,25 @@ impl RoutingOutcome { #[derive(Clone)] pub struct Driver { step_tx: mpsc::Sender>, + /// The owning algorithm's telemetry label, stamped onto every call this driver publishes. algorithm: String, + /// Run-scoped evidence shared by driver clones and attached only to a successful outcome. evidence: Arc>>, + + /// The models the algorithm can use, grouped by category. + /// Within a category they are typically ordered best-first. + models: HashMap>, } impl Driver { /// Build an empty driver with its step channel ready. Created per call by /// [`run_stream`](Algorithm::run_stream). Also returns the Step receiver. - pub(crate) fn new(algorithm: &str) -> (Self, mpsc::Receiver>) { + pub(crate) fn new( + algorithm: &str, + models: HashMap>, + ) -> (Self, mpsc::Receiver>) { // Capacity one keeps the algorithm paced by the stream consumer. It limits queued steps, // not model calls already pulled from the stream, which can still run at the same time. // A larger buffer would use more memory and let the algorithm run farther ahead with @@ -141,6 +153,7 @@ impl Driver { step_tx, algorithm: algorithm.to_string(), evidence: Arc::new(Mutex::new(None)), + models, }, step_rx, ) @@ -216,6 +229,20 @@ impl Driver { result } + /// The available models for this category, typically ordered best-first. + pub fn models_for(&self, category: Category) -> &[ModelId] { + self.models.get(&category).map_or(&[], |v| v.as_slice()) + } + + /// The first available model for `category`. + pub fn first_model_for(&self, category: Category) -> Result<&ModelId> { + self.models_for(category) + .first() + .ok_or_else(|| LibsyError::AlgorithmError { + message: format!("no models available for category {category:?}"), + }) + } + /// Emit the terminal step: [`Step::Done`] on `Ok`, or an `Err` stream /// item on failure. Internal: called once by [`run_stream`](Algorithm::run_stream) /// when the algorithm finishes. @@ -267,13 +294,14 @@ pub enum Step { pub async fn drive( algorithm: Arc, request: Request, + models: HashMap>, serve: F, ) -> Result where F: Fn(CallModel) -> Fut, Fut: Future>, { - let stream = algorithm.run_stream(request); + let stream = algorithm.run_stream(request, models); tokio::pin!(stream); let mut in_flight = futures::stream::FuturesUnordered::new(); @@ -416,8 +444,12 @@ pub trait Algorithm: Send + Sync + 'static { /// the stream aborts the spawned algorithm task. /// /// Every invocation owns a separate [`Driver`]. - fn run_stream(self: Arc, request: Request) -> StepStream { - let (driver, step_rx) = Driver::new(self.name()); + fn run_stream( + self: Arc, + request: Request, + models: HashMap>, + ) -> StepStream { + let (driver, step_rx) = Driver::new(self.name(), models); let span = observability::run_span(self.name(), &request); let handle = tokio::spawn( async move { @@ -562,7 +594,7 @@ mod tests { tokio::time::timeout(std::time::Duration::from_secs(1), async { // Distinct oneshots keep reverse-order replies paired with their producers, and a // retained call remains pending until the host responds. - let (driver, mut step_rx) = Driver::new("test"); + let (driver, mut step_rx) = Driver::new("test", HashMap::new()); let first_driver = driver.clone(); let mut first = tokio::spawn(async move { first_driver @@ -619,7 +651,7 @@ mod tests { ); // Dropping the host-facing promise closes only that call's reply channel. - let (driver, mut step_rx) = Driver::new("test"); + let (driver, mut step_rx) = Driver::new("test", HashMap::new()); let producer = tokio::spawn(async move { driver .call_model(request(), vec![ModelId::from("dropped")]) @@ -639,7 +671,7 @@ mod tests { )); // A standalone driver reports the typed step receiver disappearing at its next call. - let (driver, step_rx) = Driver::new("test"); + let (driver, step_rx) = Driver::new("test", HashMap::new()); drop(step_rx); let result = driver .call_model(request(), vec![ModelId::from("closed")]) @@ -742,7 +774,7 @@ mod tests { async fn run_offloads_via_promise_then_finishes() -> Result<()> { // Every call is offloaded via a promise the orchestrator surfaces as a // `CallModel` step for us to fulfill. - let stream = orch(target_set(&["offload/model"])).run_stream(request()); + let stream = orch(target_set(&["offload/model"])).run_stream(request(), HashMap::new()); tokio::pin!(stream); let mut saw_call = false; @@ -854,7 +886,7 @@ mod tests { // A client-less target offloads its call; we fulfill the promise with an // Err, which must flow back through `call_model_target` into the algorithm and // out as an error step — not a response. - let stream = orch(target_set(&["offload/model"])).run_stream(request()); + let stream = orch(target_set(&["offload/model"])).run_stream(request(), HashMap::new()); tokio::pin!(stream); let mut saw_error = false; @@ -926,7 +958,7 @@ mod tests { dropped: dropped.clone(), }); - let stream = algo.run_stream(request()); + let stream = algo.run_stream(request(), HashMap::new()); started_rx .recv() .await @@ -963,7 +995,7 @@ mod tests { } let algo: Arc = Arc::new(Panicky); - let stream = algo.run_stream(request()); + let stream = algo.run_stream(request(), HashMap::new()); tokio::pin!(stream); let mut saw_error = false; diff --git a/crates/libsy/src/core/testing.rs b/crates/libsy/src/core/testing.rs index aeb570988..b42a3deac 100644 --- a/crates/libsy/src/core/testing.rs +++ b/crates/libsy/src/core/testing.rs @@ -13,11 +13,14 @@ //! The closure is async so a fake can block on a barrier, wait on a notify, or never //! resolve, which is what the concurrency, hedging, and fan-out tests need. +use std::collections::HashMap; use std::future::Future; use std::sync::Arc; use futures::future::BoxFuture; -use switchyard_protocol::{LlmClientError, LlmResponse, ModelId, Request, Response, text_response}; +use switchyard_protocol::{ + Category, LlmClientError, LlmResponse, ModelId, Request, Response, text_response, +}; use crate::core::algorithm::{Algorithm, CallModel}; use crate::{LibsyError, Result}; @@ -46,10 +49,20 @@ pub(crate) async fn test_drive( algorithm: Arc, request: Request, serve: impl Serve, +) -> Result<(ModelId, Response)> { + test_drive_with_models(algorithm, request, HashMap::new(), serve).await +} + +/// Drive one request with an explicit runtime model set. +pub(crate) async fn test_drive_with_models( + algorithm: Arc, + request: Request, + models: HashMap>, + serve: impl Serve, ) -> Result<(ModelId, Response)> { let serve = Arc::new(serve); let routing_serve = Arc::clone(&serve); - let outcome = crate::drive(algorithm, request, move |call| { + let outcome = crate::drive(algorithm, request, models, move |call| { fulfill(Arc::clone(&routing_serve), call) }) .await?; diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index ad70b8da2..33b18bf62 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -31,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::append_note; pub use algorithms::util::subagent::{SubagentGate, SubagentOverride}; pub use algorithms::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignals}; @@ -39,8 +39,8 @@ pub use algorithms::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignals}; // core (scorer, picker, and the `StageClassifier`). pub use algorithms::util::stage::{ CodingAgentDimensions, DECISION_SOURCE_KEY, DecisionSource, HandoffNoteConfig, PickOutcome, - PickerMode, ScoreResult, StageClassifier, StageTargets, Tier, clear_fall_open, - dimensions_from_signal, pick_tier, score_signal, set_fall_open, + PickerMode, ScoreResult, StageClassifier, Tier, clear_fall_open, dimensions_from_signal, + pick_tier, score_signal, set_fall_open, }; mod observability; diff --git a/crates/prefill-router/tests/unit/algorithm.rs b/crates/prefill-router/tests/unit/algorithm.rs index 9e71759f2..f62a4241f 100644 --- a/crates/prefill-router/tests/unit/algorithm.rs +++ b/crates/prefill-router/tests/unit/algorithm.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -66,7 +67,7 @@ fn forward() -> ( } async fn selected(route: Arc, request: Request) -> libsy::Result { - let outcome = libsy::drive(route, request, |call| async move { + let outcome = libsy::drive(route, request, HashMap::new(), |call| async move { call.respond(Ok(switchyard_protocol::Response { llm_response: switchyard_protocol::LlmResponse::Agg( switchyard_protocol::text_response(None, "unused"), diff --git a/crates/protocol/src/category.rs b/crates/protocol/src/category.rs new file mode 100644 index 000000000..203671c59 --- /dev/null +++ b/crates/protocol/src/category.rs @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Category is a group of models + +use std::{collections::HashMap, str::FromStr}; + +use crate::ModelId; + +/// A group of models +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub enum Category { + /// When the category doesn't matter: Random, Passthrough, etc. + Any, + /// High accuracy and cost models. + Capable, + /// Lower accuracy and cost models. + Efficient, + /// Models the algorithm can use to decide. + Judge, +} + +impl Category { + /// Convenience function to create a HashMap suitable for passing to `run_stream` and family. + pub fn to_map(category: Category, names: &[&str]) -> HashMap> { + [( + category, + names.iter().map(|name| ModelId::from(*name)).collect(), + )] + .into() + } +} + +impl FromStr for Category { + type Err = String; + fn from_str(s: &str) -> Result { + let c = match s { + "capable" => Self::Capable, + "efficient" => Self::Efficient, + "judge" => Self::Judge, + "any" => Self::Any, + x => { + return Err(format!("Invalid Category '{x}'")); + } + }; + Ok(c) + } +} diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index 15352b44f..fcf2b3782 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -4,6 +4,7 @@ #![warn(missing_docs)] #![doc = include_str!("../README.md")] +pub mod category; pub mod client; pub mod envelope; pub mod format; @@ -12,6 +13,7 @@ pub mod metadata; pub mod model_id; pub mod stream; +pub use category::*; pub use client::*; pub use envelope::*; pub use format::*; diff --git a/crates/switchyard-nemo-relay-plugin/src/runtime.rs b/crates/switchyard-nemo-relay-plugin/src/runtime.rs index f9e675ef4..351733e8d 100644 --- a/crates/switchyard-nemo-relay-plugin/src/runtime.rs +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -200,7 +200,7 @@ impl SwitchyardRuntime { .unwrap_or_else(|poisoned| poisoned.into_inner()) .push(observation); }); - match route.execute(request, Some(observer)).await { + match route.execute(request, route.models(), Some(observer)).await { Ok(output) => { self.emit_observations(&mut events, take_observations(&observations), &metadata); let served_model = output.response.served_model().map(|model| model.as_str()); @@ -749,6 +749,7 @@ mod tests { None, None, Vec::new(), + HashMap::new(), ); SwitchyardRuntime { runner: Runner::new(vec![(ModelId::from(model), route)]), diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index 79adca287..11612bf1a 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -19,8 +19,8 @@ use switchyard_libsy::{ StepStream, TaskClassifierConfig, }; use switchyard_protocol::{ - LlmClientError, LlmResponse, LlmResponseStream, LlmResponseStreamEvent, Metadata, ModelId, - Request, Response, + Category, LlmClientError, LlmResponse, LlmResponseStream, LlmResponseStreamEvent, Metadata, + ModelId, Request, Response, }; use tokio::sync::Mutex; @@ -36,6 +36,10 @@ fn classify_trigger(session_affinity: bool) -> ClassifyTrigger { } } +fn parse_category(value: &str) -> PyResult { + value.parse().map_err(PyValueError::new_err) +} + /// Convert Python-owned headers into the request metadata expected by libsy. fn header_map_from_python(headers: &HashMap) -> PyResult { let mut result = http::HeaderMap::new(); @@ -187,19 +191,10 @@ struct PyLlmClassifierConfig { impl PyLlmClassifierConfig { /// Configure capability routing between efficient and capable targets. #[staticmethod] - #[pyo3(signature = (judge_target, efficient_target, capable_target, *, config))] - fn capability( - py: Python<'_>, - judge_target: String, - efficient_target: String, - capable_target: String, - config: Py, - ) -> PyResult { + #[pyo3(signature = (*, config))] + fn capability(py: Python<'_>, config: Py) -> PyResult { Ok(Self { inner: LlmClassifierConfig::Capability { - judge_target: ModelId::new(judge_target), - efficient_target: ModelId::new(efficient_target), - capable_target: ModelId::new(capable_target), config: config.bind(py).try_borrow()?.clone_core(), }, }) @@ -207,20 +202,11 @@ impl PyLlmClassifierConfig { /// Configure response-based escalation between efficient and capable targets. #[staticmethod] - #[pyo3(signature = (judge_target, efficient_target, capable_target, *, config))] - fn escalation( - py: Python<'_>, - judge_target: String, - efficient_target: String, - capable_target: String, - config: Py, - ) -> PyResult { + #[pyo3(signature = (*, config))] + fn escalation(py: Python<'_>, config: Py) -> PyResult { let config = config.bind(py).try_borrow()?; Ok(Self { inner: LlmClassifierConfig::Escalation { - judge_target: ModelId::new(judge_target), - efficient_target: ModelId::new(efficient_target), - capable_target: ModelId::new(capable_target), contract: config.contract.clone(), config: config.judge.clone(), max_output_tokens: config.max_output_tokens, @@ -228,25 +214,18 @@ impl PyLlmClassifierConfig { }) } - /// Configure schema-driven routing across named targets. + /// Configure schema-driven routing across runtime model categories. #[staticmethod] - #[pyo3(signature = (judge_target, targets, *, default_target, config))] + #[pyo3(signature = (*, default_target, config))] fn custom( py: Python<'_>, - judge_target: String, - targets: Vec<(String, String)>, default_target: String, config: Py, ) -> PyResult { let config = config.bind(py).try_borrow()?.clone_core(); Ok(Self { inner: LlmClassifierConfig::Custom { - judge_target: ModelId::new(judge_target), - targets: targets - .into_iter() - .map(|(name, target)| (name, ModelId::new(target))) - .collect(), - default_target, + default_target: parse_category(&default_target)?, config, }, }) @@ -320,14 +299,12 @@ fn classifier_contract( skip_from_py_object )] struct PyLlmFallback { - judge_target: String, config: Py, } impl PyLlmFallback { fn clone_core(&self, py: Python<'_>) -> PyResult { Ok(LlmFallback { - judge_target: ModelId::new(self.judge_target.clone()), config: self.config.bind(py).try_borrow()?.clone_core(), }) } @@ -336,12 +313,9 @@ impl PyLlmFallback { #[pymethods] impl PyLlmFallback { #[new] - #[pyo3(signature = (judge_target, *, config))] - fn new(judge_target: String, config: Py) -> Self { - Self { - judge_target, - config, - } + #[pyo3(signature = (*, config))] + fn new(config: Py) -> Self { + Self { config } } } @@ -670,11 +644,12 @@ impl PyAlgorithm { /// [`Metadata`] exactly as an HTTP host would (`Metadata::from_headers`), /// so metadata-driven algorithms see the same signals in Python as when /// served over HTTP. - #[pyo3(signature = (request, headers=None))] + #[pyo3(signature = (request, models, headers=None))] fn run_stream( &self, request: &Bound<'_, PyAny>, - headers: Option>, + models: HashMap>, + headers: Option>, ) -> PyResult { let headers = headers.as_ref().map(header_map_from_python).transpose()?; let request = Request { @@ -682,9 +657,14 @@ impl PyAlgorithm { raw_request: None, metadata: headers.map(|headers| Metadata::from_headers(&headers)), }; + let mut typed_models = HashMap::new(); + for (k, v) in models { + let category: Category = k.parse().map_err(PyValueError::new_err)?; + typed_models.insert(category, v.into_iter().map(ModelId::from).collect()); + } let stream = { let _guard = pyo3_async_runtimes::tokio::get_runtime().enter(); - Arc::clone(&self.inner).run_stream(request) + Arc::clone(&self.inner).run_stream(request, typed_models) }; Ok(PyRunStream { inner: Arc::new(Mutex::new(stream)), @@ -744,17 +724,10 @@ fn noop_algorithm() -> PyAlgorithm { /// Construct random routing over targets with optional relative weights and seed. #[pyfunction(name = "random")] -#[pyo3(signature = (targets, *, weights=None, seed=None))] -fn random_algorithm( - targets: Vec, - weights: Option>, - seed: Option, -) -> PyResult { - let model_ids = targets.into_iter().map(ModelId::new).collect(); - let algorithm = Random::new(model_ids, weights, seed).map_err(|error| match error { - RustLibsyError::NoTargets => PyValueError::new_err("random requires at least one target"), - other => PyValueError::new_err(other.to_string()), - })?; +#[pyo3(signature = (weights=None, seed=None))] +fn random_algorithm(weights: Option>, seed: Option) -> PyResult { + let algorithm = + Random::new(weights, seed).map_err(|other| PyValueError::new_err(other.to_string()))?; Ok(PyAlgorithm { inner: Arc::new(algorithm), }) @@ -771,24 +744,12 @@ fn llm_classifier_algorithm( /// Construct capability classifier routing. #[pyfunction(name = "llm_task_classifier")] -#[pyo3(signature = ( - judge_target, - efficient_target, - capable_target, - *, - config -))] +#[pyo3(signature = (*, config))] fn llm_task_classifier_algorithm( py: Python<'_>, - judge_target: String, - efficient_target: String, - capable_target: String, config: Py, ) -> PyResult { build_llm_classifier(LlmClassifierConfig::Capability { - judge_target: ModelId::new(judge_target), - efficient_target: ModelId::new(efficient_target), - capable_target: ModelId::new(capable_target), config: config.bind(py).try_borrow()?.clone_core(), }) } @@ -804,8 +765,6 @@ fn build_llm_classifier(config: LlmClassifierConfig) -> PyResult { /// Construct signal-driven stage routing with an optional LLM classifier fallback. #[pyfunction(name = "stage_router")] #[pyo3(signature = ( - capable_target, - efficient_target, *, picker, confidence_threshold, @@ -820,8 +779,6 @@ fn build_llm_classifier(config: LlmClassifierConfig) -> PyResult { #[allow(clippy::too_many_arguments)] fn stage_router_algorithm( py: Python<'_>, - capable_target: String, - efficient_target: String, picker: &str, confidence_threshold: f64, recent_window: Option, @@ -841,8 +798,6 @@ fn stage_router_algorithm( ))); } }; - let capable = ModelId::new(capable_target); - let efficient = ModelId::new(efficient_target); let mut config = StageRouterConfig::new(mode, confidence_threshold); config.recent_window = recent_window; config.handoff_notes = match (escalation_note, deescalation_note) { @@ -858,18 +813,14 @@ fn stage_router_algorithm( } (None, None) => None, }; - if let Some(prompt) = capable_system_prompt { - config.tier_prompts = config.tier_prompts.with(capable.clone(), prompt); - } - if let Some(prompt) = efficient_system_prompt { - config.tier_prompts = config.tier_prompts.with(efficient.clone(), prompt); - } + config.capable_system_prompt = capable_system_prompt; + config.efficient_system_prompt = efficient_system_prompt; config.llm_fallback = classifier .map(|classifier| classifier.bind(py).try_borrow()?.clone_core(py)) .transpose()?; - let algorithm = StageRouter::new(capable, efficient, config) - .map_err(|error| PyValueError::new_err(error.to_string()))?; + let algorithm = + StageRouter::new(config).map_err(|error| PyValueError::new_err(error.to_string()))?; Ok(PyAlgorithm { inner: Arc::new(algorithm), }) diff --git a/crates/switchyard-runner/src/algorithm.rs b/crates/switchyard-runner/src/algorithm.rs index 0e1ede289..1b32f7c68 100644 --- a/crates/switchyard-runner/src/algorithm.rs +++ b/crates/switchyard-runner/src/algorithm.rs @@ -3,7 +3,7 @@ //! Schema-neutral algorithm configuration and construction. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::error::Error; use std::fmt::{Display, Formatter}; use std::path::PathBuf; @@ -17,7 +17,7 @@ use libsy::{ StageRouter, StageRouterConfig, SubagentRouter, SubagentRouterConfig, TaskClassifierConfig, }; use serde::Deserialize; -use switchyard_protocol::ModelId; +use switchyard_protocol::{Category, ModelId}; /// Error returned when an algorithm description cannot be constructed. #[derive(Debug)] @@ -97,6 +97,7 @@ enum LlmClassifierModeConfig { #[derive(Clone, Debug)] struct CapabilityClassifierRouteConfig { + classifier_target: String, strong_target: String, weak_target: String, base_threshold: f64, @@ -111,6 +112,7 @@ struct CapabilityClassifierRouteConfig { #[derive(Clone, Debug)] struct EscalationClassifierRouteConfig { + classifier_target: String, strong_target: String, weak_target: String, prompt: Option, @@ -121,9 +123,8 @@ struct EscalationClassifierRouteConfig { #[derive(Clone, Debug)] struct CustomClassifierRouteConfig { - classifier_target: String, - targets: Vec, - default_target: String, + models: CategoryModelConfig, + default_target: Category, prompt: String, response_schema: String, policy: ClassifierPolicyConfig, @@ -133,6 +134,41 @@ struct CustomClassifierRouteConfig { max_output_tokens: u64, } +/// Runtime model groups for a custom classifier. +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct CategoryModelConfig { + /// All completion targets in fallback order. + pub any: Vec, + /// Ordered candidates for judge calls. + pub judge: Vec, + /// Ordered capable-tier models. + pub capable: Vec, + /// Ordered efficient-tier models. + pub efficient: Vec, +} + +impl CategoryModelConfig { + fn get(&self, category: Category) -> &Vec { + match category { + Category::Any => &self.any, + Category::Judge => &self.judge, + Category::Capable => &self.capable, + Category::Efficient => &self.efficient, + } + } + + fn all_names(&self) -> Vec<&str> { + self.any + .iter() + .chain(&self.judge) + .chain(&self.capable) + .chain(&self.efficient) + .map(String::as_str) + .collect() + } +} + /// Settings for an `llm_classifier` route. Which fields are required depends on /// the [`ClassifierMode`]; using a field from the wrong mode is an error. #[derive(Clone, Debug, Default, Deserialize)] @@ -171,9 +207,9 @@ pub struct LlmClassifierRouteConfig { /// Escalation mode: how many escalate verdicts latch the session, and how /// much of the transcript the judge sees. pub escalation: Option, - /// Custom mode: the target names the policy may pick from. - pub targets: Option>, - /// Custom mode: target used when the judge fails or its verdict cannot be routed. + /// Custom mode: runtime model groups. + pub models: Option, + /// Custom mode: category used when the judge fails or its verdict cannot be routed. pub default_target: Option, /// Custom mode: JSON Schema the verdict must match, written as a string. pub response_schema: Option, @@ -200,18 +236,10 @@ impl SubagentRouteConfig { match self { Self::Passthrough { target } => vec![target], Self::LlmClassifier(classifier) => classifier - .targets - .iter() - .flatten() - .map(String::as_str) - .collect(), - } - } - - fn classifier_target_name(&self) -> Option<&str> { - match self { - Self::LlmClassifier(classifier) => Some(&classifier.classifier_target), - Self::Passthrough { .. } => None, + .models + .as_ref() + .map(CategoryModelConfig::all_names) + .unwrap_or_default(), } } } @@ -445,11 +473,10 @@ impl AlgorithmSpec { .map(String::as_str) .collect(), ClassifierMode::Custom => config - .targets - .iter() - .flatten() - .map(String::as_str) - .collect(), + .models + .as_ref() + .map(CategoryModelConfig::all_names) + .unwrap_or_default(), } } Self::StageRouter { @@ -492,32 +519,24 @@ impl AlgorithmSpec { pub fn callable_target_names(&self) -> Vec<&str> { let mut names = self.routing_target_names(); match self { - Self::LlmClassifier { config, .. } => names.push(&config.classifier_target), - Self::Passthrough { - subagents: Some(subagents), - .. - } => names.extend(subagents.classifier_target_name()), - Self::StageRouter { - classifier, - subagents, - .. - } => { - if let Some(classifier) = classifier { - names.push(&classifier.target); - } - if let Some(subagents) = subagents { - names.extend(subagents.classifier_target_name()); - } + Self::LlmClassifier { config, .. } + if !matches!( + config.mode.unwrap_or(if config.escalation.is_some() { + ClassifierMode::Escalation + } else { + ClassifierMode::Capability + }), + ClassifierMode::Custom + ) => + { + names.push(&config.classifier_target) } - Self::Composite { - classifier, - subagents, + Self::StageRouter { + classifier: Some(classifier), .. - } => { + } => names.push(&classifier.target), + Self::Composite { classifier, .. } => { names.push(&classifier.target); - if let Some(subagents) = subagents { - names.extend(subagents.classifier_target_name()); - } } Self::Advisor { advisor_target, .. } => names.push(advisor_target), _ => {} @@ -525,6 +544,69 @@ impl AlgorithmSpec { names } + /// Target names grouped as the runtime [`Driver`](libsy::Driver) expects them. + pub(crate) fn runtime_model_names( + &self, + route_name: &str, + ) -> AlgorithmResult>> { + let mut models = match self { + Self::Noop { .. } => HashMap::new(), + Self::Random { targets, .. } | Self::PrefillRouter { targets, .. } => { + category_models([(Category::Any, targets.clone())]) + } + Self::Passthrough { target, .. } => { + category_models([(Category::Any, vec![target.clone()])]) + } + Self::LlmClassifier { config } => { + classifier_runtime_model_names(config.classifier_mode(route_name)?) + } + Self::StageRouter { + tiers, classifier, .. + } => { + let mut models = category_models([ + (Category::Capable, vec![tiers.capable_target.clone()]), + (Category::Efficient, vec![tiers.efficient_target.clone()]), + ( + Category::Any, + vec![tiers.capable_target.clone(), tiers.efficient_target.clone()], + ), + ]); + if let Some(classifier) = classifier { + models.insert(Category::Judge, vec![classifier.target.clone()]); + } + models + } + Self::Composite { + classifier, stage, .. + } => category_models([ + (Category::Judge, vec![classifier.target.clone()]), + (Category::Capable, vec![stage.capable_target.clone()]), + (Category::Efficient, vec![stage.efficient_target.clone()]), + ( + Category::Any, + vec![stage.capable_target.clone(), stage.efficient_target.clone()], + ), + ]), + Self::Advisor { + executor_target, .. + } => category_models([(Category::Any, vec![executor_target.clone()])]), + }; + + let subagents = match self { + Self::Passthrough { subagents, .. } + | Self::StageRouter { subagents, .. } + | Self::Composite { subagents, .. } => subagents.as_ref(), + _ => None, + }; + if let Some(subagents) = subagents { + merge_category_models( + &mut models, + subagent_runtime_model_names(subagents, route_name)?, + ); + } + Ok(models) + } + /// Response target and routing-only dependency for routers that answer while routing. pub(crate) fn routing_response_and_dependency(&self) -> Option<(&str, &str)> { match self { @@ -567,6 +649,79 @@ impl AlgorithmSpec { build_algorithm(context, self, targets) } } + +fn category_models( + entries: impl IntoIterator)>, +) -> HashMap> { + entries.into_iter().collect() +} + +fn merge_category_models( + models: &mut HashMap>, + additions: HashMap>, +) { + for (category, names) in additions { + models.entry(category).or_default().extend(names); + } +} + +fn custom_runtime_model_names(config: &CategoryModelConfig) -> HashMap> { + [ + Category::Any, + Category::Judge, + Category::Capable, + Category::Efficient, + ] + .into_iter() + .map(|category| (category, config.get(category).clone())) + .collect() +} + +fn classifier_runtime_model_names( + config: LlmClassifierModeConfig, +) -> HashMap> { + match config { + LlmClassifierModeConfig::Capability(config) => category_models([ + (Category::Judge, vec![config.classifier_target]), + (Category::Efficient, vec![config.weak_target.clone()]), + (Category::Capable, vec![config.strong_target.clone()]), + ( + Category::Any, + vec![config.weak_target, config.strong_target], + ), + ]), + LlmClassifierModeConfig::Escalation(config) => category_models([ + (Category::Judge, vec![config.classifier_target]), + (Category::Efficient, vec![config.weak_target.clone()]), + (Category::Capable, vec![config.strong_target.clone()]), + ( + Category::Any, + vec![config.strong_target, config.weak_target], + ), + ]), + LlmClassifierModeConfig::Custom(config) => custom_runtime_model_names(&config.models), + } +} + +fn subagent_runtime_model_names( + config: &SubagentRouteConfig, + route_name: &str, +) -> AlgorithmResult>> { + match config { + SubagentRouteConfig::Passthrough { target } => { + Ok(category_models([(Category::Any, vec![target.clone()])])) + } + SubagentRouteConfig::LlmClassifier(config) => { + let LlmClassifierModeConfig::Custom(config) = config.classifier_mode(route_name)? + else { + return Err(AlgorithmConfigError::new(format!( + "route {route_name}: subagents llm_classifier only supports mode custom" + ))); + }; + Ok(custom_runtime_model_names(&config.models)) + } + } +} impl LlmClassifierRouteConfig { fn classifier_mode(&self, route_name: &str) -> AlgorithmResult { let Self { @@ -583,7 +738,7 @@ impl LlmClassifierRouteConfig { response_format_type, max_output_tokens, escalation, - targets, + models, default_target, response_schema, policy, @@ -607,13 +762,14 @@ impl LlmClassifierRouteConfig { reject_custom_fields( route_name, "capability", - targets, + models, default_target, response_schema, policy, )?; Ok(LlmClassifierModeConfig::Capability( CapabilityClassifierRouteConfig { + classifier_target: classifier_target.clone(), strong_target: required_classifier_field( route_name, "strong_target", @@ -643,7 +799,7 @@ impl LlmClassifierRouteConfig { reject_custom_fields( route_name, "escalation", - targets, + models, default_target, response_schema, policy, @@ -665,6 +821,7 @@ impl LlmClassifierRouteConfig { } Ok(LlmClassifierModeConfig::Escalation( EscalationClassifierRouteConfig { + classifier_target: classifier_target.clone(), strong_target: required_classifier_field( route_name, "strong_target", @@ -683,7 +840,8 @@ impl LlmClassifierRouteConfig { )) } ClassifierMode::Custom => { - if strong_target.is_some() + if !classifier_target.is_empty() + || strong_target.is_some() || weak_target.is_some() || base_threshold.is_some() || threshold_step.is_some() @@ -696,13 +854,18 @@ impl LlmClassifierRouteConfig { } Ok(LlmClassifierModeConfig::Custom( CustomClassifierRouteConfig { - classifier_target: classifier_target.clone(), - targets: required_classifier_field(route_name, "targets", targets)?, + models: required_classifier_field(route_name, "models", models)?, default_target: required_classifier_field( route_name, "default_target", default_target, - )?, + )? + .parse() + .map_err(|error| { + AlgorithmConfigError::new(format!( + "llm_classifier route {route_name} has invalid default_target: {error}" + )) + })?, prompt: required_classifier_field(route_name, "prompt", prompt)?, response_schema: required_classifier_field( route_name, @@ -724,15 +887,12 @@ impl LlmClassifierRouteConfig { fn reject_custom_fields( route_name: &str, mode: &str, - targets: &Option>, + models: &Option, default_target: &Option, response_schema: &Option, policy: &Option, ) -> AlgorithmResult<()> { - if targets.is_some() - || default_target.is_some() - || response_schema.is_some() - || policy.is_some() + if models.is_some() || default_target.is_some() || response_schema.is_some() || policy.is_some() { return Err(AlgorithmConfigError::new(format!( "llm_classifier route {route_name} mode {mode} cannot use custom classifier fields" @@ -775,23 +935,21 @@ fn build_subagent_router_config( "route {route_name}: subagents llm_classifier only supports mode custom" ))); }; - let judge_target = - resolve_target_model_id(route_name, &config.classifier_target, targets)?; let resolved_targets = config - .targets - .iter() - .map(|name| { - resolve_target_model_id(route_name, name, targets) - .map(|target| (name.clone(), target)) - }) + .models + .all_names() + .into_iter() + .map(|name| resolve_target_model_id(route_name, name, targets)) .collect::>>()?; - let default_target = resolved_targets - .iter() - .find(|(name, _)| *name == config.default_target) - .map(|(_, target)| target.clone()) + let default_target = config + .models + .get(config.default_target) + .first() + .map(|name| resolve_target_model_id(route_name, name, targets)) + .transpose()? .ok_or_else(|| { AlgorithmConfigError::new(format!( - "route {route_name}: subagents llm_classifier default_target {:?} must be one of its configured targets", + "route {route_name}: subagents llm_classifier has no model for default category {:?}", config.default_target )) })?; @@ -811,14 +969,8 @@ fn build_subagent_router_config( ); classifier_config.recent_turn_window = config.recent_turn_window; classifier_config.max_output_tokens = config.max_output_tokens; - let subagent_targets = resolved_targets - .iter() - .map(|(_, target)| target.clone()) - .collect(); let classifier = Arc::new( LlmTaskClassifier::new(LlmClassifierConfig::Custom { - judge_target, - targets: resolved_targets, default_target: config.default_target, config: classifier_config, }) @@ -830,7 +982,7 @@ fn build_subagent_router_config( })?, ); Ok(SubagentRouterConfig { - targets: subagent_targets, + targets: resolved_targets, classifier, default_target, classify_trigger: config.classify_trigger, @@ -866,15 +1018,8 @@ fn build_algorithm( ) -> AlgorithmResult> { match config { AlgorithmSpec::Noop { .. } => Ok(Arc::new(Noop {})), - AlgorithmSpec::Random { - targets: names, - weights, - seed, - .. - } => { - let target_set = - resolve_targets(route_name, names.iter().map(String::as_str), targets)?; - let algorithm = Random::new(target_set, weights.clone(), *seed).map_err(|error| { + AlgorithmSpec::Random { weights, seed, .. } => { + let algorithm = Random::new(weights.clone(), *seed).map_err(|error| { AlgorithmConfigError::with_source( format!("random route {route_name}: {error}"), error, @@ -882,11 +1027,8 @@ fn build_algorithm( })?; Ok(Arc::new(algorithm)) } - AlgorithmSpec::Passthrough { - target, subagents, .. - } => { - let parent_target = resolve_target_model_id(route_name, target, targets)?; - let algorithm = Passthrough::new(parent_target); + AlgorithmSpec::Passthrough { subagents, .. } => { + let algorithm = Passthrough::default(); let parent: Arc = Arc::new(algorithm); attach_subagent_router(route_name, parent, subagents.as_ref(), targets) } @@ -894,14 +1036,9 @@ fn build_algorithm( config: classifier_config, .. } => { - let classifier = - resolve_target_model_id(route_name, &classifier_config.classifier_target, targets)?; let mode = classifier_config.classifier_mode(route_name)?; let algorithm = match mode { LlmClassifierModeConfig::Capability(config) => { - let strong = - resolve_target_model_id(route_name, &config.strong_target, targets)?; - let weak = resolve_target_model_id(route_name, &config.weak_target, targets)?; let classifier_config = TaskClassifierConfig { base_threshold: config.base_threshold, threshold_step: config.threshold_step, @@ -913,20 +1050,11 @@ fn build_algorithm( max_output_tokens: config.max_output_tokens, }; LlmTaskClassifier::new(LlmClassifierConfig::Capability { - judge_target: classifier, - efficient_target: weak, - capable_target: strong, config: classifier_config, }) } LlmClassifierModeConfig::Escalation(config) => { - let strong = - resolve_target_model_id(route_name, &config.strong_target, targets)?; - let weak = resolve_target_model_id(route_name, &config.weak_target, targets)?; LlmTaskClassifier::new(LlmClassifierConfig::Escalation { - judge_target: classifier, - efficient_target: weak, - capable_target: strong, contract: classifier_contract(config.prompt.as_deref()) .with_response_format_type(config.response_format_type), config: config.judge, @@ -934,14 +1062,6 @@ fn build_algorithm( }) } LlmClassifierModeConfig::Custom(config) => { - let resolved_targets = config - .targets - .iter() - .map(|name| { - resolve_target_model_id(route_name, name, targets) - .map(|target| (name.clone(), target)) - }) - .collect::>>()?; let response_schema = serde_json::from_str(&config.response_schema).map_err( |error| { AlgorithmConfigError::with_source( @@ -962,8 +1082,6 @@ fn build_algorithm( classifier_config.recent_turn_window = config.recent_turn_window; classifier_config.max_output_tokens = config.max_output_tokens; LlmTaskClassifier::new(LlmClassifierConfig::Custom { - judge_target: classifier, - targets: resolved_targets, default_target: config.default_target, config: classifier_config, }) @@ -985,36 +1103,25 @@ fn build_algorithm( .. } => { let StageTierConfig { - capable_target, - efficient_target, confidence_threshold, recent_turn_window, handoff_notes, + .. } = tiers; if matches!(picker, PickerMode::CapableFirst) { tracing::warn!( "stage_router route {route_name} uses picker \"capable_first\", which is experimental: published thresholds and routing results all come from \"efficient_first\", so there is no calibrated confidence_threshold for it and no measured accuracy or cost. Use \"efficient_first\" unless you are running your own calibration." ); } - let capable = resolve_target_model_id(route_name, capable_target, targets)?; - let efficient = resolve_target_model_id(route_name, efficient_target, targets)?; let mut config = StageRouterConfig::new(*picker, *confidence_threshold); config.recent_window = *recent_turn_window; config.handoff_notes = handoff_notes.clone(); // 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 - .as_ref() - .map(|classifier| { - resolve_target_model_id(route_name, &classifier.target, targets).map( - |judge_target| LlmFallback { - judge_target, - config: classifier.task_classifier_config(), - }, - ) - }) - .transpose()?; - let algorithm = StageRouter::new(capable, efficient, config).map_err(|error| { + config.llm_fallback = classifier.as_ref().map(|classifier| LlmFallback { + config: classifier.task_classifier_config(), + }); + let algorithm = StageRouter::new(config).map_err(|error| { AlgorithmConfigError::with_source( format!("stage_router route {route_name}: {error}"), error, @@ -1028,19 +1135,15 @@ fn build_algorithm( stage, subagents, } => { - let capable = resolve_target_model_id(route_name, &stage.capable_target, targets)?; - let efficient = resolve_target_model_id(route_name, &stage.efficient_target, targets)?; - let judge_target = resolve_target_model_id(route_name, &classifier.target, targets)?; let mut stage_config = StageRouterConfig::new(PickerMode::EfficientFirst, stage.confidence_threshold); stage_config.recent_window = stage.recent_turn_window; stage_config.handoff_notes = stage.handoff_notes.clone(); let config = CompositeRouterConfig { - judge_target, judge: classifier.task_classifier_config(), stage: stage_config, }; - let algorithm = CompositeRouter::new(capable, efficient, config).map_err(|error| { + let algorithm = CompositeRouter::new(config).map_err(|error| { AlgorithmConfigError::with_source( format!("composite route {route_name}: {error}"), error, @@ -1172,17 +1275,6 @@ fn default_classifier_max_output_tokens() -> u64 { TaskClassifierConfig::default().max_output_tokens } -fn resolve_targets<'a>( - route_name: &str, - names: impl IntoIterator, - targets: &BTreeMap, -) -> AlgorithmResult> { - names - .into_iter() - .map(|name| resolve_target_model_id(route_name, name, targets)) - .collect() -} - fn resolve_target_model_id( route_name: &str, name: &str, diff --git a/crates/switchyard-runner/src/config.rs b/crates/switchyard-runner/src/config.rs index a4bb309aa..c62bf207a 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -15,7 +15,7 @@ use switchyard_llm_client::{ AuxiliaryOperation, Backend, ClientRouter, DEFAULT_MAX_RETRIES, HttpBackendConfig, ModelConfig, TranslatingLlmClient, }; -use switchyard_protocol::{ModelId, RoutedLlmClient, WireFormat}; +use switchyard_protocol::{Category, ModelId, RoutedLlmClient, WireFormat}; use crate::{ AlgorithmSpec, AuxiliaryTarget, CallerAuthKind, DecisionTarget, ModelCapabilities, Route, @@ -208,6 +208,25 @@ impl DeploymentConfig { .into_iter() .filter_map(|name| self.decision_target(name)) .collect(); + let models = config + .algorithm + .runtime_model_names(route_name) + .map_err(|error| RunnerError::configuration_source(error.to_string(), error))? + .into_iter() + .map(|(category, names)| { + let models = names + .into_iter() + .map(|name| { + targets.get(&name).cloned().ok_or_else(|| { + RunnerError::configuration(format!( + "route references unknown target {name}" + )) + }) + }) + .collect::>>()?; + Ok((category, models)) + }) + .collect::>>>()?; let route = Route::new( algorithm, route_clients, @@ -216,6 +235,7 @@ impl DeploymentConfig { anthropic_auxiliary_target, responses_auxiliary_target, decision_targets, + models, ); routes.push((config.id.clone(), route)); } @@ -698,7 +718,20 @@ target = "weak" fn public_runner_from_toml_builds_a_deployment() -> RunnerResult<()> { let runner = Runner::from_toml(VALID_CONFIG)?; - assert!(runner.route("switchyard/classifier").is_some()); + let classifier = runner + .route("switchyard/classifier") + .expect("classifier route should exist"); + let models = classifier.models(); + assert_eq!( + models[&Category::Judge], + [ModelId::from("classifier/model")] + ); + assert_eq!(models[&Category::Efficient], [ModelId::from("weak/model")]); + assert_eq!(models[&Category::Capable], [ModelId::from("strong/model")]); + assert_eq!( + models[&Category::Any], + [ModelId::from("weak/model"), ModelId::from("strong/model")] + ); assert!(runner.route("switchyard/passthrough").is_some()); Ok(()) } @@ -716,11 +749,10 @@ target = "weak" configured.push_str( r#"type = "llm_classifier" mode = "custom" -classifier_target = "classifier" -targets = ["strong", "weak"] -default_target = "weak" +models = { judge = ["classifier"], capable = ["strong"], efficient = ["weak"], any = ["strong", "weak"] } +default_target = "efficient" prompt = "Select a target for this delegated task." -response_schema = '{"type":"object","properties":{"target":{"type":"string","enum":["strong","weak"]}},"required":["target"],"additionalProperties":false}' +response_schema = '{"type":"object","properties":{"target":{"type":"string","enum":["capable","efficient"]}},"required":["target"],"additionalProperties":false}' policy = { type = "target_selector", selector = "/target" } classify_trigger = "new_session""#, ); @@ -1061,20 +1093,6 @@ classifier_magic = true ), "unknown target missing", ), - ( - VALID_CONFIG.replace( - "targets = [\"strong\", \"weak\"]", - "targets = [\"strong\", \"strong\"]", - ), - "random targets must be unique", - ), - ( - VALID_CONFIG.replace( - "targets = [\"strong\", \"weak\"]", - "targets = [\"strong\", \"weak\"]\nweights = [1]", - ), - "expected 2 weights, got 1", - ), ( VALID_CONFIG.replace( "targets = [\"strong\", \"weak\"]", diff --git a/crates/switchyard-runner/src/lib.rs b/crates/switchyard-runner/src/lib.rs index efb0441e4..aae5cd25b 100644 --- a/crates/switchyard-runner/src/lib.rs +++ b/crates/switchyard-runner/src/lib.rs @@ -10,7 +10,7 @@ mod route; mod runner; pub use algorithm::{ - AdvisorTriggerConfig, AlgorithmConfigError, AlgorithmSpec, ClassifierMode, + AdvisorTriggerConfig, AlgorithmConfigError, AlgorithmSpec, CategoryModelConfig, ClassifierMode, ClassifierPolicyConfig, LlmClassifierRouteConfig, StageClassifierConfig, SubagentRouteConfig, }; pub use failure::{RouteErrorKind, RouteErrorPhase, RouteErrorSummary, stream_error_summary}; diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index 89f6e722c..e0bb02552 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -3,13 +3,13 @@ //! One configured algorithm and the clients that serve its targets. -use std::error::Error; use std::sync::Arc; +use std::{collections::HashMap, error::Error}; use libsy::{Algorithm, LibsyError, RoutingOutcome}; use serde_json::Value; use switchyard_llm_client::{AuxiliaryOperation, ClientRouter, RunObserver, TranslatingLlmClient}; -use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat}; +use switchyard_protocol::{Category, LlmClientError, ModelId, Request, Response, WireFormat}; use thiserror::Error; use crate::DecisionTarget; @@ -125,6 +125,7 @@ pub struct Route { anthropic_auxiliary_target: Option, responses_auxiliary_target: Option, decision_targets: Vec, + models: HashMap>, } /// The selected model and untouched response produced by a route execution. @@ -135,6 +136,7 @@ pub struct RunOutput { impl Route { /// Creates a fully configured execution route. + #[allow(clippy::too_many_arguments)] pub fn new( algorithm: Arc, clients: ClientRouter, @@ -143,6 +145,7 @@ impl Route { anthropic_auxiliary_target: Option, responses_auxiliary_target: Option, decision_targets: Vec, + models: HashMap>, ) -> Self { Self { algorithm, @@ -152,6 +155,7 @@ impl Route { anthropic_auxiliary_target, responses_auxiliary_target, decision_targets, + models, } } @@ -178,6 +182,11 @@ impl Route { .cloned() } + /// Returns the models grouped for one algorithm execution. + pub fn models(&self) -> HashMap> { + self.models.clone() + } + /// Rejects a caller format incompatible with forwarded credentials. pub fn check_caller_format(&self, input_format: WireFormat) -> Result<(), RunnerError> { if let Some(kind) = self.caller_auth @@ -192,12 +201,14 @@ impl Route { pub async fn execute( &self, request: Request, + models: HashMap>, observer: Option, ) -> Result { let (selected_model, response) = switchyard_llm_client::run( Arc::clone(&self.algorithm), self.clients.clone(), request, + models, observer, ) .await?; @@ -208,10 +219,19 @@ impl Route { } /// Completes routing-time calls without serving a post-routing completion. - pub async fn decide(&self, request: Request) -> Result { - switchyard_llm_client::decide(Arc::clone(&self.algorithm), self.clients.clone(), request) - .await - .map_err(Into::into) + pub async fn decide( + &self, + request: Request, + models: HashMap>, + ) -> Result { + switchyard_llm_client::decide( + Arc::clone(&self.algorithm), + self.clients.clone(), + request, + models, + ) + .await + .map_err(Into::into) } /// Executes a model-bearing provider operation through a compatible target. diff --git a/crates/switchyard-runner/src/runner.rs b/crates/switchyard-runner/src/runner.rs index 36ecb6410..ff30757ee 100644 --- a/crates/switchyard-runner/src/runner.rs +++ b/crates/switchyard-runner/src/runner.rs @@ -86,6 +86,10 @@ impl Runner { }) } + pub fn model_ids(&self) -> impl Iterator { + self.routes.iter().map(|(id, _)| id) + } + /// Returns the validated API root used for unmatched HTTP requests. pub fn fallback_base_url(&self) -> Option<&str> { self.fallback_base_url.as_deref() diff --git a/crates/switchyard-runner/tests/route.rs b/crates/switchyard-runner/tests/route.rs index df03c147d..724c444e3 100644 --- a/crates/switchyard-runner/tests/route.rs +++ b/crates/switchyard-runner/tests/route.rs @@ -9,8 +9,8 @@ use async_trait::async_trait; use futures_util::StreamExt; use switchyard_llm_client::{ClientRouter, RunObservation}; use switchyard_protocol::{ - LlmClientError, LlmResponse, ModelId, Request, Response, RoutedLlmClient, text_request, - text_response, + Category, LlmClientError, LlmResponse, ModelId, Request, Response, RoutedLlmClient, + text_request, text_response, }; use switchyard_runner::{AlgorithmSpec, ModelCapabilities, Route}; @@ -54,6 +54,7 @@ fn plugin_route(client: Arc) -> Route { None, None, Vec::new(), + [(Category::Any, vec![ModelId::from("semantic-target")])].into(), ) } @@ -69,9 +70,9 @@ async fn plugin_shaped_route_executes_without_runner_model_or_toml() { llm_request: text_request(Some("arbitrary-upstream-model".to_string()), "hello"), ..Request::default() }; - + let models = [(Category::Any, vec![ModelId::from("semantic-target")])].into(); let output = route - .execute(request, Some(observer)) + .execute(request, models, Some(observer)) .await .expect("route should execute"); @@ -133,8 +134,9 @@ async fn route_returns_stream_without_polling_it() { ..Request::default() }; + let models = [(Category::Any, vec![ModelId::from("semantic-target")])].into(); let output = route - .execute(request, None) + .execute(request, models, None) .await .expect("stream handle should be returned"); @@ -145,25 +147,6 @@ async fn route_returns_stream_without_polling_it() { assert_eq!(polls.load(Ordering::SeqCst), 0); } -#[test] -fn algorithm_build_reports_unknown_configured_target() { - let spec = AlgorithmSpec::Random { - targets: vec!["missing".to_string()], - weights: None, - seed: None, - }; - - let error = match spec.build("plugin", &BTreeMap::new()) { - Ok(_) => panic!("unknown target should fail"), - Err(error) => error, - }; - - assert_eq!( - error.to_string(), - "route plugin references unknown target missing" - ); -} - // Preserve checkpoint target order and explicit TOML overrides. #[test] fn prefill_router_config_preserves_target_order_and_overrides() { diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 8b6f06de8..3e2ac9204 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -13,7 +13,7 @@ mod sse; mod stats; mod usage_metrics; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::error::Error; use std::fmt::{Display, Formatter}; use std::future::Future; @@ -194,6 +194,7 @@ impl ServerState { None, None, Vec::new(), + HashMap::new(), ), ) }) @@ -725,7 +726,8 @@ async fn decision( .as_deref() .map(ModelId::from) .unwrap_or_default(); - let mut outcome = match route.decide(request).await { + + let mut outcome = match route.decide(request, route.models()).await { Ok(outcome) => outcome, Err(error) => return runner_error(error), }; @@ -1031,7 +1033,8 @@ async fn handle_llm_request( state.stats.clone(), state.routing_log.clone().zip(routing_log_context.clone()), ); - let output = match route.execute(request, Some(observer)).await { + + let output = match route.execute(request, route.models(), Some(observer)).await { Ok(output) => output, Err(error) => return runner_error(error), }; diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 0f6567092..16518827b 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -22,8 +22,9 @@ use serde_json::{Value, json}; use switchyard_llm_client::{ Backend, ClientRouter, HttpBackendConfig, ModelConfig, TranslatingLlmClient, }; -use switchyard_protocol::ModelId; use switchyard_protocol::RoutedLlmClient; +use switchyard_protocol::{Category, ModelId, WireFormat}; +use switchyard_runner::{DecisionTarget, ModelCapabilities, Route, Runner}; use switchyard_server::config::load_server_state; use switchyard_server::{ DEFAULT_MAX_REQUEST_BODY_BYTES, ServerState, build_llm_router, build_switchyard_router, @@ -305,7 +306,7 @@ async fn upstream_chat( if requests_invalid_verdict { r#"{"decision":{"target":"unknown"}}"# } else { - r#"{"decision":{"target":"premium"}}"# + r#"{"decision":{"target":"efficient"}}"# } } else if body .pointer("/response_format/json_schema/schema/properties/escalate") @@ -527,16 +528,37 @@ fn random_state_with_retries( let entries = routes .iter() .map(|(route_model, targets)| { - let target_set = targets.iter().map(|model| ModelId::from(*model)).collect(); - let algorithm: Arc = Arc::new(Random::new(target_set, None, None)?); + let algorithm: Arc = Arc::new(Random::new(None, None)?); + let decision_targets = targets + .iter() + .map(|model| DecisionTarget { + target: (*model).to_string(), + model: ModelId::from(*model), + format: WireFormat::OpenAiChat, + base_url: base_url.to_string(), + extra_body: BTreeMap::new(), + }) + .collect(); Ok(( ModelId::from(*route_model), - algorithm, - ClientRouter::single(Arc::clone(&client)), + Route::new( + algorithm, + ClientRouter::single(Arc::clone(&client)), + None, + ModelCapabilities::default(), + None, + None, + decision_targets, + [( + Category::Any, + targets.iter().map(|model| ModelId::from(*model)).collect(), + )] + .into(), + ), )) }) .collect::>>()?; - Ok(ServerState::new(entries)?) + ServerState::from_runner(Runner::new(entries)).map_err(Into::into) } async fn test_app(routes: &[(&str, &[&str])]) -> TestResult<(MockUpstream, Router)> { @@ -1426,8 +1448,7 @@ base_threshold = 0.5 } #[tokio::test] -async fn custom_classifier_routes_four_targets_and_falls_back_on_an_invalid_verdict() -> TestResult -{ +async fn custom_classifier_uses_categories_and_falls_back_on_an_invalid_verdict() -> TestResult { let upstream = MockUpstream::start().await?; let state = load_test_config(&format!( r#" @@ -1461,9 +1482,8 @@ llm_client = "upstream" id = "switchyard/custom" type = "llm_classifier" mode = "custom" -classifier_target = "classifier" -targets = ["weak", "middle", "strong", "premium"] -default_target = "strong" +models = {{ judge = ["classifier"], capable = ["strong", "premium"], efficient = ["middle", "weak"], any = ["weak", "middle", "strong", "premium"] }} +default_target = "capable" prompt = "CUSTOM MULTI TARGET" response_schema = ''' {{ @@ -1472,7 +1492,7 @@ response_schema = ''' "decision": {{ "type": "object", "properties": {{ - "target": {{"type": "string", "enum": ["weak", "middle", "strong", "premium"]}} + "target": {{"type": "string", "enum": ["any", "judge", "capable", "efficient"]}} }}, "required": ["target"], "additionalProperties": false @@ -1492,7 +1512,7 @@ selector = "/decision/target" let app = build_switchyard_router(state); for (task, selected) in [ - ("route this task", "model/premium"), + ("route this task", "model/middle"), ("return an invalid verdict", "model/strong"), ] { let response = send( @@ -1533,7 +1553,7 @@ selector = "/decision/target" assert_eq!( judge_call["response_format"]["json_schema"]["schema"]["properties"]["decision"]["properties"] ["target"]["enum"], - json!(["weak", "middle", "strong", "premium"]) + json!(["any", "judge", "capable", "efficient"]) ); Ok(()) } diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index 2708b68db..78efe59ca 100644 --- a/docs/reference/toml_schema.md +++ b/docs/reference/toml_schema.md @@ -180,12 +180,12 @@ checkpoint = "/models/router.pt" ### `llm_classifier` Runs one of three judge-backed modes: `capability`, `escalation`, or `custom`. -`classifier_target` and `max_output_tokens` apply to all three. +`max_output_tokens` applies to all three. | Key | Required | Default | Meaning | |---|:---:|---|---| | `mode` | No | `capability` | Classifier behavior. Set it explicitly for new configurations. | -| `classifier_target` | Yes | — | Target the judge is called through. Not a routing destination. | +| `classifier_target` | Capability, escalation | — | Target the judge is called through. Not a routing destination. Custom mode uses `models.judge`. | | `max_output_tokens` | No | `4096` | Maximum completion tokens for the judge verdict. Must be at least `1`. | | `response_format_type` | No | `json_schema` | Structured-output mode for capability and escalation judges. Use `json_object` when the provider does not support JSON Schema; Switchyard adds the schema to the prompt and validates the verdict locally. Custom mode always uses its configured JSON Schema. | @@ -218,12 +218,15 @@ Escalation mode serves the weak target first and judges the completed turn. See Existing configurations that contain `escalation` but omit `mode` remain valid. Custom mode validates the judge's JSON against `response_schema`, resolves the -policy selector, and routes to any configured target label. +policy selector, and routes to a runtime model category. | Key | Required | Default | Meaning | |---|:---:|---|---| -| `targets` | Yes | — | Two or more target names available to the policy. | -| `default_target` | Yes | — | Target used when the judge fails or its verdict cannot be routed. | +| `models.any` | Yes | — | All completion targets in fallback order. | +| `models.judge` | Yes | — | Ordered judge candidates. | +| `models.capable` | Yes | — | Ordered capable candidates. The first is selected. | +| `models.efficient` | Yes | — | Ordered efficient candidates. The first is selected. | +| `default_target` | Yes | — | Category used when the judge fails or its verdict cannot be routed. Must be `any`, `judge`, `capable`, or `efficient`. | | `prompt` | Yes | — | Judge system prompt. The configured inner schema is sent separately as structured-output configuration. | | `response_schema` | Yes | — | Inner JSON Schema encoded as a TOML string. Switchyard adds the provider wrapper. | | `policy` | Yes | — | Policy table. `target_selector` accepts a JSON Pointer such as `/decision/target`. | @@ -231,6 +234,8 @@ policy selector, and routes to any configured target label. | `message_hash_fallback` | No | `false` | Keys affinity on the first user message. Requires `classify_trigger = "new_session"`. | | `recent_turn_window` | No | unset | When unset, the judge sees the opening task and latest user follow-up, when present. When set, it also sees trailing turns. | +The selected JSON label must be `any`, `judge`, `capable`, or `efficient`. + Classifier prompts must not contain `{{RESPONSE_SCHEMA}}`. Switchyard supplies the schema automatically: through the structured-output request in `json_schema` mode, or in the prompt in `json_object` mode. diff --git a/docs/routing_algorithms/llm_classifier_routing.md b/docs/routing_algorithms/llm_classifier_routing.md index 4706175e3..b87560798 100644 --- a/docs/routing_algorithms/llm_classifier_routing.md +++ b/docs/routing_algorithms/llm_classifier_routing.md @@ -139,18 +139,16 @@ packaged `crux`, `primary_rule`, `capability_boundary`, and `p_solve` fields. ## Custom multi-target routing Custom mode accepts an inner JSON Schema and a policy that reads the validated -verdict. This example routes across four configured targets: +verdict. The policy selects one of four runtime categories. ```toml [routes.smart] id = "smart" type = "llm_classifier" mode = "custom" -classifier_target = "classifier" -targets = ["fast", "balanced", "reasoning", "premium"] -default_target = "premium" +default_target = "capable" prompt = """ -Choose the best configured target for this request. +Choose the capable or efficient category for this request. Return JSON matching the response schema supplied with the request. """ response_schema = ''' @@ -162,7 +160,7 @@ response_schema = ''' "properties": { "target": { "type": "string", - "enum": ["fast", "balanced", "reasoning", "premium"] + "enum": ["capable", "efficient"] } }, "required": ["target"], @@ -174,15 +172,22 @@ response_schema = ''' } ''' +[routes.smart.models] +judge = ["classifier"] +capable = ["fast", "premium"] +efficient = ["balanced"] +any = ["fast", "premium", "balanced"] + [routes.smart.policy] type = "target_selector" selector = "/decision/target" ``` -The names in `targets` reference existing target tables. Switchyard passes the +The names in `models` reference existing target tables. Switchyard passes the schema to the provider in a strict structured-output wrapper and validates the -returned JSON again. `jsonptr` resolves the selector against that verdict. A -missing, non-string, or unknown target falls back to `default_target`. +returned JSON again. `jsonptr` resolves the selector against that verdict. The +only valid labels are `any`, `judge`, `capable`, and `efficient`. A missing, +non-string, or unknown label falls back to `default_target`. This separation applies to every classifier mode. Prompts containing the legacy `{{RESPONSE_SCHEMA}}` placeholder are rejected during configuration validation. diff --git a/docs/routing_algorithms/subagent_routing.md b/docs/routing_algorithms/subagent_routing.md index 2533e80ea..27dc4f88b 100644 --- a/docs/routing_algorithms/subagent_routing.md +++ b/docs/routing_algorithms/subagent_routing.md @@ -39,16 +39,15 @@ reasoning = true [routes.agent.subagents] type = "llm_classifier" mode = "custom" -classifier_target = "classifier" -targets = ["worker", "reviewer"] -default_target = "worker" +models = { judge = ["classifier"], capable = ["reviewer"], efficient = ["worker"], any = ["worker", "reviewer"] } +default_target = "efficient" classify_trigger = "new_session" max_output_tokens = 64 prompt = """ Select exactly one target for the delegated task. -- Select "reviewer" for code review, critique, auditing, or correctness analysis. -- Select "worker" for implementation, research, explanation, and other delegated work. +- Select "capable" for code review, critique, auditing, or correctness analysis. +- Select "efficient" for implementation, research, explanation, and other delegated work. Return only JSON matching the response schema. """ @@ -56,7 +55,7 @@ response_schema = ''' { "type": "object", "properties": { - "target": {"type": "string", "enum": ["worker", "reviewer"]} + "target": {"type": "string", "enum": ["capable", "efficient"]} }, "required": ["target"], "additionalProperties": false diff --git a/examples/libsy.py b/examples/libsy.py index fd6b8b684..42da987df 100644 --- a/examples/libsy.py +++ b/examples/libsy.py @@ -58,12 +58,11 @@ async def main() -> None: } client = EchoClient() algorithm = algorithms.random( - ["fast", "quality"], weights=[1, 3], seed=42, ) - async for step in algorithm.run_stream(request): + async for step in algorithm.run_stream(request, {"any": ["fast", "quality"]}): match step: case Step.CallModel(call): call.respond(await client.call(call.request, call.models[0])) diff --git a/examples/litellm/src/switchyard_litellm/plugins/stage_routing_plugin.py b/examples/litellm/src/switchyard_litellm/plugins/stage_routing_plugin.py index a17e23f1b..ac4b6a117 100644 --- a/examples/litellm/src/switchyard_litellm/plugins/stage_routing_plugin.py +++ b/examples/litellm/src/switchyard_litellm/plugins/stage_routing_plugin.py @@ -53,8 +53,6 @@ async def run(self, context: RoutingContext) -> RoutingContext: ) plugin = SwitchyardRoutingPlugin( algorithms.stage_router( - candidates[0], - candidates[1], picker=self._picker, confidence_threshold=self._confidence_threshold, recent_window=self._recent_window, @@ -63,7 +61,12 @@ async def run(self, context: RoutingContext) -> RoutingContext: only_on_wrong_signal_escalation=self._only_on_wrong_signal_escalation, capable_system_prompt=self._capable_system_prompt, efficient_system_prompt=self._efficient_system_prompt, - ) + ), + models={ + "any": candidates, + "capable": [candidates[0]], + "efficient": [candidates[1]], + }, ) return await plugin.run(context) diff --git a/examples/litellm/src/switchyard_litellm/plugins/switchyard_routing_plugin.py b/examples/litellm/src/switchyard_litellm/plugins/switchyard_routing_plugin.py index 333d9a80e..31c2fc1ca 100644 --- a/examples/litellm/src/switchyard_litellm/plugins/switchyard_routing_plugin.py +++ b/examples/litellm/src/switchyard_litellm/plugins/switchyard_routing_plugin.py @@ -159,16 +159,35 @@ class SwitchyardRoutingPlugin(LiteLLMRequestRewriter): selected deployment by the object's LiteLLM callback role. """ - def __init__(self, algorithm: Algorithm) -> None: + def __init__( + self, + algorithm: Algorithm, + models: Mapping[str, Sequence[str]] | None = None, + ) -> None: super().__init__() self._algorithm = algorithm + self._models = ( + {category: list(names) for category, names in models.items()} + if models is not None + else None + ) async def run(self, context: RoutingContext) -> RoutingContext: """Run Switchyard and retain only its selected LiteLLM candidate.""" candidates = list(context.candidate_models) request = _request(context.structured_messages) - async for step in self._algorithm.run_stream(request): + models = ( + self._models + if self._models is not None + else { + "any": candidates, + "judge": candidates, + "capable": candidates, + "efficient": candidates, + } + ) + async for step in self._algorithm.run_stream(request, models): match step: case Step.CallModel(_): raise ValueError( diff --git a/examples/litellm/tests/unit/test_switchyard_routing_plugin.py b/examples/litellm/tests/unit/test_switchyard_routing_plugin.py index e0cca081c..d513778ad 100644 --- a/examples/litellm/tests/unit/test_switchyard_routing_plugin.py +++ b/examples/litellm/tests/unit/test_switchyard_routing_plugin.py @@ -36,13 +36,16 @@ def stage_plugin(**kwargs: object) -> SwitchyardRoutingPlugin: """Build the supported signal-only Stage configuration.""" return SwitchyardRoutingPlugin( algorithms.stage_router( - SOL, - TERRA, picker="efficient_first", confidence_threshold=0.5, recent_window=3, **kwargs, - ) + ), + models={ + "any": [SOL, TERRA], + "capable": [SOL], + "efficient": [TERRA], + }, ) @@ -152,13 +155,18 @@ async def test_litellm_conversion_preserves_stage_tool_signal_input() -> None: direct_outcome = None direct_algorithm = algorithms.stage_router( - SOL, - TERRA, picker="efficient_first", confidence_threshold=0.5, recent_window=3, ) - async for step in direct_algorithm.run_stream(original_request): + async for step in direct_algorithm.run_stream( + original_request, + { + "capable": [SOL], + "efficient": [TERRA], + "any": [SOL, TERRA], + }, + ): match step: case Step.Done(outcome): direct_outcome = outcome @@ -230,19 +238,9 @@ async def test_unsupported_structured_messages_fail_closed( await stage_plugin().run(routing_context(messages)) -async def test_selection_outside_current_litellm_pool_fails_closed() -> None: - plugin = SwitchyardRoutingPlugin(algorithms.random(["openrouter/openai/not-allowed"])) - - with pytest.raises(ValueError, match="not in LiteLLM's candidate pool"): - await plugin.run(routing_context([{"role": "user", "content": "hello"}])) - - async def test_classifier_backed_algorithm_fails_on_intermediate_model_call() -> None: plugin = SwitchyardRoutingPlugin( algorithms.llm_task_classifier( - "openrouter/openai/gpt-5.6-judge", - TERRA, - SOL, config=TaskClassifierConfig(0.5), ) ) @@ -434,7 +432,12 @@ def test_request_patch_rejects_unsafe_or_unrepresentable_overrides( class EmptyAlgorithm: """Algorithm-shaped test double whose stream violates the terminal-step contract.""" - async def run_stream(self, request: dict[str, object]) -> AsyncIterator[object]: + async def run_stream( + self, + request: dict[str, object], + models: dict[str, list[str]], + ) -> AsyncIterator[object]: + del models if request: return yield object() diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index 9756b1c6b..4346b0dd4 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -60,7 +60,7 @@ def __init__(self, stream: AsyncIterator[Mapping[str, object]]) -> None: ... @final class CustomClassifierConfig: - """Configure schema-validated routing across named targets. + """Configure schema-validated routing across runtime model categories. ``max_output_tokens`` must be positive. Enabling ``message_hash_fallback`` requires ``session_affinity``. @@ -180,9 +180,6 @@ class LlmClassifierConfig: @staticmethod def capability( - judge_target: str, - efficient_target: str, - capable_target: str, *, config: TaskClassifierConfig, ) -> LlmClassifierConfig: @@ -191,9 +188,6 @@ def capability( @staticmethod def escalation( - judge_target: str, - efficient_target: str, - capable_target: str, *, config: EscalationClassifierConfig, ) -> LlmClassifierConfig: @@ -202,20 +196,17 @@ def escalation( @staticmethod def custom( - judge_target: str, - targets: Sequence[tuple[str, str]], *, default_target: str, config: CustomClassifierConfig, ) -> LlmClassifierConfig: - """Route among named targets using a schema-selected label.""" + """Route among runtime categories using a schema-selected label.""" ... @final class LlmFallback: def __init__( self, - judge_target: str, *, config: TaskClassifierConfig, ) -> None: ... @@ -225,14 +216,13 @@ class Algorithm: def run_stream( self, request: Mapping[str, object], + models: Mapping[str, Sequence[str]], headers: Mapping[str, str] | None = None, ) -> AsyncIterator[Step.CallModel | Step.Done]: ... def noop() -> Algorithm: ... def random( - targets: Sequence[str], - *, weights: Sequence[float] | None = None, seed: int | None = None, ) -> Algorithm: ... @@ -242,16 +232,11 @@ def llm_classifier(config: LlmClassifierConfig) -> Algorithm: ... def llm_task_classifier( - judge_target: str, - efficient_target: str, - capable_target: str, *, config: TaskClassifierConfig, ) -> Algorithm: ... def stage_router( - capable_target: str, - efficient_target: str, *, picker: str, confidence_threshold: float, diff --git a/tests/test_libsy_minimal_bindings.py b/tests/test_libsy_minimal_bindings.py index b693d4d87..a83263acb 100644 --- a/tests/test_libsy_minimal_bindings.py +++ b/tests/test_libsy_minimal_bindings.py @@ -58,10 +58,14 @@ async def run_algorithm( algorithm: Algorithm, clients: dict[str, Any] | None = None, *, + models: dict[str, list[str]] | None = None, request: dict[str, Any] | None = None, headers: dict[str, str] | None = None, ) -> tuple[str, dict[str, Any]]: - async for step in algorithm.run_stream(request or request_body(), headers=headers): + runtime_models = models if models is not None else {"any": list((clients or {}).keys())} + async for step in algorithm.run_stream( + request or request_body(), runtime_models, headers=headers + ): match step: case Step.CallModel(call): for index, target in enumerate(call.models): @@ -103,11 +107,12 @@ 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() outcome: RoutingOutcome | None = None variants: list[str] = [] - async for step in algorithm.run_stream(request_body()): + models = {"any": ["fast"]} + async for step in algorithm.run_stream(request_body(), models): match step: case Step.Done(done): variants.append("done") @@ -123,9 +128,7 @@ async def test_random_streams_complex_steps_and_accepts_a_dictionary_response() assert outcome.metadata.evidence is None response = await client.call(outcome.request) assert client.calls[0]["model"] == "fast" - assert client.calls[0]["messages"][0]["content"] == [ - {"type": "text", "text": "hello"} - ] + assert client.calls[0]["messages"][0]["content"] == [{"type": "text", "text": "hello"}] assert response["model"] == "fast" assert response["outputs"][0]["content"] == [{"type": "text", "text": "fast"}] @@ -134,7 +137,7 @@ async def test_routing_call_accepts_a_streamed_response() -> None: async def events() -> AsyncIterator[dict[str, object]]: for chunk in [ {"MessageStart": {"id": "response-1", "model": "judge"}}, - {"TextDelta": {"index": 0, "text": '{"target":"balanced"}'}}, + {"TextDelta": {"index": 0, "text": '{"target":"efficient"}'}}, {"MessageStop": {"reason": "end_turn"}}, ]: yield {"preservation": None, "normalized": [chunk]} @@ -143,19 +146,23 @@ async def events() -> AsyncIterator[dict[str, object]]: "type": "object", "additionalProperties": False, "required": ["target"], - "properties": {"target": {"type": "string", "enum": ["fast", "balanced"]}}, + "properties": {"target": {"type": "string", "enum": ["capable", "efficient"]}}, } algorithm = algorithms.llm_classifier( LlmClassifierConfig.custom( - "judge", - [("fast", "model-a"), ("balanced", "model-b")], - default_target="fast", + default_target="capable", config=CustomClassifierConfig("Choose a target.", schema, "/target"), ) ) outcome: RoutingOutcome | None = None - async for step in algorithm.run_stream(request_body()): + models = { + "judge": ["judge"], + "capable": ["model-a"], + "efficient": ["model-b"], + "any": ["model-a", "model-b"], + } + async for step in algorithm.run_stream(request_body(), models): match step: case Step.CallModel(call): call.respond(LlmResponse.Stream(events())) @@ -195,9 +202,6 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: weak = EchoClient("weak") algorithm = algorithms.llm_classifier( LlmClassifierConfig.capability( - "judge", - "weak", - "strong", config=TaskClassifierConfig( 0.5, threshold_step=0.1, @@ -207,7 +211,13 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: ) outcome: RoutingOutcome | None = None - async for step in algorithm.run_stream(request_body()): + models={ + "judge": ["judge"], + "efficient": ["weak"], + "capable": ["strong"], + "any": ["weak", "strong"], + } + async for step in algorithm.run_stream(request_body(), models): match step: case Step.CallModel(call): call.respond(LlmResponse.Agg(await judge.call(call.request))) @@ -226,9 +236,9 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: prompt = judge.calls[0]["instructions"][0]["content"][0]["text"] assert prompt == "Custom capability rubric." - assert judge.calls[0]["output"]["response_format"]["json_schema"]["schema"][ - "properties" - ]["p_solve"] + assert judge.calls[0]["output"]["response_format"]["json_schema"]["schema"]["properties"][ + "p_solve" + ] assert response["model"] == "weak" @@ -241,7 +251,7 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: "outputs": [ { "role": "assistant", - "content": [{"type": "text", "text": '{"target":"balanced"}'}], + "content": [{"type": "text", "text": '{"target":"efficient"}'}], "stop_reason": "end_turn", } ], @@ -251,13 +261,11 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: "type": "object", "additionalProperties": False, "required": ["target"], - "properties": {"target": {"type": "string", "enum": ["fast", "balanced", "best"]}}, + "properties": {"target": {"type": "string", "enum": ["capable", "efficient"]}}, } algorithm = algorithms.llm_classifier( LlmClassifierConfig.custom( - "judge", - [("fast", "model-a"), ("balanced", "model-b"), ("best", "model-c")], - default_target="fast", + default_target="capable", config=CustomClassifierConfig("Choose a target.", schema, "/target"), ) ) @@ -269,6 +277,12 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: "model-a": EchoClient("model-a"), "model-b": EchoClient("model-b"), "model-c": EchoClient("model-c"), + }, + models={ + "judge": ["judge"], + "capable": ["model-a", "model-c"], + "efficient": ["model-b"], + "any": ["model-a", "model-b", "model-c"], }, ) @@ -303,9 +317,6 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: judge = JudgeClient("judge") weak = EchoClient("weak") algorithm = algorithms.llm_task_classifier( - "judge", - "weak", - "strong", config=TaskClassifierConfig(0.5, response_format_type="json_object"), ) @@ -316,6 +327,12 @@ async def call(self, request: dict[str, Any]) -> dict[str, Any]: "weak": weak, "strong": EchoClient("strong"), }, + models={ + "judge": ["judge"], + "efficient": ["weak"], + "capable": ["strong"], + "any": ["weak", "strong"], + }, ) assert judge.calls[0]["output"]["response_format"] == {"type": "json_object"} @@ -338,7 +355,6 @@ def test_classifier_config_rejects_unknown_response_format() -> None: async def test_random_weights_and_seed_are_reproducible() -> None: def algorithm(): return algorithms.random( - ["fast", "capable"], weights=[1, 3], seed=42, ) @@ -354,10 +370,8 @@ def algorithm(): def test_random_rejects_invalid_weights() -> None: - targets = ["fast", "capable"] - - with pytest.raises(ValueError, match="expected 2 weights, got 1"): - algorithms.random(targets, weights=[1]) + with pytest.raises(ValueError, match="finite and nonnegative"): + algorithms.random(weights=[-1]) async def test_noop_needs_no_client() -> None: @@ -376,7 +390,7 @@ async def test_noop_needs_no_client() -> None: ) def test_algorithm_rejects_invalid_headers(headers: dict[str, str], message: str) -> None: with pytest.raises(ValueError, match=message): - algorithms.noop().run_stream(request_body(), headers=headers) + algorithms.noop().run_stream(request_body(), {}, headers=headers) async def test_algorithm_accepts_case_insensitive_duplicate_names() -> None: @@ -391,7 +405,7 @@ def test_algorithm_rejects_header_map_capacity_overflow() -> None: headers = {f"x-header-{index}": "value" for index in range(32_769)} with pytest.raises(ValueError, match="max size reached"): - algorithms.noop().run_stream(request_body(), headers=headers) + algorithms.noop().run_stream(request_body(), {}, headers=headers) def test_algorithm_exposes_only_streaming_execution() -> None: @@ -401,38 +415,49 @@ def test_algorithm_exposes_only_streaming_execution() -> None: assert not hasattr(algorithm, "run") -def test_random_requires_a_target() -> None: - with pytest.raises(ValueError, match="at least one target"): - algorithms.random([]) - - def test_invalid_request_is_rejected_at_the_boundary() -> None: - algorithm = algorithms.random(["fast"]) + algorithm = algorithms.random() with pytest.raises(ValueError, match="unknown variant"): algorithm.run_stream( { "model": "auto", "messages": [{"role": "invalid", "content": []}], - } + }, + {"any": ["fast"]}, ) async def test_context_window_failure_falls_back_to_the_next_model() -> None: class OverflowClient: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + async def call(self, request: dict[str, Any]) -> dict[str, Any]: + self.calls.append(request) raise ContextWindowExceededError("request exceeds context window") + overflow = OverflowClient() algorithm = algorithms.stage_router( - "strong", - "fast", picker="efficient_first", confidence_threshold=0.5, + efficient_system_prompt="Use the efficient tier.", ) selected_model, response = await run_algorithm( algorithm, - {"fast": OverflowClient(), "strong": EchoClient("strong")}, + {"fast": overflow, "strong": EchoClient("strong")}, + models={ + "efficient": ["fast"], + "capable": ["strong"], + "any": ["fast", "strong"], + }, ) assert selected_model == "fast" assert response["model"] == "strong" + assert overflow.calls[0]["instructions"] == [ + { + "role": "system", + "content": [{"type": "text", "text": "Use the efficient tier."}], + } + ]