diff --git a/README.md b/README.md index fa972f4cc..90c482e7a 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, ) @@ -205,8 +203,15 @@ async def call_with_fallback(request: dict, models: list[str], clients: dict) -> raise error or RuntimeError("no candidate models") +runtime_models = { + "efficient": ["fast"], + "capable": ["quality"], + "any": ["quality", "fast"], +} + + async def route(request: dict, clients: dict) -> LlmResponse.Agg | LlmResponse.Stream: - async for step in algorithm.run_stream(request): + async for step in algorithm.run_stream(request, runtime_models): match step: case Step.CallModel(call): try: @@ -222,7 +227,8 @@ async def route(request: dict, clients: dict) -> LlmResponse.Agg | LlmResponse.S raise RuntimeError("algorithm ended without a decision") ``` -`clients` maps each target name to your existing client; each `call` takes a +`runtime_models` groups the model IDs available for this request by category. +`clients` maps each model ID to your existing client; each `call` takes a normalized request dict and returns a normalized response dict. `call.models` and `outcome.selected_model_ids` list candidates in order, so the helper tries each one before giving up. `outcome.request` is the request to send, which may 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..191d3c5b8 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,20 +73,20 @@ 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 +ROUTING BIAS: the efficient tier is the DEFAULT — it handles nearly all support work end-to-end: lookups, standard actions, troubleshooting with known steps, collecting info, confirmations, relaying tool results, single-policy checks. -ROUTE TO WEAK UNLESS YOU CAN NAME A SPECIFIC REASON IT WILL FAIL. +ROUTE TO efficient UNLESS YOU CAN NAME A SPECIFIC REASON IT WILL FAIL. -Escalate to STRONG only for a CONCRETE hard property: reconciling multiple +Escalate to capable only for a CONCRETE hard property: reconciling multiple conflicting policy conditions, multi-account/line arithmetic, diagnosis still unexplained after tool checks, or an exception decision where a wrong answer -harms the customer. When in doubt, choose WEAK. +harms the customer. When in doubt, choose efficient. Set abstain=true only when truly unclassifiable. No markdown, no chain-of-thought — only the JSON object. 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 9a98ea614..508177cb1 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -19,7 +19,9 @@ use std::time::Instant; use http::StatusCode; use parking_lot::Mutex; -use switchyard_libsy::{Algorithm, CallModel, LibsyError, Result, RoutingOutcome, drive}; +use switchyard_libsy::{ + Algorithm, CallModel, LibsyError, Result, RoutingOutcome, RuntimeModels, drive, +}; use switchyard_protocol::{ LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason, }; @@ -44,6 +46,7 @@ pub async fn run( algorithm: Arc, clients: ClientRouter, request: Request, + models: Arc, observer: Option, ) -> Result<(ModelId, Response)> { let algorithm_name = algorithm.name().to_string(); @@ -52,7 +55,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 +113,10 @@ pub async fn decide( algorithm: Arc, clients: ClientRouter, request: Request, + models: Arc, ) -> 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?; @@ -452,10 +456,9 @@ mod tests { use wiremock::{Mock, MockServer, ResponseTemplate}; use crate::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; + use switchyard_protocol::Category; - struct CandidateAlgorithm { - models: Vec, - } + struct CandidateAlgorithm {} struct AnsweredAlgorithm { model: ModelId, @@ -469,13 +472,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 +607,29 @@ 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]) -> Arc { + Arc::new(RuntimeModels::new( + [( + 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 +647,7 @@ mod tests { }), ClientRouter::single(client.clone()), request(), + Arc::new(RuntimeModels::default()), Some(observer), ) .await?; @@ -697,11 +712,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 +747,7 @@ mod tests { }), clients, request(), + Arc::new(RuntimeModels::default()), ) .await?; @@ -763,11 +778,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?; @@ -905,10 +919,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(()) @@ -997,9 +1016,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 { @@ -1007,7 +1024,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..fd48f790e 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -34,12 +34,12 @@ use tracing_subscriber::layer::{Context as LayerContext, SubscriberExt}; use tracing_subscriber::registry::LookupSpan; use switchyard_libsy::{ - AffinityRouter, Algorithm, Classifier, ClassifyTrigger, Driver, LibsyError, - LlmClassifierConfig, LlmTaskClassifier, PickerMode, RoutingOutcome, StageRouter, - StageRouterConfig, Step, TaskClassifierConfig, + Algorithm, ClassifyTrigger, Driver, LibsyError, LlmClassifierConfig, LlmTaskClassifier, + PickerMode, RoutingOutcome, RuntimeModels, StageRouter, 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, + Arc::new(RuntimeModels::default()), + 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,42 @@ fn classifier_router( )?)) } +fn classifier_models( + judge_model: &str, + efficient_model: &str, + capable_model: &str, +) -> Arc { + Arc::new(RuntimeModels::new( + [ + (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"), @@ -666,9 +702,15 @@ async fn affinity_warns_once_when_request_has_no_usable_identity() -> switchyard let _guard = serialize_test().lock().await; let (store, _, _, _, _) = telemetry(); let event_count = store.events().len(); - let router = AffinityRouter::new().with_message_hash_fallback(); - let mut state = (); - let mut request = Request { + let router: Arc = + Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability { + config: TaskClassifierConfig { + classify_trigger: ClassifyTrigger::NewSession, + message_hash_fallback: true, + ..TaskClassifierConfig::default() + }, + })?); + let request = Request { llm_request: LlmRequest { messages: vec![Message { role: Role::User, @@ -685,7 +727,11 @@ async fn affinity_warns_once_when_request_has_no_usable_identity() -> switchyard }; for _ in 0..2 { - router.score(&mut state, &mut request, None).await?; + let mut steps = router.clone().run_stream( + request.clone(), + classifier_models("warning/judge", "warning/efficient", "warning/capable"), + ); + let _ = steps.next().await; } let events = store.events(); @@ -716,9 +762,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 +774,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 +789,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 +1099,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 +1140,24 @@ 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, + Arc::new(RuntimeModels::new( + [ + (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 +1198,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"), + Arc::new(RuntimeModels::default()), Some(observer), ) .await?; @@ -1299,8 +1372,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 +1408,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"), + Arc::new(RuntimeModels::default()), + ); tokio::pin!(stream); let mut saw_error_step = false; @@ -1420,9 +1498,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 +1598,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 +1643,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"), + Arc::new(RuntimeModels::default()), + ); tokio::pin!(stream); let attributes = [("algorithm", ALGO)]; @@ -1619,7 +1701,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"), + Arc::new(RuntimeModels::default()), + ); 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/advisor_gate.rs b/crates/libsy/src/algorithms/advisor_gate.rs index 72fc2a230..da4062ad0 100644 --- a/crates/libsy/src/algorithms/advisor_gate.rs +++ b/crates/libsy/src/algorithms/advisor_gate.rs @@ -34,8 +34,8 @@ use std::sync::Arc; use std::time::Instant; use switchyard_protocol::{ - ContentBlock, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Request, Role, - SamplingParams, + Category, ContentBlock, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Request, + Role, SamplingParams, }; use crate::core::algorithm::{Algorithm, Driver, RoutingOutcome}; @@ -153,8 +153,6 @@ impl Default for AdvisorGateConfig { /// turn, which a stronger advisor reviews once per scope budget (APPROVE /// releases it, REDO feeds the plan back and re-invokes the executor). pub struct AdvisorGate { - executor: ModelId, - advisor: ModelId, config: AdvisorGateConfig, /// Folds request- and response-side facts into the per-turn [`GateSignals`]. signals: GateSignalProcessor, @@ -166,8 +164,8 @@ pub struct AdvisorGate { } impl AdvisorGate { - /// Validates ranges and compiles the trigger and verdict patterns. - pub fn new(executor: ModelId, advisor: ModelId, config: AdvisorGateConfig) -> Result { + /// Validates the config. Models are supplied when each request runs. + pub fn new(config: AdvisorGateConfig) -> Result { if config.max_reviews < 1 { return Err(algorithm_error("max_reviews must be at least 1")); } @@ -183,8 +181,6 @@ impl AdvisorGate { })?; let budget = ReviewBudget::new(config.max_reviews); Ok(Self { - executor, - advisor, config, signals: GateSignalProcessor, trigger, @@ -201,14 +197,21 @@ impl AdvisorGate { request: Request, scope: &ScopeKey, ) -> Result { + let executor_models = driver.models_for(&Category::Efficient).to_vec(); + let executor = executor_models + .first() + .ok_or_else(|| LibsyError::AlgorithmError { + message: "no models available for category Efficient".to_string(), + })?; + // Spent budget (or failure cap): pure passthrough — live stream, // verbatim preserved-body replay, zero buffering. Executor errors // (including ContextWindowExceeded) propagate for the host's // client-visible mapping. if self.budget.check_exhausted(scope) { return Ok(RoutingOutcome::route_to( - self.executor.clone(), - Vec::new(), + executor.clone(), + executor_models[1..].to_vec(), request, )); } @@ -221,7 +224,7 @@ impl AdvisorGate { &mut signals, Event::Request { request: &mut request, - driver: Some(driver), + driver, }, ) .await?; @@ -229,9 +232,13 @@ impl AdvisorGate { // Gated phase: generate the turn once, fully buffered, so the gate // can inspect it before the client sees anything. let response = driver - .call_model(request.clone(), vec![self.executor.clone()]) + .call_model(request.clone(), executor_models.clone()) .await?; - let turn = buffer_turn(self.executor.as_str(), response).await?; + let served_executor = response + .served_model() + .cloned() + .unwrap_or_else(|| executor.clone()); + let turn = buffer_turn(served_executor.as_str(), response).await?; // Response-side signals fold in after it: the terminal turn never // appears on a later request, so the trigger runs on this event. @@ -249,14 +256,14 @@ impl AdvisorGate { && self.budget.try_mark_stall_fired(stall_key(&request)); if decision.fired.is_none() && !stall { return Ok(RoutingOutcome::answered( - self.executor.clone(), + served_executor.clone(), request, turn.into_response(), )); } if !self.budget.try_reserve(scope) { return Ok(RoutingOutcome::answered( - self.executor.clone(), + served_executor.clone(), request, turn.into_response(), )); @@ -277,7 +284,7 @@ impl AdvisorGate { "trigger": trigger_label, })); Ok(RoutingOutcome::answered( - self.executor.clone(), + served_executor.clone(), request, turn.into_response(), )) @@ -288,7 +295,14 @@ impl AdvisorGate { "verdict": "redo", "trigger": trigger_label, })); - Ok(self.redo(request, turn, &plan)) + Ok(self.redo( + executor, + &executor_models[1..], + &served_executor, + request, + turn, + &plan, + )) } Ok(ConsultOutcome::Failed { reason }) => { self.budget.refund_failure(scope); @@ -299,7 +313,7 @@ impl AdvisorGate { "reason_code": reason, })); Ok(RoutingOutcome::answered( - self.executor.clone(), + served_executor, request, turn.into_response(), )) @@ -314,9 +328,17 @@ impl AdvisorGate { /// REDO: the client never sees the gated turn. Its text (or reasoning) is /// echoed as an assistant message, the advisor's plan follows as user /// feedback, and the executor continues as a pure passthrough call. - fn redo(&self, request: Request, turn: GatedTurn, plan: &str) -> RoutingOutcome { + fn redo( + &self, + executor: &ModelId, + executor_fallbacks: &[ModelId], + served_executor: &ModelId, + request: Request, + turn: GatedTurn, + plan: &str, + ) -> RoutingOutcome { record_discarded(&turn.agg.usage); - emit_discarded_audit(self.executor.as_str(), &turn.agg.usage); + emit_discarded_audit(served_executor.as_str(), &turn.agg.usage); let echo = visible_text(&turn.agg) .or_else(|| reasoning_text(&turn.agg)) .unwrap_or_else(|| EMPTY_ECHO_PLACEHOLDER.to_string()); @@ -332,7 +354,7 @@ impl AdvisorGate { // preserved pre-surgery body verbatim and the feedback never reaches // the executor. crate::algorithms::util::prompts::drop_exact_replay(&mut redo); - RoutingOutcome::route_to(self.executor.clone(), Vec::new(), redo) + RoutingOutcome::route_to(executor.clone(), executor_fallbacks.to_vec(), redo) } /// Consults the advisor over the buffered transcript and parses the @@ -366,15 +388,27 @@ impl AdvisorGate { ); let consult_request = self.build_consult_request(base, transcript); let started = Instant::now(); - let reply = match driver - .call_model(consult_request, vec![self.advisor.clone()]) - .await - { - Ok(response) => response - .llm_response - .into_agg() - .await - .map_err(|source| LibsyError::client_call(self.advisor.clone(), source)), + // An unresolvable advisor is treated like any other consult failure, so + // fail_open still returns the buffered executor turn to the client. + let reply = match driver.first_model_for(&Category::Judge) { + Ok(advisor) => { + let advisor = advisor.clone(); + let advisor_models = driver.models_for(&Category::Judge).to_vec(); + match driver.call_model(consult_request, advisor_models).await { + Ok(response) => { + let served_advisor = response + .served_model() + .cloned() + .unwrap_or_else(|| advisor.clone()); + response + .llm_response + .into_agg() + .await + .map_err(|source| LibsyError::client_call(served_advisor, source)) + } + Err(error) => Err(error), + } + } Err(error) => Err(error), }; let latency_ms = started.elapsed().as_secs_f64() * 1000.0; diff --git a/crates/libsy/src/algorithms/advisor_gate/signals.rs b/crates/libsy/src/algorithms/advisor_gate/signals.rs index 3b44f19bc..95f7549d1 100644 --- a/crates/libsy/src/algorithms/advisor_gate/signals.rs +++ b/crates/libsy/src/algorithms/advisor_gate/signals.rs @@ -56,6 +56,7 @@ impl Processor for GateSignalProcessor { #[cfg(test)] mod tests { use super::*; + use crate::core::testing::empty_driver; use switchyard_protocol::{ AggLlmResponse, ContentBlock, LlmRequest, Message, ModelId, Request, ResponseOutput, Role, ToolCall, ToolResult, @@ -116,7 +117,7 @@ mod tests { &mut state, Event::Request { request: &mut request, - driver: None, + driver: &empty_driver(), }, ) .await?; @@ -153,6 +154,8 @@ mod tests { Event::Decision { request: &mut request, selected_model_id: &selected, + category: None, + driver: &empty_driver(), }, ) .await?; diff --git a/crates/libsy/src/algorithms/advisor_gate/tests.rs b/crates/libsy/src/algorithms/advisor_gate/tests.rs index 464396f0e..5d0fe26c9 100644 --- a/crates/libsy/src/algorithms/advisor_gate/tests.rs +++ b/crates/libsy/src/algorithms/advisor_gate/tests.rs @@ -3,6 +3,7 @@ //! Behavior tests for the advisor review gate. +use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use switchyard_protocol::{ResponseOutput, ToolCall, ToolResult, completion_text}; @@ -15,19 +16,36 @@ use switchyard_protocol::{ use super::transcript::{NO_TEXT_PLACEHOLDER, TRUNCATION_MARKER, middle_drop}; use super::*; -use crate::core::testing::{reply, test_drive}; +use crate::RuntimeModels; +use crate::core::testing::{Serve, reply, test_drive_with_models}; const EXECUTOR: &str = "executor"; const ADVISOR: &str = "advisor"; +const EXECUTOR_FALLBACK: &str = "executor-fallback"; +const ADVISOR_FALLBACK: &str = "advisor-fallback"; fn target(name: &str) -> ModelId { ModelId::new(name) } fn gate(config: AdvisorGateConfig) -> Arc { - Arc::new( - AdvisorGate::new(target(EXECUTOR), target(ADVISOR), config).expect("test config is valid"), - ) + Arc::new(AdvisorGate::new(config).expect("test config is valid")) +} + +fn runtime_models() -> HashMap> { + [ + (Category::Efficient, vec![target(EXECUTOR)]), + (Category::Judge, vec![target(ADVISOR)]), + ] + .into() +} + +async fn test_drive( + algorithm: Arc, + request: Request, + serve: impl Serve, +) -> Result<(ModelId, Response)> { + test_drive_with_models(algorithm, request, runtime_models(), serve).await } fn request(messages: Vec) -> Request { @@ -281,6 +299,59 @@ async fn approved_terminal_turn_returns_buffered_body() { assert_eq!(selected_model, EXECUTOR); } +#[tokio::test] +async fn calls_preserve_candidates_and_attribute_the_serving_executor() { + let gate = gate(AdvisorGateConfig::default()); + let models = RuntimeModels::new( + [ + ( + Category::Efficient, + vec![target(EXECUTOR), target(EXECUTOR_FALLBACK)], + ), + ( + Category::Judge, + vec![target(ADVISOR), target(ADVISOR_FALLBACK)], + ), + ] + .into(), + ); + let calls = Arc::new(parking_lot::Mutex::new(Vec::new())); + let observed = Arc::clone(&calls); + let outcome = crate::drive(gate, task_request(), Arc::new(models), move |call| { + let observed = Arc::clone(&observed); + async move { + let candidates = call.models.clone(); + observed.lock().push(candidates.clone()); + let (text, served) = if candidates[0] == target(EXECUTOR) { + ("all done", target(EXECUTOR_FALLBACK)) + } else { + ("APPROVE", target(ADVISOR_FALLBACK)) + }; + let mut response = reply(text); + response.set_served_model(&served); + call.respond(Ok(response)) + } + }) + .await + .expect("routes"); + + assert_eq!( + *calls.lock(), + vec![ + vec![target(EXECUTOR), target(EXECUTOR_FALLBACK)], + vec![target(ADVISOR), target(ADVISOR_FALLBACK)], + ] + ); + assert_eq!( + outcome.selected_model_id().expect("selected model"), + &target(EXECUTOR_FALLBACK) + ); + assert_eq!( + outcome.request.model_id().as_deref(), + Some(EXECUTOR_FALLBACK) + ); +} + #[tokio::test] async fn redo_appends_echo_and_feedback_then_reinvokes() { let script = Script::new(); @@ -1153,9 +1224,7 @@ fn transcript_middle_drop() { #[test] fn new_validation_errors() { let invalid = |config: AdvisorGateConfig, needle: &str| { - let error = AdvisorGate::new(target(EXECUTOR), target(ADVISOR), config) - .err() - .expect("config rejected"); + let error = AdvisorGate::new(config).err().expect("config rejected"); assert!(error.to_string().contains(needle), "{error}"); }; invalid( diff --git a/crates/libsy/src/algorithms/composite.rs b/crates/libsy/src/algorithms/composite.rs index af75f7b64..3afe61486 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, Request}; const COMPOSITE: &str = "composite"; @@ -33,13 +33,20 @@ 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(category: Option) -> Option { + match category { + Some(Category::Capable) => Some(Tier::Capable), + Some(Category::Efficient) => Some(Tier::Efficient), + _ => 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 +83,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(winner.category) { set_fall_open(state, tier); if let Some(identity) = identity { @@ -90,9 +97,7 @@ impl Processor for TierSetter { if let Some(tier) = identity.and_then(|identity| self.tiers.lock().get(&identity).copied()) { set_fall_open(state, tier); - if let Some(driver) = driver { - driver.set_evidence_if_empty(serde_json::json!({"source": "retained"})); - } + driver.set_evidence_if_empty(serde_json::json!({"source": "retained"})); } Ok(()) } @@ -100,8 +105,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 +123,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 +140,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 +172,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, ModelId, 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,66 +213,109 @@ 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 { .. }) )); } + /// A route may point its judge at the same target it serves capable turns on. + /// The tier then cannot be recovered by looking the served model up in the + /// runtime groups — one id, two categories — so it comes off the verdict itself. + #[tokio::test] + async fn a_judge_sharing_the_capable_model_still_latches_the_tier() -> Result<()> { + let models: HashMap> = [ + (Category::Judge, vec![ModelId::from("strong")]), + (Category::Capable, vec![ModelId::from("strong")]), + (Category::Efficient, vec![ModelId::from("weak")]), + ( + Category::Any, + vec![ModelId::from("strong"), ModelId::from("weak")], + ), + ] + .into(); + // The judge runs first as a request-side processor, so the opening call is its own. + let calls = Arc::new(Mutex::new(0u32)); + let serve = { + let calls = Arc::clone(&calls); + move |target: ModelId, _request: Request| { + let calls = Arc::clone(&calls); + async move { + let mut calls = calls.lock(); + *calls += 1; + let completion = if *calls == 1 { + // p_solve below the threshold: the judge does not trust the + // efficient tier, so the verdict is capable. + r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.1}"#.to_string() + } else { + target.to_string() + }; + Ok(crate::core::testing::reply(completion)) + } + } + }; + let router = router()?; + + test_drive_with_models( + router.clone(), + user_turn_request(), + models.clone(), + serve.clone(), + ) + .await?; + // A tool step is not a user turn, so nothing re-judges and the latched tier decides. + let (selected, _) = + test_drive_with_models(router, turn_request(false), models, serve).await?; + + assert_eq!(selected, ModelId::from("strong")); + Ok(()) + } + #[tokio::test] async fn a_session_without_an_id_keys_on_the_message_hash() -> Result<()> { let recorder = Arc::new(Recorder::default()); *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 +339,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..825e32428 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) @@ -82,13 +68,10 @@ impl Classifier for EscalationClassifier { &self, state: &mut State, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { - let Some(driver) = driver else { - return Err(LibsyError::AlgorithmError { - 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 +79,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 +87,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 +103,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 +116,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. @@ -154,15 +137,12 @@ impl Classifier for EscalationClassifier { metadata: efficient_response.metadata, }; - let (classification, _) = self - .judge - .score(state, &mut judge_request, Some(driver)) - .await?; + let (classification, _) = self.judge.score(state, &mut judge_request, driver).await?; 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 +156,7 @@ impl Classifier for EscalationClassifier { "source": "escalation", "verdict": "escalate", })); - return Ok((decisive(&self.capable), None)); + return Ok((decisive(&capable), None)); } if escalate { @@ -186,25 +166,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 +239,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 +271,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 +286,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 +325,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 +333,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 +344,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 +367,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 +394,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 +436,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 +462,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..8d6f5e71d 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -13,8 +13,9 @@ //! private state value across turns with the same session ID. Requests without a session ID use //! unretained per-run state. //! -//! The selected target is offered first, followed by every other configured target. The consumer -//! may fall through that ordered candidate list when a model call fails. +//! The selected target is offered first, then the rest of the category it was drawn from, then +//! every other runtime target. The consumer may fall through that ordered candidate list when a +//! model call fails. use std::{ collections::HashMap, @@ -27,10 +28,10 @@ use parking_lot::Mutex; use tokio::sync::Mutex as AsyncMutex; use crate::core::algorithm::{self, Algorithm, Driver}; -use crate::core::classifier::{Classification, Classifier, Score}; +use crate::core::classifier::{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>, @@ -48,47 +49,6 @@ const SESSION_STATE_TTL: Duration = Duration::from_secs(60 * 60); /// Run the expired session cleanup code this often. const SESSION_CLEANUP_INTERVAL: Duration = Duration::from_secs(60 * 60); -/// Terminal classifier for a cascade whose classifiers may all abstain. -/// -/// A classifier abstains when it cannot decide, which lets the next one try. The -/// last has no next, so a cascade that could abstain all the way through needs a -/// decider that never does. Which target that is belongs to whoever assembles the -/// cascade, not to the classifiers in it. -pub struct DefaultTarget { - target: ModelId, -} - -impl DefaultTarget { - /// Close a cascade with `target`. - pub fn new(target: impl Into) -> Self { - Self { - target: target.into(), - } - } -} - -#[async_trait] -impl Classifier for DefaultTarget { - async fn score( - &self, - _state: &mut S, - _request: &mut Request, - driver: Option<&Driver>, - ) -> Result<(Classification, Option)> { - // Zero confidence: this is a fallback, not a judgement. - if let Some(driver) = driver { - driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"})); - } - Ok(( - Classification::Scores(vec![Score { - target: self.target.clone(), - confidence: 0.0, - }]), - None, - )) - } -} - /// Processor chain → classifier cascade → routed model call. See the module docs. /// /// The generic state type is shared by every processor and classifier in the composition. @@ -96,19 +56,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() -> Self { Self { name: "fall_through".to_string(), processors: Vec::new(), classifiers: Vec::new(), - targets, session_states: None, cleanup_started: Once::new(), } @@ -120,12 +78,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() -> 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(), } @@ -180,7 +137,7 @@ where // it, later components see the rewrite, and the final value reaches the model. let mut request = request; let session_state = self.session_state(&request); - let (target, served) = match session_state { + let (score, served) = match session_state { Some(state) => { let mut state = state.lock().await; self.route(&mut state, &driver, &mut request).await? @@ -196,10 +153,21 @@ where // for twice. // Nothing reads it on the way out: streamed or buffered, it reaches the caller // untouched. + let target = score.target; match served { Some(response) => Ok(RoutingOutcome::answered(target, request, response)), None => { - let fallback_models = self.fallbacks(&target); + // The rest of the category the decision came from leads: those are the + // alternatives to the model the classifier picked, in the order it ranked + // them. Every other runtime target follows as a last resort, so a route + // whose category holds one model still fails over to the other tier. + let chosen = driver.models_for(score.category.as_ref().unwrap_or(&Category::Any)); + let mut fallback_models: Vec = Vec::new(); + for candidate in chosen.iter().chain(driver.models_for(&Category::Any)) { + if *candidate != target && !fallback_models.contains(candidate) { + fallback_models.push(candidate.clone()); + } + } Ok(RoutingOutcome::route_to(target, fallback_models, request)) } } @@ -212,15 +180,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()?; @@ -240,14 +199,11 @@ where state: &mut S, driver: &Driver, request: &mut Request, - ) -> Result<(ModelId, Option)> { + ) -> Result<(Score, Option)> { // 1. Processor chain accumulates request-side facts into the composition's state. // The driver is offered here so a processor may consult a model first. for processor in &self.processors { - let event = Event::Request { - request, - driver: Some(driver), - }; + let event = Event::Request { request, driver }; processor.process(state, event).await?; } @@ -255,7 +211,7 @@ where // per-request driver is offered to each — driver-backed classifiers use it. let mut routed = None; for classifier in &self.classifiers { - let (scores, response) = classifier.score(state, request, Some(driver)).await?; + let (scores, response) = classifier.score(state, request, driver).await?; if let Some(score) = scores.argmax(false)? { // Only the deciding classifier's response answers the turn; an abstaining // classifier selected nothing for it to be the answer to. @@ -270,7 +226,7 @@ where }; // 3. Resolve the target and log the choice. - algorithm::ensure_model_is_target(&self.targets, &score.target)?; + algorithm::ensure_model_is_target(driver.models_for(&Category::Any), &score.target)?; let target = score.target.clone(); tracing::info!(algorithm=self.name, target=%score.target, confidence=score.confidence, "Model selected"); @@ -280,11 +236,13 @@ where let event = Event::Decision { request, selected_model_id: &target, + category: score.category.clone(), + driver, }; processor.process(state, event).await?; } - Ok((target, served)) + Ok((score, served)) } } @@ -339,11 +297,10 @@ where #[cfg(test)] mod tests { use super::*; + use crate::Classification; + use crate::algorithms::llm_class::DefaultCategoryClassifier; use crate::algorithms::util::prompts; - use crate::core::classifier::Classification; - use crate::{SystemPromptProcessor, TargetPrompts}; - - use crate::core::testing::{Serve, echo, reply, test_drive}; + use crate::core::testing::{Serve, category_models, echo, reply, test_drive_with_models}; use switchyard_protocol::{LlmRequest, Message, Metadata, Role, completion_text, text_request}; #[derive(Debug, thiserror::Error)] @@ -368,95 +325,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); @@ -466,7 +340,7 @@ mod tests { &self, _state: &mut (), _request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { Ok(( Classification::Scores( @@ -475,6 +349,7 @@ mod tests { .map(|s| Score { confidence: s.confidence, target: s.target.clone(), + category: None, }) .collect(), ), @@ -487,6 +362,7 @@ mod tests { Score { confidence, target: ModelId::from(target), + category: None, } } @@ -518,7 +394,13 @@ mod tests { where S: Default + Send + 'static, { - let (selected_model, response) = test_drive(router.clone(), request, serve).await?; + let (selected_model, response) = test_drive_with_models( + router.clone(), + request, + category_models(Category::Any, &["strong", "weak"]), + serve, + ) + .await?; let text = response .llm_response .into_agg() @@ -546,64 +428,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,27 +444,82 @@ mod tests { } } - let recorder = Arc::new(PromptRecorder::default()); - let router = FallThrough::new(target_set(&["capable", "efficient"])) + let captured = Arc::new(Mutex::new(None)); + let router = FallThrough::new() .with_processor(Arc::new(Noting)) - .with_classifier(Arc::new(DefaultTarget::new("capable"))); + .with_classifier(Arc::new(DefaultCategoryClassifier(Category::Any))); - let call = routed_prompt_call(&recorder, router).await?; + test_drive_with_models( + Arc::new(router), + Request { + llm_request: text_request(Some("auto".to_string()), "fix the build"), + raw_request: None, + metadata: None, + }, + category_models(Category::Any, &["capable", "efficient"]), + 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(()) } #[tokio::test] - async fn selected_target_leads_the_ordered_candidate_list() -> Result<()> { + async fn the_selected_category_leads_the_candidate_list() -> Result<()> { use futures::StreamExt; + // The decision came from `capable`, so the rest of that category is tried before + // the models it does not contain. let router = Arc::new( - FallThrough::<()>::new(target_set(&["weak", "mid", "strong"])) - .with_classifier(fixed(vec![score("mid", 0.9)])), + FallThrough::<()>::new() + .with_classifier(Arc::new(DefaultCategoryClassifier(Category::Capable))), + ); + let models = crate::RuntimeModels::new( + [ + (Category::Capable, target_set(&["premium", "reasoning"])), + (Category::Any, target_set(&["fast", "premium", "reasoning"])), + ] + .into(), ); - let stream = router.run_stream(request()); + + let stream = router.run_stream(request(), Arc::new(models)); + tokio::pin!(stream); + while let Some(step) = stream.next().await { + if let crate::Step::Done(outcome) = step? { + assert_eq!( + outcome.selected_model_ids, + target_set(&["premium", "reasoning", "fast"]) + ); + return Ok(()); + } + } + Err(test_error("expected a Done step")) + } + + #[tokio::test] + async fn selected_target_leads_the_ordered_candidate_list() -> Result<()> { + use futures::StreamExt; + + let router = + Arc::new(FallThrough::<()>::new().with_classifier(fixed(vec![score("mid", 0.9)]))); + + let models = category_models(Category::Any, &["weak", "mid", "strong"]); + let stream = router.run_stream(request(), Arc::new(models.into())); tokio::pin!(stream); while let Some(step) = stream.next().await { if let crate::Step::Done(outcome) = step? { @@ -658,7 +537,7 @@ mod tests { #[tokio::test] async fn argmax_picks_the_highest_confidence_target() -> Result<()> { - let router = FallThrough::<()>::new(target_set(&["strong", "weak"])) + let router = FallThrough::<()>::new() .with_classifier(fixed(vec![score("weak", 0.2), score("strong", 0.9)])); let (model, selected_model) = run_with(router, echo()).await?; assert_eq!(model, "strong"); @@ -669,7 +548,7 @@ mod tests { #[tokio::test] async fn falls_through_the_first_abstaining_classifier() -> Result<()> { // First classifier abstains (empty); the second decides. - let router = FallThrough::<()>::new(target_set(&["strong", "weak"])) + let router = FallThrough::<()>::new() .with_classifier(fixed(vec![])) .with_classifier(fixed(vec![score("weak", 1.0)])); let (model, _) = run_with(router, echo()).await?; @@ -680,7 +559,7 @@ mod tests { #[tokio::test] async fn first_deciding_classifier_wins_the_cascade() -> Result<()> { // The first classifier decides; the second is never consulted. - let router = FallThrough::<()>::new(target_set(&["strong", "weak"])) + let router = FallThrough::<()>::new() .with_classifier(fixed(vec![score("strong", 0.6)])) .with_classifier(fixed(vec![score("weak", 1.0)])); let (model, _) = run_with(router, echo()).await?; @@ -690,8 +569,7 @@ mod tests { #[tokio::test] async fn all_abstaining_is_an_error() -> Result<()> { - let router = - FallThrough::<()>::new(target_set(&["strong", "weak"])).with_classifier(fixed(vec![])); + let router = FallThrough::<()>::new().with_classifier(fixed(vec![])); let error = run_with(router, echo()) .await .err() @@ -703,34 +581,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn classifiers_receive_the_per_request_driver() -> Result<()> { - // A classifier that only decides when handed a driver — proving the cascade offers - // the per-request driver to every classifier (driver-backed ones need it). - struct NeedsDriver; - - #[async_trait] - impl Classifier for NeedsDriver { - async fn score( - &self, - _state: &mut (), - _request: &mut Request, - driver: Option<&Driver>, - ) -> Result<(Classification, Option)> { - match driver { - Some(_) => Ok((Classification::Scores(vec![score("strong", 1.0)]), None)), - None => Err(test_error("expected a driver")), - } - } - } - - let router = FallThrough::<()>::new(target_set(&["strong", "weak"])) - .with_classifier(Arc::new(NeedsDriver)); - let (model, _) = run_with(router, echo()).await?; - assert_eq!(model, "strong"); - Ok(()) - } - #[tokio::test] async fn processor_observes_request_then_decision() -> Result<()> { use parking_lot::Mutex; @@ -753,7 +603,7 @@ mod tests { } let seen = Arc::new(Mutex::new(Vec::new())); - let router = FallThrough::<()>::new(target_set(&["strong", "weak"])) + let router = FallThrough::<()>::new() .with_processor(Arc::new(RecordingProcessor(seen.clone()))) .with_classifier(fixed(vec![score("strong", 1.0)])); run_with(router, echo()).await?; @@ -792,7 +642,7 @@ mod tests { &self, _state: &mut (), request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { *self.0.lock() = request .llm_request @@ -810,8 +660,7 @@ mod tests { let seen_by_classifier = Arc::new(Mutex::new(Vec::new())); let seen_by_model = Arc::new(Mutex::new(None)); - let targets = target_set(&["strong"]); - let router = FallThrough::new(targets) + let router = FallThrough::new() .with_processor(Arc::new(Appender("first"))) .with_processor(Arc::new(Appender("second"))) .with_classifier(Arc::new(TrailClassifier(seen_by_classifier.clone()))); @@ -865,7 +714,7 @@ mod tests { &self, state: &mut TurnState, _request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { let target = if state.count >= 2 { "strong" } else { "weak" }; Ok((Classification::Scores(vec![score(target, 1.0)]), None)) @@ -873,7 +722,7 @@ mod tests { } let router = Arc::new( - FallThrough::::new_with_state(target_set(&["strong", "weak"])) + FallThrough::::new_with_state() .with_processor(Arc::new(CountingProcessor)) .with_classifier(Arc::new(ThresholdClassifier)), ); @@ -921,7 +770,7 @@ mod tests { #[tokio::test] async fn final_session_is_removed_when_routing_fails() { - let router = Arc::new(FallThrough::::new_with_state(target_set(&["strong"]))); + let router = Arc::new(FallThrough::::new_with_state()); let final_request = Request { metadata: Some(Metadata { session_id: Some("session-1".to_string()), @@ -931,7 +780,13 @@ mod tests { ..request() }; - let result = test_drive(router.clone(), final_request, echo()).await; + let result = test_drive_with_models( + router.clone(), + final_request, + category_models(Category::Any, &["strong"]), + echo(), + ) + .await; assert!(matches!(result, Err(LibsyError::AlgorithmError { .. }))); let states = router @@ -944,7 +799,7 @@ mod tests { #[test] fn cleanup_removes_only_inactive_idle_sessions() { - let router = FallThrough::::new_with_state(target_set(&["strong"])); + let router = FallThrough::::new_with_state(); let _active_state = router .session_state(&request()) // session-1 .expect("session state was inserted"); diff --git a/crates/libsy/src/algorithms/llm_class.rs b/crates/libsy/src/algorithms/llm_class.rs index e247566eb..4bb5e7ff0 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}; @@ -196,24 +196,14 @@ impl ClassifierInput for TaskInput { } } -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, } @@ -228,28 +218,38 @@ 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 + }; + // The chosen category may have no models configured. That is ambiguous, not an + // error, so the surrounding router applies its configured fallback. + let Some(target) = driver.models_for(&category).first().cloned() else { + return Ok(Classification::Ambiguous(vec![])); }; - Classification::Scores(vec![Score { - target: target.clone(), + Ok(Classification::Scores(vec![Score { + target, confidence: 1.0, - }]) + category: Some(category), + }])) } } @@ -410,7 +410,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, @@ -418,7 +418,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(), @@ -486,18 +486,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, @@ -524,34 +523,47 @@ pub struct LlmTaskClassifier { } struct ClassifierRouteConfig { - default_target: ModelId, + default_target: Category, classify_trigger: ClassifyTrigger, message_hash_fallback: bool, } +/// Terminal classifier for a cascade whose classifiers may all abstain. +/// Closes a cascade with the first runtime model in `category`. +pub struct DefaultCategoryClassifier(pub Category); + +#[async_trait] +impl Classifier for DefaultCategoryClassifier { + async fn score( + &self, + _state: &mut S, + _request: &mut Request, + driver: &Driver, + ) -> Result<(Classification, Option)> { + let target = driver.first_model_for(&self.0)?; + driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"})); + Ok(( + Classification::Scores(vec![Score { + target: target.clone(), + confidence: 0.0, + category: Some(self.0.clone()), + }]), + 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. @@ -559,14 +571,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, }, @@ -581,49 +589,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, @@ -632,75 +617,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, @@ -713,9 +645,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( @@ -725,15 +655,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, }, @@ -741,24 +669,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() .with_name(ALGORITHM_NAME) .with_classifier(Arc::clone(&inner)), inner, @@ -772,11 +689,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" @@ -784,7 +699,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().with_name(ALGORITHM_NAME); if let Some(affinity) = affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref() { @@ -793,7 +708,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()) @@ -803,25 +718,13 @@ 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( &self, state: &mut State, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { self.inner.score(state, request, driver).await } @@ -844,6 +747,7 @@ impl Algorithm for LlmTaskClassifier { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; use parking_lot::Mutex; @@ -851,16 +755,18 @@ 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; + type CapabilityJudge = StructuredJudge>; + fn test_config(base_threshold: f64) -> TaskClassifierConfig { TaskClassifierConfig { base_threshold, @@ -869,7 +775,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", Arc::new(runtime_models().into())).0 } fn verdict( @@ -890,7 +813,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 { @@ -988,9 +911,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), }, )?)) @@ -1031,8 +951,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!( @@ -1046,10 +971,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(), @@ -1071,16 +1003,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(()) @@ -1090,9 +1025,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."), @@ -1100,7 +1032,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); @@ -1112,18 +1050,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(()) @@ -1133,9 +1075,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, @@ -1144,10 +1083,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?; @@ -1156,6 +1103,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| -> HashMap> { + [ + ( + 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(); @@ -1169,8 +1187,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(()) @@ -1197,9 +1215,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(), @@ -1228,21 +1243,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), })?; } @@ -1267,7 +1271,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()); @@ -1277,14 +1281,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..581b01168 100644 --- a/crates/libsy/src/algorithms/passthrough.rs +++ b/crates/libsy/src/algorithms/passthrough.rs @@ -5,24 +5,13 @@ 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(), - } - } -} +pub struct Passthrough; #[async_trait::async_trait] impl Algorithm for Passthrough { @@ -30,11 +19,16 @@ 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 models = driver.models_for(&Category::Any).to_vec(); + // Selected is the first one. The rest are fallbacks. + let Some(target) = models.first() else { + return Err(LibsyError::NoTargets); + }; + tracing::info!(target = %target, "passthrough selected target"); Ok(RoutingOutcome::route_to( - self.target.clone(), - Vec::new(), + target.clone(), + models[1..].to_vec(), request, )) } @@ -46,8 +40,8 @@ mod tests { 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::{category_models, echo, test_drive_with_models}; + use switchyard_protocol::{Category, Request, completion_text, text_request}; #[tokio::test] async fn test_passthrough() -> crate::Result<()> { @@ -57,8 +51,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); + let models = category_models(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..7b9b86a67 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,12 @@ 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>, + weight_count: Option, rng: Mutex, } @@ -37,63 +37,38 @@ 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 if 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 weight_count = weights.as_ref().map(Vec::len); + 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, + weight_count, 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 +86,35 @@ where &self, _state: &mut S, _request: &mut Request, - _driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { + // All the available models + let options = driver.models_for(&Category::Any); + if options.is_empty() { + return Err(LibsyError::NoTargets); + } + if let Some(weight_count) = self.weight_count + && weight_count != options.len() + { + return Err(invalid_weights(format!( + "{weight_count} weights were provided but Category::Any has {} runtime models.", + options.len() + ))); + } + 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, + category: Some(Category::Any), }]), None, )) @@ -129,18 +127,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() .with_name("random") .with_classifier(classifier); Ok(Self { inner }) @@ -165,12 +155,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::{category_models, echo, test_drive_with_models}; use switchyard_protocol::Request; fn request() -> Request { @@ -191,22 +181,24 @@ mod tests { } } - fn target_set(names: &[&str]) -> Vec { - names.iter().map(|name| ModelId::from(*name)).collect() + fn algorithm(weights: Option>, seed: Option) -> Result { + Random::new(weights, seed) } - fn algorithm(names: &[&str], weights: Option>, seed: Option) -> Result { - Random::new(target_set(names), weights, seed) + fn shared_algorithm() -> Result> { + Ok(Arc::new(algorithm(None, None)?)) } - fn shared_algorithm(names: &[&str]) -> Result> { - Ok(Arc::new(algorithm(names, 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 @@ -220,8 +212,10 @@ mod tests { #[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 = category_models(Category::Any, &["only/model"]); + let (selected_model, response) = + test_drive_with_models(algorithm, request(), models, echo()).await?; assert_eq!( response @@ -237,12 +231,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 = category_models(Category::Any, &["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 +259,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 = category_models(Category::Any, &["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 +281,24 @@ mod tests { #[tokio::test] async fn affinity_reuses_the_initial_random_selection() -> Result<()> { let names = ["a/model", "b/model"]; + let models = category_models(Category::Any, &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() .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() @@ -316,8 +307,9 @@ mod tests { let mut state = (); let mut request = request_for_session("session-1"); + let driver = Driver::new("test", Arc::new(models.clone().into())).0; let retained = affinity - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, &driver) .await? .0 .argmax(false)?; @@ -326,7 +318,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 +335,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 +349,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 = category_models(Category::Any, &["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 2f9adf945..4638f5e9b 100644 --- a/crates/libsy/src/algorithms/stage.rs +++ b/crates/libsy/src/algorithms/stage.rs @@ -19,21 +19,47 @@ 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, ToolSemantics, ToolSignalProcessor}; use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::{Classification, Classifier, Score}; +use crate::core::processor::{Event, Processor}; use crate::core::state::State; use crate::{LibsyError, Result}; -use switchyard_protocol::{ModelId, Request, Response}; +use switchyard_protocol::{Category, Request, Response}; /// Telemetry name for a router this module assembles. const STAGE_ROUTER: &str = "stage_router"; +struct TierPromptProcessor { + capable: Option, + efficient: Option, +} + +#[async_trait] +impl Processor for TierPromptProcessor { + async fn process(&self, _state: &mut State, event: Event<'_>) -> Result<()> { + if let Event::Decision { + request, category, .. + } = event + { + let prompt = match category { + Some(Category::Capable) => self.capable.as_deref(), + Some(Category::Efficient) => self.efficient.as_deref(), + _ => 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. /// /// The classifiers themselves are composition-agnostic and write no state; only @@ -49,18 +75,16 @@ impl Classifier for SourceStamp { &self, state: &mut State, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { let (classification, served) = self.inner.score(state, request, driver).await?; // An abstaining classifier passes the turn on, so it is not its to claim. if let Some(winner) = classification.argmax(false)? { record_decision_source(state, self.source); record_routing_decision(self.source, &winner.target); - if let Some(driver) = driver { - driver.set_evidence_if_empty(serde_json::json!({ - "source": self.source.as_str(), - })); - } + driver.set_evidence_if_empty(serde_json::json!({ + "source": self.source.as_str(), + })); } Ok((classification, served)) } @@ -68,7 +92,6 @@ impl Classifier for SourceStamp { /// Closes the cascade at zero confidence: a fallback, not a judgement. struct FallOpen { - targets: StageTargets, default_tier: Tier, } @@ -78,14 +101,19 @@ impl Classifier for FallOpen { &self, state: &mut State, _request: &mut Request, - _driver: Option<&Driver>, + driver: &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.first_model_for(&category)?.clone(); Ok(( Classification::Scores(vec![Score { target, confidence: 0.0, + category: Some(category), }]), None, )) @@ -94,9 +122,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 — @@ -118,12 +143,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, } @@ -137,7 +162,8 @@ impl StageRouterConfig { recent_window: None, tool_semantics: ToolSemantics::default(), handoff_notes: None, - tier_prompts: TargetPrompts::default(), + capable_system_prompt: None, + efficient_system_prompt: None, llm_fallback: None, } } @@ -151,14 +177,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)?, }) } } @@ -180,11 +205,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!( @@ -194,16 +215,10 @@ pub(crate) fn build_stage_route( }); } config.tool_semantics.validate()?; - // 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); } @@ -211,19 +226,13 @@ pub(crate) fn build_stage_route( recent_window: config.recent_window.unwrap_or(DEFAULT_RECENT_WINDOW), tool_semantics: config.tool_semantics, }; - let target_set = vec![capable.clone(), efficient.clone()]; - let mut router = FallThrough::::new_with_state(target_set) + let mut router = FallThrough::::new_with_state() .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, @@ -234,15 +243,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; @@ -251,10 +263,9 @@ 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::{empty_driver, 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); @@ -265,12 +276,13 @@ mod tests { &self, _state: &mut State, _request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { Ok(( Classification::Scores(vec![Score { target: ModelId::from(self.0), confidence: 1.0, + category: None, }]), None, )) @@ -286,7 +298,7 @@ mod tests { &self, _state: &mut State, _request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { Ok((Classification::Ambiguous(vec![]), None)) } @@ -299,7 +311,7 @@ mod tests { }; let mut state = State::default(); stamp - .score(&mut state, &mut Request::default(), None) + .score(&mut state, &mut Request::default(), &empty_driver()) .await?; Ok(match state.extra.get(DECISION_SOURCE_KEY) { Some(StateValue::String(source)) => Some(source.clone()), @@ -327,12 +339,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 { .. }) )); } @@ -341,21 +370,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(()) } @@ -364,11 +392,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 { @@ -381,7 +405,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), @@ -419,15 +442,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!( @@ -454,12 +481,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: {:?}", @@ -481,8 +520,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!( @@ -503,7 +547,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), @@ -518,9 +568,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"); @@ -543,7 +605,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(()) @@ -554,7 +622,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..7fdc2c1bd 100644 --- a/crates/libsy/src/algorithms/subagent.rs +++ b/crates/libsy/src/algorithms/subagent.rs @@ -5,24 +5,23 @@ use std::sync::Arc; -use switchyard_protocol::{Metadata, ModelId, Request}; +use switchyard_protocol::{Category, Metadata, Request}; -use super::fall_through::{DefaultTarget, FallThrough}; +use super::fall_through::FallThrough; use super::util::affinity::{AffinityRouter, ClassifyTrigger}; use super::util::subagent::SubagentGate; -use crate::core::algorithm::{self, Algorithm, Driver}; +use crate::algorithms::llm_class::DefaultCategoryClassifier; +use crate::core::algorithm::{Algorithm, Driver}; use crate::core::classifier::Classifier; use crate::core::state::State; use crate::{LibsyError, Result, RoutingOutcome}; /// Runtime components for delegated sub-agent routing. pub struct SubagentRouterConfig { - /// Targets the delegated-work classifier may select. - pub targets: Vec, /// Classifier invoked for delegated work according to `classify_trigger`. pub classifier: Arc>, - /// Child target used when `classifier` abstains. - pub default_target: ModelId, + /// Child model category used when `classifier` abstains. + pub default_target: Category, /// Controls whether each child is classified once or on every request. pub classify_trigger: ClassifyTrigger, /// Unsupported for child routing because child identity must come from harness metadata. @@ -30,13 +29,11 @@ pub struct SubagentRouterConfig { } impl SubagentRouterConfig { - /// Routes delegated work directly to one fixed target. - pub fn fixed_target(target: impl Into) -> Self { - let target = target.into(); + /// Routes all delegated work to the first model in the sub-agent `Any` category. + pub fn fixed_target() -> Self { Self { - targets: vec![target.clone()], - classifier: Arc::new(DefaultTarget::new(target.clone())), - default_target: target, + classifier: Arc::new(DefaultCategoryClassifier(Category::Any)), + default_target: Category::Any, classify_trigger: ClassifyTrigger::EveryRequest, message_hash_fallback: false, } @@ -54,10 +51,8 @@ impl SubagentRouter { /// /// # Errors /// - /// Returns an error when the child default is not a child target or when the affinity - /// settings cannot identify delegated children safely. + /// Returns an error when the affinity settings cannot identify delegated children safely. pub fn new(parent: Arc, config: SubagentRouterConfig) -> Result { - algorithm::ensure_model_is_target(&config.targets, &config.default_target)?; if config.message_hash_fallback { return Err(LibsyError::AlgorithmError { message: "sub-agent routing cannot use message_hash_fallback".to_string(), @@ -65,12 +60,10 @@ impl SubagentRouter { } let mut subagent = match config.classify_trigger { - ClassifyTrigger::EveryRequest => { - FallThrough::new_with_state(config.targets).with_name("subagent") - } + ClassifyTrigger::EveryRequest => FallThrough::new_with_state().with_name("subagent"), ClassifyTrigger::NewSession => { let affinity = Arc::new(AffinityRouter::for_subagents()); - FallThrough::new_with_state(config.targets) + FallThrough::new_with_state() .with_name("subagent") .with_processor(affinity.clone()) .with_classifier(affinity) @@ -84,7 +77,7 @@ impl SubagentRouter { }; subagent = subagent .with_classifier(Arc::new(SubagentGate::new(config.classifier))) - .with_classifier(Arc::new(DefaultTarget::new(config.default_target))); + .with_classifier(Arc::new(DefaultCategoryClassifier(config.default_target))); Ok(Self { parent, subagent }) } @@ -102,7 +95,8 @@ impl Algorithm for SubagentRouter { .as_ref() .is_some_and(Metadata::is_subagent_work) { - self.subagent.execute(driver, request).await + // Delegated work routes over the sub-agent's own models, never the parent's. + self.subagent.execute(driver.for_subagent()?, request).await } else { self.parent.clone().route(driver, request).await } @@ -118,18 +112,17 @@ 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, + LlmClassifierConfig, LlmTaskClassifier, RuntimeModels, State, }; struct ScriptedClassifier { @@ -142,18 +135,20 @@ mod tests { &self, _state: &mut State, _request: &mut Request, - _driver: Option<&Driver>, + driver: &Driver, ) -> crate::Result<(Classification, Option)> { - let scores = match self.calls.fetch_add(1, Ordering::Relaxed) { - 0 => vec![Score { - confidence: 1.0, - target: ModelId::from("worker"), - }], - 1 => vec![Score { + let category = match self.calls.fetch_add(1, Ordering::Relaxed) { + 0 => Some(Category::Capable), + 1 => Some(Category::Efficient), + _ => None, + }; + let scores = match category { + Some(category) => vec![Score { confidence: 1.0, - target: ModelId::from("reviewer"), + target: driver.first_model_for(&category)?.clone(), + category: Some(category), }], - _ => Vec::new(), + None => Vec::new(), }; Ok((Classification::Scores(scores), None)) } @@ -177,17 +172,12 @@ mod tests { })) } - fn parent() -> Arc { - Arc::new(Passthrough::new("parent")) - } - fn configured(classifier: Arc>) -> crate::Result> { Ok(Arc::new(SubagentRouter::new( - parent(), + Arc::new(Passthrough), SubagentRouterConfig { - targets: vec![ModelId::from("worker"), ModelId::from("reviewer")], classifier, - default_target: ModelId::from("worker"), + default_target: Category::Capable, classify_trigger: ClassifyTrigger::NewSession, message_hash_fallback: false, }, @@ -201,11 +191,33 @@ 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?; + // The parent and its children route over separate model groups. + let models = RuntimeModels::new([(Category::Any, vec![ModelId::from("parent")])].into()) + .with_subagent( + [ + ( + Category::Any, + vec![ModelId::from("worker"), ModelId::from("reviewer")], + ), + (Category::Capable, vec![ModelId::from("worker")]), + (Category::Efficient, vec![ModelId::from("reviewer")]), + ] + .into(), + ); + let (selected_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,33 +225,36 @@ 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!(selected_parent, "parent"); assert_eq!(first, "worker"); assert_eq!(same_child, "worker"); assert_eq!(sibling, "reviewer"); assert_eq!(defaulted, "worker"); assert_eq!(maintenance, "parent"); assert_eq!(classifier.calls.load(Ordering::Relaxed), 3); + + let fixed = Arc::new(SubagentRouter::new( + Arc::new(Passthrough), + SubagentRouterConfig::fixed_target(), + )?); + let (fixed, _) = test_drive_with_models(fixed, child("fixed"), models, echo()).await?; + assert_eq!(fixed, "worker"); Ok(()) } #[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 +285,33 @@ 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 = RuntimeModels::new([(Category::Any, vec![ModelId::from("parent")])].into()) + .with_subagent( + [ + (Category::Judge, vec![ModelId::from("judge")]), + (Category::Capable, vec![ModelId::from("worker")]), + (Category::Efficient, vec![ModelId::from("reviewer")]), + ( + Category::Any, + vec![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..6e490810d 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. @@ -31,12 +33,13 @@ impl Classifier for AlwaysOrchestrator { &self, _state: &mut (), _request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { Ok(( Classification::Scores(vec![Score { confidence: 0.5, target: ModelId::from("orchestrator"), + category: None, }]), None, )) @@ -63,7 +66,7 @@ fn request(headers: &[(&str, &str)]) -> Request { fn router() -> Arc { let affinity = Arc::new(AffinityRouter::for_subagents()); Arc::new( - FallThrough::<()>::new(targets()) + FallThrough::<()>::new() .with_processor(affinity.clone()) .with_classifier(affinity) .with_classifier(Arc::new(SubagentOverride::new("worker"))) @@ -73,7 +76,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() @@ -140,7 +145,7 @@ async fn harness_maintenance_turns_are_not_forced_to_the_worker() -> Result<()> /// Builds a cascade whose override scores `worker`, sharing `affinity` across instances. fn router_overriding_to(affinity: Arc, worker: &str) -> Arc { Arc::new( - FallThrough::<()>::new(targets()) + FallThrough::<()>::new() .with_processor(affinity.clone()) .with_classifier(affinity) .with_classifier(Arc::new(SubagentOverride::new(worker))) @@ -200,7 +205,7 @@ async fn a_cascade_without_the_override_still_routes_root_traffic() -> Result<() // cascade, which is the point of composing them rather than nesting one in the other. let affinity = Arc::new(AffinityRouter::for_subagents()); let router = Arc::new( - FallThrough::<()>::new(targets()) + FallThrough::<()>::new() .with_processor(affinity.clone()) .with_classifier(affinity) .with_classifier(Arc::new(AlwaysOrchestrator)), diff --git a/crates/libsy/src/algorithms/util.rs b/crates/libsy/src/algorithms/util.rs index 3f1202e26..9a290920f 100644 --- a/crates/libsy/src/algorithms/util.rs +++ b/crates/libsy/src/algorithms/util.rs @@ -23,6 +23,7 @@ pub(crate) fn decisive(target: &ModelId) -> Classification { Classification::Scores(vec![Score { target: target.clone(), confidence: 1.0, + category: None, }]) } diff --git a/crates/libsy/src/algorithms/util/affinity.rs b/crates/libsy/src/algorithms/util/affinity.rs index f557a28d6..0d04f297e 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}; @@ -177,6 +177,7 @@ where if let Event::Decision { request, selected_model_id, + .. } = event && let Some(key) = self.affinity_key(request) { @@ -232,7 +233,7 @@ where &self, _state: &mut S, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> crate::Result<(Classification, Option)> { let Some(key) = self.affinity_key(request) else { return Ok((Classification::Scores(Vec::new()), None)); @@ -242,17 +243,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(); - if assigned.is_some() - && let Some(driver) = driver - { + // An empty `any` group carries no information about which models are still + // available, so it must not be read as "every assignment is now stale". + let available = driver.models_for(&Category::Any); + let mut assignments = self.assignments.lock(); + let assigned = assignments.get(&key).cloned(); + let assigned = match assigned.as_ref() { + Some(target) if !available.is_empty() && !available.contains(target) => { + assignments.remove(&key); + None + } + assigned => assigned, + }; + if assigned.is_some() { driver.set_evidence(serde_json::json!({"source": "retained"})); } Ok(( Classification::Scores(match assigned { Some(target) => vec![Score { confidence: 1.0, - target, + target: target.clone(), + category: None, }], None => Vec::new(), }), @@ -278,6 +289,8 @@ mod tests { use switchyard_protocol::{LlmRequest, Metadata, ToolResult, text_request}; + use crate::core::algorithm::RuntimeModels; + /// Boxed, thread-safe error type keeping the test helpers ergonomic. type BoxErr = Box; @@ -285,6 +298,22 @@ mod tests { ModelId::from(target) } + fn driver() -> Driver { + Driver::new( + "test", + Arc::new(RuntimeModels::new( + [( + Category::Any, + ["model-a", "model-b", "weak", "strong"] + .map(ModelId::from) + .to_vec(), + )] + .into(), + )), + ) + .0 + } + fn request(metadata: Metadata) -> Request { Request { llm_request: text_request(Some("auto".to_string()), "hi"), @@ -350,6 +379,8 @@ mod tests { Event::Decision { request, selected_model_id: &selected_model_id, + category: None, + driver: &driver(), }, ) .await?; @@ -362,7 +393,7 @@ mod tests { state: &mut (), request: &mut Request, ) -> Result, BoxErr> { - match classifier.score(state, request, None).await?.0 { + match classifier.score(state, request, &driver()).await?.0 { Classification::Scores(scores) => Ok(scores), Classification::Ambiguous(_) => Err("affinity never returns ambiguous scores".into()), } @@ -643,6 +674,8 @@ mod tests { Event::Decision { request: &mut first, selected_model_id: &fixed_model("model-a"), + category: None, + driver: &driver(), }, ) .await?; @@ -668,6 +701,8 @@ mod tests { Event::Decision { request: &mut unkeyed, selected_model_id: &fixed_model("model-a"), + category: None, + driver: &driver(), }, ) .await?; @@ -692,6 +727,8 @@ mod tests { Event::Decision { request: &mut second, selected_model_id: &fixed_model("model-b"), + category: None, + driver: &driver(), }, ) .await?; @@ -701,6 +738,8 @@ mod tests { Event::Decision { request: &mut first, selected_model_id: &fixed_model("model-a"), + category: None, + driver: &driver(), }, ) .await?; @@ -862,4 +901,29 @@ mod tests { ); Ok(()) } + + #[tokio::test] + async fn an_unprovisioned_any_group_does_not_evict_assignments() -> Result<(), BoxErr> { + let router = AffinityRouter::new(); + let mut state = (); + let mut req = request(session("session-1", "agent-a")); + retain(&router, &mut state, &mut req, "model-a").await?; + + let empty = Driver::new("test", Arc::new(RuntimeModels::default())).0; + let (classification, _) = router.score(&mut state, &mut req, &empty).await?; + let Classification::Scores(retained) = classification else { + return Err("affinity never returns ambiguous scores".into()); + }; + assert_eq!(retained.first().map(|s| s.target.as_str()), Some("model-a")); + + // The assignment survived, so a later turn with the group present still latches. + assert_eq!( + scores(&router, &mut state, &mut req) + .await? + .first() + .map(|s| s.target.as_str()), + Some("model-a") + ); + Ok(()) + } } diff --git a/crates/libsy/src/algorithms/util/escalation.rs b/crates/libsy/src/algorithms/util/escalation.rs index fb093f52a..b0979e286 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}; @@ -124,15 +125,16 @@ pub(crate) type EscalationJudge = StructuredJudge) -> Classification { + fn to_classification( + &self, + verdict: Option<&EscalationVerdict>, + driver: &Driver, + ) -> Result { if let Some(verdict) = verdict { tracing::debug!( escalate = verdict.escalate, @@ -141,15 +143,17 @@ impl JudgePolicy for EscalationPolicy { ); } 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(), + category: Some(Category::Capable), + }])), + Some(_) => Ok(Classification::Scores(vec![Score { + target: driver.first_model_for(&Category::Efficient)?.clone(), confidence: 1.0, - }]), - None => Classification::Ambiguous(Vec::new()), + category: Some(Category::Efficient), + }])), + None => Ok(Classification::Ambiguous(Vec::new())), } } } @@ -167,14 +171,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, @@ -189,8 +190,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..f12c7a0dd 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| { @@ -338,19 +345,16 @@ where &self, state: &mut State, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { - // A missing driver is a broken composition, not an unavailable judge. - let Some(driver) = driver else { + let judge_models = driver.models_for(&Category::Judge); + if judge_models.is_empty() { return Err(LibsyError::AlgorithmError { - message: format!( - "judge classifier for target {:?} requires a driver to call it", - self.target - ), + message: "no models available for category Judge".to_string(), }); - }; - let verdict = self.verdict(state, request, driver).await; - let classification = self.policy.to_classification(verdict.as_ref()); + } + 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())) @@ -390,6 +394,8 @@ fn strip_json_fence(text: &str) -> &str { #[cfg(test)] mod tests { use super::*; + use crate::core::algorithm::RuntimeModels; + use std::sync::Arc; use futures::StreamExt; use http::StatusCode; @@ -428,21 +434,26 @@ 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, - }]) + category: None, + }])) } } fn classifier() -> JudgeClassifier { - JudgeClassifier::new(TestJudge, ModelId::from("judge"), TestPolicy) + JudgeClassifier::new(TestJudge, TestPolicy) } fn request() -> Request { @@ -557,7 +568,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 = RuntimeModels::new([(Category::Judge, vec![ModelId::from("judge")])].into()); + let (driver, step_rx) = Driver::new("test", Arc::new(models)); let mut steps = tokio_stream::wrappers::ReceiverStream::new(step_rx); let classifier = classifier(); let mut state = State::default(); @@ -568,10 +580,8 @@ mod tests { let _ = call.respond(reply); } }; - let (classification, ()) = tokio::join!( - classifier.score(&mut state, &mut request, Some(&driver)), - serve - ); + let (classification, ()) = + tokio::join!(classifier.score(&mut state, &mut request, &driver), serve); let (classification, _) = classification?; selected(classification) } @@ -694,24 +704,6 @@ mod tests { assert_eq!(libsy_error_reason(&error), "call_error"); } - #[tokio::test] - async fn a_missing_driver_is_an_error_not_a_fallback() -> Result<()> { - let mut request = request(); - let error = classifier() - .score(&mut State::default(), &mut request, None) - .await - .err() - .ok_or_else(|| LibsyError::AlgorithmError { - message: "expected a missing-driver error".to_string(), - })?; - - assert!( - matches!(&error, LibsyError::AlgorithmError { message } if message.contains("judge")), - "unexpected error: {error}" - ); - Ok(()) - } - #[test] fn fenced_replies_parse_as_verdicts() -> Result<()> { let judge = TestJudge; 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 f9e3ef806..b654ee8b6 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")] @@ -557,18 +511,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, @@ -605,7 +557,7 @@ impl Classifier for StageClassifier { &self, state: &mut State, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { let tool_signals = &state.tool_signals; let Some(signal) = tool_signals else { @@ -623,24 +575,26 @@ 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 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()); @@ -648,6 +602,7 @@ impl Classifier for StageClassifier { Classification::Scores(vec![Score { target: target.clone(), confidence: conf, + category: Some(category), }]), None, )) @@ -660,8 +615,24 @@ impl Classifier for StageClassifier { #[cfg(test)] mod tests { use super::*; + use crate::core::algorithm::RuntimeModels; use serde_json::json; - use switchyard_protocol::{Metadata, Request, WireFormat, text_request}; + use std::sync::Arc; + use switchyard_protocol::{Metadata, ModelId, Request, WireFormat, text_request}; + + fn driver() -> Driver { + Driver::new( + "stage_test", + Arc::new(RuntimeModels::new( + [ + (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})); @@ -800,11 +771,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 { @@ -818,8 +784,8 @@ 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) - .score(&mut state, &mut Request::default(), None) + let classification = StageClassifier::new(PickerMode::EfficientFirst, 0.5) + .score(&mut state, &mut Request::default(), &driver()) .await?; assert!(classification.0.argmax(false)?.is_none()); assert!(matches!( @@ -837,8 +803,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(), &driver) .await?; match classification.0 { Classification::Scores(scores) => { @@ -868,8 +835,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(), &driver) .await?; match classification.0 { Classification::Scores(scores) => { @@ -886,8 +854,8 @@ 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) - .score(&mut state, &mut Request::default(), None) + let classification = StageClassifier::new(PickerMode::EfficientFirst, 0.5) + .score(&mut state, &mut Request::default(), &driver()) .await?; assert!(classification.0.argmax(false)?.is_none()); assert!(matches!( @@ -956,7 +924,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, @@ -993,9 +961,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, &driver) .await?; assert_eq!(trailing_text(&request), Some(format!("hi|{ESCALATION}"))); @@ -1008,10 +977,11 @@ 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, &driver).await?; assert_eq!(trailing_text(&request), Some(format!("hi|{ESCALATION}"))); } Ok(()) @@ -1028,9 +998,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, &driver) .await?; assert_eq!(trailing_text(&request), Some(format!("hi|{DEESCALATION}"))); @@ -1045,7 +1016,7 @@ mod tests { let mut request = request(); let classification = noting_classifier(PickerMode::CapableFirst) - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, &driver()) .await?; assert!(matches!(classification.0, Classification::Ambiguous(_))); @@ -1057,9 +1028,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, &driver) .await?; assert_eq!(trailing_text(&request), Some("hi".to_string())); diff --git a/crates/libsy/src/algorithms/util/subagent.rs b/crates/libsy/src/algorithms/util/subagent.rs index 10447b65e..a7c9c5482 100644 --- a/crates/libsy/src/algorithms/util/subagent.rs +++ b/crates/libsy/src/algorithms/util/subagent.rs @@ -78,7 +78,7 @@ where &self, state: &mut S, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { if !request .metadata @@ -123,7 +123,7 @@ where &self, _state: &mut S, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)> { // Delegated *work* only. A harness maintenance turn (e.g. Codex `compact`) carries // sub-agent lineage but is not delegated work, so it abstains and routes normally. @@ -131,7 +131,7 @@ where .metadata .as_ref() .is_some_and(Metadata::is_subagent_work); - if is_delegated_work && let Some(driver) = driver { + if is_delegated_work { driver.set_evidence(serde_json::json!({"source": "subagent"})); } Ok(( @@ -139,6 +139,7 @@ where vec![Score { confidence: 1.0, target: self.worker.clone(), + category: None, }] } else { Vec::new() @@ -151,6 +152,7 @@ where #[cfg(test)] mod tests { use super::*; + use crate::core::testing::empty_driver; use parking_lot::Mutex; use switchyard_protocol::{slice_to_header_map, text_request}; @@ -165,13 +167,14 @@ mod tests { &self, _state: &mut (), request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { self.requests.lock().push(request.clone()); Ok(( Classification::Scores(vec![Score { confidence: 1.0, target: ModelId::from("worker"), + category: None, }]), None, )) @@ -192,7 +195,7 @@ mod tests { async fn selected(headers: &[(&str, &str)]) -> Result> { let mut state = (); let classification = SubagentOverride::new("worker") - .score(&mut state, &mut request(headers), None) + .score(&mut state, &mut request(headers), &empty_driver()) .await?; Ok(classification.0.argmax(false)?.map(|score| score.target)) } @@ -243,7 +246,7 @@ mod tests { .score( &mut state, &mut request(&[("x-openai-subagent", "review")]), - None, + &empty_driver(), ) .await?; match classification.0 { @@ -264,7 +267,9 @@ mod tests { request.llm_request.messages = vec![Message::text(Role::Assistant, "no user prompt")]; let mut state = (); - let (classification, response) = gate.score(&mut state, &mut request, None).await?; + let (classification, response) = gate + .score(&mut state, &mut request, &empty_driver()) + .await?; assert!(classification.argmax(false)?.is_none()); assert!(response.is_none()); diff --git a/crates/libsy/src/algorithms/util/target_selector.rs b/crates/libsy/src/algorithms/util/target_selector.rs index 4e36b648a..76a50127f 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,80 +29,99 @@ 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 { + fn to_classification( + &self, + verdict: Option<&Self::Verdict>, + driver: &Driver, + ) -> Result { let target = verdict .and_then(|verdict| self.selector.resolve(verdict).ok()) .and_then(Value::as_str) - .and_then(|label| self.targets.get(label)); + .and_then(|label| label.parse::().ok()) + // The judge decides the turn; it is not somewhere to route it. + .filter(|category| *category != Category::Judge) + .and_then(|category| { + let target = driver.models_for(&category).first()?.clone(); + Some((category, target)) + }); match target { - Some(target) => Classification::Scores(vec![Score { - target: target.clone(), + Some((category, target)) => Ok(Classification::Scores(vec![Score { + target, confidence: 1.0, - }]), - None => Classification::Ambiguous(vec![]), + category: Some(category), + }])), + None => Ok(Classification::Ambiguous(vec![])), } } } #[cfg(test)] mod tests { + use crate::core::algorithm::RuntimeModels; use serde_json::json; + use std::sync::Arc; + use switchyard_protocol::ModelId; use super::*; use crate::Result; + fn driver() -> Driver { + Driver::new( + "test", + Arc::new(RuntimeModels::new( + [(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"))]), - )?; + fn a_missing_unknown_or_unavailable_target_abstains() -> Result<()> { + let policy = TargetSelectorPolicy::new("/target")?; + let driver = driver(); - assert_eq!( - policy - .to_classification(Some(&json!({"target": "unknown"}))) - .argmax(false)?, - None - ); - assert_eq!( - policy - .to_classification(Some(&json!({"reason": "missing"}))) - .argmax(false)?, - None - ); + for verdict in [ + json!({"target": "efficient"}), + json!({"target": "unknown"}), + json!({"reason": "missing"}), + ] { + assert_eq!( + policy + .to_classification(Some(&verdict), &driver)? + .argmax(false)?, + None + ); + } Ok(()) } #[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..e2ca8e488 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}; @@ -30,6 +33,65 @@ use crate::{DriverError, LibsyError, Result, observability}; /// `Arc` object-safe. pub type StepStream = Pin> + Send>>; +/// The models one algorithm run may use, grouped by [`Category`]. Within a +/// category they are ordered best-first. +/// +/// Delegated sub-agent work gets its own groups, reachable only through +/// [`Driver::for_subagent`]. Keeping them separate is what stops a sub-agent's +/// `capable` from resolving to the parent's, and stops the parent falling back +/// onto a model only its sub-agents were given. +/// +/// One run's driver clones all read the same value, so it is passed as +/// `Arc` rather than cloned per driver. +#[derive(Clone, Debug, Default)] +pub struct RuntimeModels { + /// When using subagents this is the parent agent category. + by_category: HashMap>, + subagent: Option>>, +} + +impl RuntimeModels { + /// The models available to the algorithm itself. + pub fn new(by_category: HashMap>) -> Self { + Self { + by_category, + subagent: None, + } + } + + /// Adds the groups used for delegated sub-agent work. + pub fn with_subagent(mut self, models: HashMap>) -> Self { + self.subagent = Some(models); + self + } + + /// The models in `category`, ordered best-first. + pub fn models_for(&self, category: &Category) -> &[ModelId] { + self.by_category.get(category).map_or(&[], Vec::as_slice) + } + + /// The models delegated sub-agent work uses for `category`, ordered best-first. + pub fn subagent_models_for(&self, category: &Category) -> &[ModelId] { + self.subagent + .as_ref() + .and_then(|models| models.get(category)) + .map_or(&[], Vec::as_slice) + } +} + +impl From>> for RuntimeModels { + fn from(by_category: HashMap>) -> Self { + Self::new(by_category) + } +} + +/// Which of a [`RuntimeModels`]' groups a driver reads. +#[derive(Clone, Copy)] +enum Scope { + Parent, + Subagent, +} + /// An offloaded model call, surfaced inside [`Step::CallModel`]. /// /// The host reads the public fields, performs (or delegates) the model call, and fulfills it @@ -121,16 +183,27 @@ 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>>, + + /// Every group this run may route over, shared by all driver clones. + models: Arc, + + /// Which of those groups this driver reads. + scope: Scope, } 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: Arc, + ) -> (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 +214,8 @@ impl Driver { step_tx, algorithm: algorithm.to_string(), evidence: Arc::new(Mutex::new(None)), + models, + scope: Scope::Parent, }, step_rx, ) @@ -216,6 +291,37 @@ impl Driver { result } + /// The available models for this category, typically ordered best-first. + pub fn models_for(&self, category: &Category) -> &[ModelId] { + match self.scope { + Scope::Parent => self.models.models_for(category), + Scope::Subagent => self.models.subagent_models_for(category), + } + } + + /// 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.as_str()), + }) + } + + /// A driver scoped to delegated sub-agent work: its categories are the + /// sub-agent's own, and the parent's are no longer reachable through it. + pub fn for_subagent(&self) -> Result { + if self.models.subagent.is_none() { + return Err(LibsyError::AlgorithmError { + message: "delegated work has no sub-agent models".to_string(), + }); + } + Ok(Self { + scope: Scope::Subagent, + ..self.clone() + }) + } + /// 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 +373,14 @@ pub enum Step { pub async fn drive( algorithm: Arc, request: Request, + models: Arc, 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 +523,8 @@ 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: Arc) -> 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 +669,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", Arc::new(RuntimeModels::default())); let first_driver = driver.clone(); let mut first = tokio::spawn(async move { first_driver @@ -619,7 +726,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", Arc::new(RuntimeModels::default())); let producer = tokio::spawn(async move { driver .call_model(request(), vec![ModelId::from("dropped")]) @@ -639,7 +746,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", Arc::new(RuntimeModels::default())); drop(step_rx); let result = driver .call_model(request(), vec![ModelId::from("closed")]) @@ -742,7 +849,8 @@ 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(), Arc::new(RuntimeModels::default())); tokio::pin!(stream); let mut saw_call = false; @@ -854,7 +962,8 @@ 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(), Arc::new(RuntimeModels::default())); tokio::pin!(stream); let mut saw_error = false; @@ -926,7 +1035,7 @@ mod tests { dropped: dropped.clone(), }); - let stream = algo.run_stream(request()); + let stream = algo.run_stream(request(), Arc::new(RuntimeModels::default())); started_rx .recv() .await @@ -963,7 +1072,7 @@ mod tests { } let algo: Arc = Arc::new(Panicky); - let stream = algo.run_stream(request()); + let stream = algo.run_stream(request(), Arc::new(RuntimeModels::default())); tokio::pin!(stream); let mut saw_error = false; diff --git a/crates/libsy/src/core/classifier.rs b/crates/libsy/src/core/classifier.rs index 41aa04903..6af9bcbe4 100644 --- a/crates/libsy/src/core/classifier.rs +++ b/crates/libsy/src/core/classifier.rs @@ -4,7 +4,7 @@ use crate::core::algorithm::Driver; use crate::{LibsyError, Result}; use async_trait::async_trait; -use switchyard_protocol::{ModelId, Request, Response}; +use switchyard_protocol::{Category, ModelId, Request, Response}; /// One classifier's recommendation of a routing `target`, with a `[0.0, 1.0]` confidence. #[derive(Debug, Clone, PartialEq)] @@ -13,6 +13,10 @@ pub struct Score { pub confidence: f64, /// The target (model / tier) being recommended. pub target: ModelId, + /// The category `target` was drawn from, when the classifier picked one. The rest of + /// that category is what the turn falls through on failure, so a decision made without + /// a category — an affinity replay, say — leaves this `None`. + pub category: Option, } /// A classifier's verdict for a request: a set of target [`Score`]s, flagged by how @@ -73,8 +77,7 @@ fn argmax(scores: &[Score]) -> Result> { pub trait Classifier: Send + Sync { /// Score the classifier's targets given the current state and request. /// - /// When present, `driver` lets a classifier offload model calls. It is `None` - /// when the classifier is evaluated outside an algorithm run. + /// `driver` lets a classifier inspect runtime models and offload model calls. /// /// `request` is borrowed mutably so a classifier may rewrite it in place — inject a /// system prompt, drop tools, compact history. The edit is not scoped to this call: @@ -85,13 +88,14 @@ pub trait Classifier: Send + Sync { &self, state: &mut S, request: &mut Request, - driver: Option<&Driver>, + driver: &Driver, ) -> Result<(Classification, Option)>; } #[cfg(test)] mod tests { use super::*; + use crate::core::testing::empty_driver; use switchyard_protocol::text_request; /// Terse `Score` builder for the assertions below. @@ -99,6 +103,7 @@ mod tests { Score { target: ModelId::from(target), confidence, + category: None, } } @@ -180,7 +185,7 @@ mod tests { &self, state: &mut bool, request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { *state = true; let target = request.model_id().unwrap_or(ModelId::from("auto")); @@ -188,6 +193,7 @@ mod tests { Classification::Scores(vec![Score { target, confidence: 1.0, + category: None, }]), None, )) @@ -202,9 +208,8 @@ mod tests { raw_request: None, metadata: None, }; - // A `None` driver is valid: the classifier scored without offloading a model call. let (classification, _) = RecordingClassifier - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, &empty_driver()) .await?; assert_eq!( classification.argmax(false)?.map(|s| s.target), @@ -223,13 +228,14 @@ mod tests { &self, _state: &mut (), request: &mut Request, - _driver: Option<&Driver>, + _driver: &Driver, ) -> Result<(Classification, Option)> { request.llm_request.model = Some("rewritten".to_string()); Ok(( Classification::Scores(vec![Score { target: ModelId::from("rewritten"), confidence: 1.0, + category: None, }]), None, )) @@ -246,7 +252,7 @@ mod tests { }; RewritingClassifier - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, &empty_driver()) .await?; // The rewrite outlives the call: later classifiers in the cascade score this value, diff --git a/crates/libsy/src/core/processor.rs b/crates/libsy/src/core/processor.rs index 04d38eda4..538f11157 100644 --- a/crates/libsy/src/core/processor.rs +++ b/crates/libsy/src/core/processor.rs @@ -4,7 +4,7 @@ use crate::Result; use crate::core::algorithm::Driver; use async_trait::async_trait; -use switchyard_protocol::{AggLlmResponse, ModelId, Request}; +use switchyard_protocol::{AggLlmResponse, Category, ModelId, Request}; /// An event observed by the algorithm. Events are consumed by [`Processor`] to mutate state. /// @@ -17,7 +17,7 @@ pub enum Event<'a> { /// The request, rewritable in place. request: &'a mut Request, /// Offered so a processor can consult a model before the cascade runs. - driver: Option<&'a Driver>, + driver: &'a Driver, }, /// A routing decision paired with the request that produced it. /// @@ -28,6 +28,12 @@ pub enum Event<'a> { request: &'a mut Request, /// The model selected for `request`. selected_model_id: &'a ModelId, + /// The category `selected_model_id` was drawn from, when the deciding + /// classifier picked one. `None` for a decision made without a category, + /// such as an affinity replay. + category: Option, + /// Offered so a processor can inspect the runtime model categories. + driver: &'a Driver, }, /// A buffered response received back from a model. ModelResponse(&'a AggLlmResponse), @@ -44,6 +50,7 @@ pub trait Processor: Send + Sync { #[cfg(test)] mod tests { use super::*; + use crate::core::testing::empty_driver; use std::collections::HashMap; use switchyard_protocol::{text_request, text_response}; @@ -95,7 +102,7 @@ mod tests { &mut state, Event::Request { request: &mut req, - driver: None, + driver: &empty_driver(), }, ) .await?; @@ -108,6 +115,8 @@ mod tests { Event::Decision { request: &mut req, selected_model_id: &selected_model_id, + category: None, + driver: &empty_driver(), }, ) .await?; @@ -129,7 +138,7 @@ mod tests { &mut state, Event::Request { request: &mut req, - driver: None, + driver: &empty_driver(), }, ) .await?; @@ -166,7 +175,7 @@ mod tests { &mut state, Event::Request { request: &mut req, - driver: None, + driver: &empty_driver(), }, ) .await?; diff --git a/crates/libsy/src/core/testing.rs b/crates/libsy/src/core/testing.rs index aeb570988..175db2ba7 100644 --- a/crates/libsy/src/core/testing.rs +++ b/crates/libsy/src/core/testing.rs @@ -13,15 +13,34 @@ //! 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::core::algorithm::{Algorithm, CallModel, Driver, RuntimeModels}; use crate::{LibsyError, Result}; +/// Builds one runtime model category for a test. +pub(crate) fn category_models( + category: Category, + names: &[&str], +) -> HashMap> { + [( + category, + names.iter().map(|name| ModelId::from(*name)).collect(), + )] + .into() +} + +pub(crate) fn empty_driver() -> Driver { + Driver::new("test", Arc::new(RuntimeModels::default())).0 +} + /// The result a fake client hands back for one offloaded call. pub(crate) type ServeResult = std::result::Result; @@ -46,10 +65,20 @@ pub(crate) async fn test_drive( algorithm: Arc, request: Request, serve: impl Serve, +) -> Result<(ModelId, Response)> { + test_drive_with_models(algorithm, request, RuntimeModels::default(), serve).await +} + +/// Drive one request with an explicit runtime model set. +pub(crate) async fn test_drive_with_models( + algorithm: Arc, + request: Request, + models: impl Into, + 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, Arc::new(models.into()), move |call| { fulfill(Arc::clone(&routing_serve), call) }) .await?; diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 7999f2428..6e84c456a 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -5,7 +5,9 @@ #![doc = include_str!("../README.md")] mod core; -pub use core::algorithm::{Algorithm, CallModel, Driver, RoutingOutcome, Step, StepStream, drive}; +pub use core::algorithm::{ + Algorithm, CallModel, Driver, RoutingOutcome, RuntimeModels, Step, StepStream, drive, +}; pub use core::classifier::{Classification, Classifier, Score}; pub use core::outcome_metadata::OutcomeMetadata; pub use core::processor::{Event, Processor}; @@ -31,7 +33,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, ToolSemantics, ToolSignals}; @@ -39,8 +41,8 @@ pub use algorithms::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSemantics, T // 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/src/algorithm.rs b/crates/prefill-router/src/algorithm.rs index f0979c96a..b0ed4efa9 100644 --- a/crates/prefill-router/src/algorithm.rs +++ b/crates/prefill-router/src/algorithm.rs @@ -114,13 +114,13 @@ impl Algorithm for PrefillRouterAlgo { async fn route( self: Arc, - _driver: Driver, + driver: Driver, mut request: Request, ) -> libsy::Result { let mut state = (); if let Some(target) = self .affinity - .score(&mut state, &mut request, None) + .score(&mut state, &mut request, &driver) .await? .0 .argmax(false) @@ -181,6 +181,8 @@ impl Algorithm for PrefillRouterAlgo { Event::Decision { request: &mut request, selected_model_id: &target, + category: None, + driver: &driver, }, ) .await?; diff --git a/crates/prefill-router/tests/unit/algorithm.rs b/crates/prefill-router/tests/unit/algorithm.rs index 9e71759f2..4209e3302 100644 --- a/crates/prefill-router/tests/unit/algorithm.rs +++ b/crates/prefill-router/tests/unit/algorithm.rs @@ -4,9 +4,9 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use libsy::{Algorithm, LibsyError}; +use libsy::{Algorithm, LibsyError, RuntimeModels}; use switchyard_protocol::{ - ContentBlock, LlmRequest, Message, ModelId, Request, Role, ToolResult, text_request, + Category, ContentBlock, LlmRequest, Message, ModelId, Request, Role, ToolResult, text_request, }; use crate::{PrefillForward, PrefillRouterAlgo, Result}; @@ -66,14 +66,19 @@ fn forward() -> ( } async fn selected(route: Arc, request: Request) -> libsy::Result { - let outcome = libsy::drive(route, request, |call| async move { - call.respond(Ok(switchyard_protocol::Response { - llm_response: switchyard_protocol::LlmResponse::Agg( - switchyard_protocol::text_response(None, "unused"), - ), - metadata: None, - })) - }) + let outcome = libsy::drive( + route, + request, + Arc::new(RuntimeModels::new([(Category::Any, target_set())].into())), + |call| async move { + call.respond(Ok(switchyard_protocol::Response { + llm_response: switchyard_protocol::LlmResponse::Agg( + switchyard_protocol::text_response(None, "unused"), + ), + metadata: None, + })) + }, + ) .await?; Ok(outcome.selected_model_id()?.to_string()) } diff --git a/crates/protocol/src/category.rs b/crates/protocol/src/category.rs new file mode 100644 index 000000000..3f228be13 --- /dev/null +++ b/crates/protocol/src/category.rs @@ -0,0 +1,56 @@ +// 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::str::FromStr; +use std::sync::Arc; + +/// A group of models +#[derive(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, + /// A deployment-defined group. Only a custom classifier's policy selects one, + /// and no algorithm ascribes meaning to the name. + Named(Arc), +} + +impl Category { + /// Returns the lowercase category name used in configuration. + pub fn as_str(&self) -> &str { + match self { + Self::Any => "any", + Self::Capable => "capable", + Self::Efficient => "efficient", + Self::Judge => "judge", + Self::Named(name) => name, + } + } +} + +impl FromStr for Category { + type Err = String; + + /// The four reserved names are matched first. Were `"capable"` allowed to + /// become a [`Category::Named`] there would be two keys that compare unequal + /// but print the same, and a lookup of [`Category::Capable`] would silently + /// miss the configured group. + fn from_str(s: &str) -> Result { + let c = match s { + "capable" => Self::Capable, + "efficient" => Self::Efficient, + "judge" => Self::Judge, + "any" => Self::Any, + "" => return Err("Category name cannot be empty".to_string()), + name => Self::Named(Arc::from(name)), + }; + 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..7c8ca677e 100644 --- a/crates/switchyard-nemo-relay-plugin/src/runtime.rs +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -728,6 +728,7 @@ fn identity_metadata(metadata: Option<&Metadata>) -> Json { #[cfg(test)] mod tests { use std::collections::{BTreeMap, HashMap}; + use switchyard_runner::RuntimeModels; use switchyard_llm_client::ClientRouter; use switchyard_protocol::{LlmClientError, LlmResponseStreamEvent, ModelId, Usage}; @@ -749,6 +750,7 @@ mod tests { None, None, Vec::new(), + RuntimeModels::default(), ); 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 cc7801a42..eea067f41 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -15,12 +15,12 @@ use switchyard_libsy::{ Algorithm, CallModel, ClassifierContractConfig, ClassifierResponseFormat, ClassifyTrigger, CustomClassifierConfig, CustomClassifierPolicy, EscalationJudgeConfig, HandoffNoteConfig, LibsyError as RustLibsyError, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, - PickerMode, Random, RoutingOutcome, StageRouter, StageRouterConfig, Step as RustStep, - StepStream, TaskClassifierConfig, ToolSemantics, + PickerMode, Random, RoutingOutcome, RuntimeModels, StageRouter, StageRouterConfig, + Step as RustStep, StepStream, TaskClassifierConfig, ToolSemantics, }; 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,29 @@ fn classify_trigger(session_affinity: bool) -> ClassifyTrigger { } } +fn parse_category(value: &str) -> PyResult { + value.parse().map_err(PyValueError::new_err) +} + +/// Convert one scope's `{group: [model_id]}` mapping into its typed form. +/// +/// `any`, `capable`, `efficient`, and `judge` are the groups algorithms reason +/// about. Any other key is a deployment-defined group, selectable by name only +/// by a custom classifier's policy. +fn category_models_from_python( + models: HashMap>, +) -> PyResult>> { + models + .into_iter() + .map(|(category, ids)| { + Ok(( + parse_category(&category)?, + ids.into_iter().map(ModelId::from).collect(), + )) + }) + .collect() +} + /// 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 +210,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 +221,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 +233,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 +318,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 +332,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 } } } @@ -666,15 +659,31 @@ struct PyAlgorithm { impl PyAlgorithm { /// Run the algorithm as routing-time model calls followed by one terminal outcome. /// + /// Models are supplied per request, not at algorithm construction, so the + /// same algorithm object can route against a different pool on every call. + /// + /// `models` maps a category name to the model ids in it, ordered best-first: + /// the algorithm picks the first entry and treats the rest as fallbacks. + /// `any`, `capable`, `efficient`, and `judge` are the categories algorithms + /// reason about. Any other name is a deployment-defined group that only a + /// custom classifier's policy can select by name. An empty category name + /// raises `ValueError`. + /// + /// `subagent_models` is the same mapping for the delegated sub-agent scope. + /// Algorithms that route delegated work read it instead of `models`; when it + /// is omitted that scope is empty. + /// /// `headers`, when given, is normalized into the request's correlation /// [`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, subagent_models=None, headers=None))] fn run_stream( &self, request: &Bound<'_, PyAny>, - headers: Option>, + models: HashMap>, + subagent_models: Option>>, + headers: Option>, ) -> PyResult { let headers = headers.as_ref().map(header_map_from_python).transpose()?; let request = Request { @@ -682,9 +691,14 @@ impl PyAlgorithm { raw_request: None, metadata: headers.map(|headers| Metadata::from_headers(&headers)), }; + let mut runtime_models = RuntimeModels::new(category_models_from_python(models)?); + if let Some(subagent_models) = subagent_models { + runtime_models = + runtime_models.with_subagent(category_models_from_python(subagent_models)?); + } 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, Arc::new(runtime_models)) }; Ok(PyRunStream { inner: Arc::new(Mutex::new(stream)), @@ -744,17 +758,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 +778,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 +799,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, @@ -821,8 +814,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, @@ -843,8 +834,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) { @@ -860,12 +849,8 @@ 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; if let Some(mut semantics) = tool_semantics { config.tool_semantics = ToolSemantics { observe: semantics.remove("observe").unwrap_or_default(), @@ -883,8 +868,8 @@ fn stage_router_algorithm( .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 21a8fa9a0..626c150af 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, BTreeSet, HashMap}; use std::error::Error; use std::fmt::{Display, Formatter}; use std::path::PathBuf; @@ -18,7 +18,7 @@ use libsy::{ ToolSemantics, }; use serde::Deserialize; -use switchyard_protocol::ModelId; +use switchyard_protocol::{Category, ModelId}; /// Error returned when an algorithm description cannot be constructed. #[derive(Debug)] @@ -98,6 +98,7 @@ enum LlmClassifierModeConfig { #[derive(Clone, Debug)] struct CapabilityClassifierRouteConfig { + classifier_target: String, strong_target: String, weak_target: String, base_threshold: f64, @@ -112,6 +113,7 @@ struct CapabilityClassifierRouteConfig { #[derive(Clone, Debug)] struct EscalationClassifierRouteConfig { + classifier_target: String, strong_target: String, weak_target: String, prompt: Option, @@ -122,9 +124,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, @@ -134,6 +135,89 @@ struct CustomClassifierRouteConfig { max_output_tokens: u64, } +/// Runtime model groups for a custom classifier, keyed by group name. +/// +/// `any` and `judge` are required; `capable` and `efficient` carry tier meaning +/// when present. Any other key is a deployment-defined group the policy may +/// select by name, which is what lets one route choose between more than two +/// models. Ordered so error messages and derived target lists do not depend on +/// hash iteration. +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(transparent)] +pub struct CategoryModelConfig(BTreeMap>); + +impl CategoryModelConfig { + fn validate(&self, route_name: &str, default_target: &Category) -> AlgorithmResult<()> { + for category in [Category::Any, Category::Judge] { + if self.get(&category).is_empty() { + return Err(AlgorithmConfigError::new(format!( + "llm_classifier route {route_name} models.{} must contain at least one target", + category.as_str() + ))); + } + } + if self.get(default_target).is_empty() { + return Err(AlgorithmConfigError::new(format!( + "llm_classifier route {route_name} models.{} must contain at least one target because it is the default_target", + default_target.as_str() + ))); + } + // Every group name parses, so a typo would otherwise build fine and then + // abstain on every request. `any` is the fallback pool the router checks + // the selected target against, so a group outside it can never be served. + let any = self.get(&Category::Any); + for (name, models) in &self.0 { + if name == Category::Judge.as_str() || name == Category::Any.as_str() { + continue; + } + if let Some(missing) = models.iter().find(|model| !any.contains(model)) { + return Err(AlgorithmConfigError::new(format!( + "llm_classifier route {route_name} models.{name} lists target {missing}, which must also appear in models.any" + ))); + } + } + Ok(()) + } + + fn get(&self, category: &Category) -> &[String] { + self.0.get(category.as_str()).map_or(&[], Vec::as_slice) + } + + /// Every configured target name, judge included. + fn all_names(&self) -> Vec<&str> { + Self::deduped(self.0.values().flatten()) + } + + /// The completion targets: every group except the judge's own candidates. + /// + /// `any` leads because it is the deployment's stated fallback order, and + /// callers read this order to pick a route's representative target. + fn routing_names(&self) -> Vec<&str> { + let others = self + .0 + .iter() + .filter(|(name, _)| { + *name != Category::Judge.as_str() && *name != Category::Any.as_str() + }) + .flat_map(|(_, models)| models); + Self::deduped(self.get(&Category::Any).iter().chain(others)) + } + + fn deduped<'a>(names: impl Iterator) -> Vec<&'a str> { + let mut seen = BTreeSet::new(); + names + .map(String::as_str) + .filter(|name| seen.insert(*name)) + .collect() + } + + fn groups(&self) -> impl Iterator)> + '_ { + self.0 + .iter() + .filter_map(|(name, models)| Some((name.parse::().ok()?, models.clone()))) + } +} + /// 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)] @@ -172,9 +256,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, @@ -201,18 +285,24 @@ impl SubagentRouteConfig { match self { Self::Passthrough { target } => vec![target], Self::LlmClassifier(classifier) => classifier - .targets - .iter() - .flatten() - .map(String::as_str) - .collect(), + .models + .as_ref() + .map(CategoryModelConfig::routing_names) + .unwrap_or_default(), } } - fn classifier_target_name(&self) -> Option<&str> { + fn judge_target_names(&self) -> Vec<&str> { match self { - Self::LlmClassifier(classifier) => Some(&classifier.classifier_target), - Self::Passthrough { .. } => None, + Self::Passthrough { .. } => Vec::new(), + Self::LlmClassifier(classifier) => classifier + .models + .as_ref() + .map(|models| models.get(&Category::Judge)) + .unwrap_or_default() + .iter() + .map(String::as_str) + .collect(), } } } @@ -440,32 +530,25 @@ impl AlgorithmSpec { } names } - Self::LlmClassifier { config, .. } => { - match config.mode.unwrap_or(if config.escalation.is_some() { - ClassifierMode::Escalation - } else { - ClassifierMode::Capability - }) { - ClassifierMode::Capability => config - .weak_target - .iter() - .chain(&config.strong_target) - .map(String::as_str) - .collect(), - ClassifierMode::Escalation => config - .strong_target - .iter() - .chain(&config.weak_target) - .map(String::as_str) - .collect(), - ClassifierMode::Custom => config - .targets - .iter() - .flatten() - .map(String::as_str) - .collect(), - } - } + Self::LlmClassifier { config, .. } => match config.classifier_mode() { + ClassifierMode::Capability => config + .weak_target + .iter() + .chain(&config.strong_target) + .map(String::as_str) + .collect(), + ClassifierMode::Escalation => config + .strong_target + .iter() + .chain(&config.weak_target) + .map(String::as_str) + .collect(), + ClassifierMode::Custom => config + .models + .as_ref() + .map(CategoryModelConfig::routing_names) + .unwrap_or_default(), + }, Self::StageRouter { tiers, subagents, .. } => { @@ -510,51 +593,137 @@ 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()); + // Custom mode names its judge in `models.judge`; the other two modes + // use the top-level `classifier_target`. + Self::LlmClassifier { config, .. } => { + if matches!(config.classifier_mode(), ClassifierMode::Custom) { + names.extend( + config + .models + .as_ref() + .map(|models| models.get(&Category::Judge)) + .unwrap_or_default() + .iter() + .map(String::as_str), + ); + } else { + 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), _ => {} } + // A sub-agent classifier calls its own judge, which is never a completion target. + if let Self::Passthrough { + subagents: Some(subagents), + .. + } + | Self::StageRouter { + subagents: Some(subagents), + .. + } + | Self::Composite { + subagents: Some(subagents), + .. + } = self + { + names.extend(subagents.judge_target_names()); + } names } + /// Target names grouped as the runtime [`Driver`](libsy::Driver) expects them. + pub(crate) fn runtime_model_names( + &self, + route_name: &str, + ) -> AlgorithmResult { + let parent = 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.validated_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::Auto { + capable_target, + efficient_target, + } => category_models([ + (Category::Capable, vec![capable_target.clone()]), + (Category::Efficient, vec![efficient_target.clone()]), + ( + Category::Any, + vec![capable_target.clone(), efficient_target.clone()], + ), + ]), + 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, + advisor_target, + .. + } => category_models([ + (Category::Efficient, vec![executor_target.clone()]), + (Category::Any, vec![executor_target.clone()]), + (Category::Judge, vec![advisor_target.clone()]), + ]), + }; + + let subagents = match self { + Self::Passthrough { subagents, .. } + | Self::StageRouter { subagents, .. } + | Self::Composite { subagents, .. } => subagents.as_ref(), + _ => None, + }; + // Sub-agent groups stay separate from the parent's. Merging them would let a + // sub-agent's `capable` resolve to the parent's, and would put models only the + // sub-agents were given into the parent's `Any` fallback list. + let subagent = subagents + .map(|subagents| subagent_runtime_model_names(subagents, route_name)) + .transpose()?; + Ok(RuntimeModelNames { parent, subagent }) + } + /// 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 { Self::LlmClassifier { config, .. } - if matches!( - config.mode.unwrap_or(if config.escalation.is_some() { - ClassifierMode::Escalation - } else { - ClassifierMode::Capability - }), - ClassifierMode::Escalation - ) => + if matches!(config.classifier_mode(), ClassifierMode::Escalation) => { Some(( config.weak_target.as_deref()?, @@ -586,8 +755,85 @@ impl AlgorithmSpec { build_algorithm(context, self, targets) } } + +/// One route's target names, grouped by category and by routing scope. +pub(crate) struct RuntimeModelNames { + /// Groups the algorithm itself routes over. + pub(crate) parent: HashMap>, + /// Groups delegated sub-agent work routes over, when the route has a `subagents` table. + pub(crate) subagent: Option>>, +} + +fn category_models( + entries: impl IntoIterator)>, +) -> HashMap> { + entries.into_iter().collect() +} + +fn custom_runtime_model_names(config: &CategoryModelConfig) -> HashMap> { + config.groups().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.validated_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 { + fn classifier_mode(&self) -> ClassifierMode { + let default = if self.escalation.is_some() { + ClassifierMode::Escalation + } else { + ClassifierMode::Capability + }; + self.mode.unwrap_or(default) + } + + fn validated_classifier_mode( + &self, + route_name: &str, + ) -> AlgorithmResult { let Self { classifier_target, mode, @@ -602,7 +848,7 @@ impl LlmClassifierRouteConfig { response_format_type, max_output_tokens, escalation, - targets, + models, default_target, response_schema, policy, @@ -626,13 +872,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", @@ -662,7 +909,7 @@ impl LlmClassifierRouteConfig { reject_custom_fields( route_name, "escalation", - targets, + models, default_target, response_schema, policy, @@ -684,6 +931,7 @@ impl LlmClassifierRouteConfig { } Ok(LlmClassifierModeConfig::Escalation( EscalationClassifierRouteConfig { + classifier_target: classifier_target.clone(), strong_target: required_classifier_field( route_name, "strong_target", @@ -702,7 +950,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() @@ -713,15 +962,28 @@ impl LlmClassifierRouteConfig { "llm_classifier route {route_name} mode custom cannot use capability or escalation fields and response_format_type must be 'json_schema'" ))); } + let models = required_classifier_field(route_name, "models", models)?; + let default_target: Category = 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}" + )) + })?; + if default_target == Category::Judge { + return Err(AlgorithmConfigError::new(format!( + "llm_classifier route {route_name} default_target cannot be judge" + ))); + } + models.validate(route_name, &default_target)?; Ok(LlmClassifierModeConfig::Custom( CustomClassifierRouteConfig { - classifier_target: classifier_target.clone(), - targets: required_classifier_field(route_name, "targets", targets)?, - default_target: required_classifier_field( - route_name, - "default_target", - default_target, - )?, + models, + default_target, prompt: required_classifier_field(route_name, "prompt", prompt)?, response_schema: required_classifier_field( route_name, @@ -743,15 +1005,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" @@ -784,36 +1043,27 @@ fn build_subagent_router_config( targets: &BTreeMap, ) -> AlgorithmResult { match config { - SubagentRouteConfig::Passthrough { target } => Ok(SubagentRouterConfig::fixed_target( - resolve_target_model_id(route_name, target, targets)?, - )), + SubagentRouteConfig::Passthrough { target } => { + resolve_target_model_id(route_name, target, targets)?; + Ok(SubagentRouterConfig::fixed_target()) + } SubagentRouteConfig::LlmClassifier(config) => { - let LlmClassifierModeConfig::Custom(config) = config.classifier_mode(route_name)? + let LlmClassifierModeConfig::Custom(config) = + config.validated_classifier_mode(route_name)? else { return Err(AlgorithmConfigError::new(format!( "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)) - }) - .collect::>>()?; - let default_target = resolved_targets - .iter() - .find(|(name, _)| *name == config.default_target) - .map(|(_, target)| target.clone()) - .ok_or_else(|| { - AlgorithmConfigError::new(format!( - "route {route_name}: subagents llm_classifier default_target {:?} must be one of its configured targets", - config.default_target - )) - })?; + for name in config.models.all_names() { + resolve_target_model_id(route_name, name, targets)?; + } + if config.models.get(&config.default_target).is_empty() { + return Err(AlgorithmConfigError::new(format!( + "route {route_name}: subagents llm_classifier has no model for default category {}", + config.default_target.as_str() + ))); + } let response_schema = serde_json::from_str(&config.response_schema).map_err(|error| { AlgorithmConfigError::with_source( @@ -830,15 +1080,9 @@ 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, + default_target: config.default_target.clone(), config: classifier_config, }) .map_err(|error| { @@ -849,9 +1093,8 @@ fn build_subagent_router_config( })?, ); Ok(SubagentRouterConfig { - targets: subagent_targets, classifier, - default_target, + default_target: config.default_target, classify_trigger: config.classify_trigger, message_hash_fallback: config.message_hash_fallback, }) @@ -889,11 +1132,25 @@ fn build_algorithm( 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| { + // The algorithm only sees its targets at request time, so a bad pairing + // would otherwise fail every request instead of failing to start. + if let Some(weights) = weights + && weights.len() != names.len() + { + return Err(AlgorithmConfigError::new(format!( + "random route {route_name}: expected {} weights, got {}", + names.len(), + weights.len() + ))); + } + let mut seen = BTreeSet::new(); + if let Some(duplicate) = names.iter().find(|name| !seen.insert(*name)) { + return Err(AlgorithmConfigError::new(format!( + "random route {route_name}: targets must be unique, {duplicate} is repeated" + ))); + } + let algorithm = Random::new(weights.clone(), *seed).map_err(|error| { AlgorithmConfigError::with_source( format!("random route {route_name}: {error}"), error, @@ -901,11 +1158,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; let parent: Arc = Arc::new(algorithm); attach_subagent_router(route_name, parent, subagents.as_ref(), targets) } @@ -913,14 +1167,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 mode = classifier_config.validated_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, @@ -932,20 +1181,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, @@ -953,14 +1193,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( @@ -981,8 +1213,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, }) @@ -1004,38 +1234,27 @@ fn build_algorithm( .. } => { let StageTierConfig { - capable_target, - efficient_target, confidence_threshold, recent_turn_window, tool_semantics, 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.tool_semantics = tool_semantics.clone(); 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, @@ -1044,14 +1263,9 @@ fn build_algorithm( let parent: Arc = Arc::new(algorithm); attach_subagent_router(route_name, parent, subagents.as_ref(), targets) } - AlgorithmSpec::Auto { - capable_target, - efficient_target, - } => { - let capable = resolve_target_model_id(route_name, capable_target, targets)?; - let efficient = resolve_target_model_id(route_name, efficient_target, targets)?; + AlgorithmSpec::Auto { .. } => { let config = StageRouterConfig::new(PickerMode::EfficientFirst, 0.5); - let algorithm = StageRouter::new(capable, efficient, config).map_err(|error| { + let algorithm = StageRouter::new(config).map_err(|error| { AlgorithmConfigError::with_source( format!("auto route {route_name}: {error}"), error, @@ -1064,20 +1278,16 @@ 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.tool_semantics = stage.tool_semantics.clone(); 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, @@ -1087,8 +1297,6 @@ fn build_algorithm( attach_subagent_router(route_name, parent, subagents.as_ref(), targets) } AlgorithmSpec::Advisor { - executor_target, - advisor_target, reviewer_system_prompt, redo_feedback_prefix, gate_trigger, @@ -1102,8 +1310,6 @@ fn build_algorithm( fail_open, .. } => { - let executor = resolve_target_model_id(route_name, executor_target, targets)?; - let advisor = resolve_target_model_id(route_name, advisor_target, targets)?; // A pattern set under the default trigger would be silently // ignored; reject the misconfiguration instead. if *gate_trigger == AdvisorTriggerConfig::NoToolCall && gate_trigger_pattern.is_some() { @@ -1132,7 +1338,7 @@ fn build_algorithm( config.advisor_temperature = *advisor_temperature; config.transcript_max_chars = *transcript_max_chars; config.fail_open = *fail_open; - let algorithm = AdvisorGate::new(executor, advisor, config).map_err(|error| { + let algorithm = AdvisorGate::new(config).map_err(|error| { AlgorithmConfigError::with_source( format!("advisor route {route_name}: {error}"), error, @@ -1150,8 +1356,10 @@ fn build_algorithm( } => { #[cfg(feature = "prefill-router")] { - let targets = - resolve_targets(route_name, names.iter().map(String::as_str), targets)?; + let targets = names + .iter() + .map(|name| resolve_target_model_id(route_name, name, targets)) + .collect::>>()?; let mut config = prefill_router::PrefillRouterConfig::new(targets, checkpoint); config.device.clone_from(device); config.cache_dir.clone_from(cache_dir); @@ -1209,17 +1417,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 867a8e53a..316754723 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -8,6 +8,7 @@ use std::fs; use std::path::Path; use std::sync::Arc; +use libsy::RuntimeModels; use serde::de::DeserializeOwned; use serde::{Deserialize, Deserializer}; use serde_json::Value; @@ -15,7 +16,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, @@ -227,6 +228,14 @@ impl DeploymentConfig { .into_iter() .filter_map(|name| self.decision_target(name)) .collect(); + let names = config + .algorithm + .runtime_model_names(route_name) + .map_err(|error| RunnerError::configuration_source(error.to_string(), error))?; + let mut models = RuntimeModels::new(resolve_category_models(names.parent, &targets)?); + if let Some(subagent) = names.subagent { + models = models.with_subagent(resolve_category_models(subagent, &targets)?); + } let route = Route::new( algorithm, route_clients, @@ -235,6 +244,7 @@ impl DeploymentConfig { anthropic_auxiliary_target, responses_auxiliary_target, decision_targets, + models, ); routes.push((config.id.clone(), route)); } @@ -556,6 +566,30 @@ impl ClientFormat { } } } + +/// Resolves one scope's configured target names to the models the driver serves. +fn resolve_category_models( + names: HashMap>, + targets: &BTreeMap, +) -> RunnerResult>> { + names + .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() +} + fn build_backend( client_name: &str, config: &LlmClientConfig, @@ -739,7 +773,26 @@ 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.models_for(&Category::Judge), + [ModelId::from("classifier/model")] + ); + assert_eq!( + models.models_for(&Category::Efficient), + [ModelId::from("weak/model")] + ); + assert_eq!( + models.models_for(&Category::Capable), + [ModelId::from("strong/model")] + ); + assert_eq!( + models.models_for(&Category::Any), + [ModelId::from("weak/model"), ModelId::from("strong/model")] + ); assert!(runner.route("switchyard/passthrough").is_some()); Ok(()) } @@ -757,11 +810,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""#, ); @@ -769,6 +821,28 @@ classify_trigger = "new_session""#, configured } + #[test] + fn subagent_models_stay_separate_from_the_parent_tiers() -> RunnerResult<()> { + // The sub-agent target is also the parent's capable tier. Merged into one group it + // would be indistinguishable from that tier, and delegated work would follow the + // parent's ordering instead of its own configured target. + let runner = runner_from_toml(&with_subagent_passthrough(&stage_config(), "stage"))?; + let models = runner + .route("switchyard/stage") + .expect("stage route should exist") + .models(); + + assert_eq!( + models.subagent_models_for(&Category::Any), + [ModelId::from("strong/model")] + ); + assert_eq!( + models.models_for(&Category::Any), + [ModelId::from("strong/model"), ModelId::from("weak/model")] + ); + Ok(()) + } + fn with_subagent_passthrough(config: &str, route: &str) -> String { format!("{config}\n[routes.{route}.subagents]\ntype = \"passthrough\"\ntarget = \"strong\"") } @@ -1202,23 +1276,23 @@ classifier_magic = true ( VALID_CONFIG.replace( "targets = [\"strong\", \"weak\"]", - "targets = [\"strong\", \"strong\"]", + "targets = [\"strong\", \"weak\"]\nweights = [0, 0]", ), - "random targets must be unique", + "at least one weight must be positive", ), ( VALID_CONFIG.replace( "targets = [\"strong\", \"weak\"]", - "targets = [\"strong\", \"weak\"]\nweights = [1]", + "targets = [\"strong\", \"strong\"]", ), - "expected 2 weights, got 1", + "targets must be unique, strong is repeated", ), ( VALID_CONFIG.replace( "targets = [\"strong\", \"weak\"]", - "targets = [\"strong\", \"weak\"]\nweights = [0, 0]", + "targets = [\"strong\", \"weak\"]\nweights = [1]", ), - "at least one weight must be positive", + "expected 2 weights, got 1", ), ( VALID_CONFIG.replace("base_threshold = 0.5", "base_threshold = 1.5"), @@ -1653,6 +1727,24 @@ advisor_target = "advisor" #[test] fn advisor_route_parses_with_defaults_and_builds() -> RunnerResult<()> { let state = runner_from_toml(ADVISOR_CONFIG)?; + let route = state + .route("switchyard/advisor") + .expect("advisor route should exist"); + let models = route.models(); + // The gate calls the executor through `efficient`; `any` keeps it in the + // route's last-resort pool. + assert_eq!( + models.models_for(&Category::Efficient), + [ModelId::from("executor/model")] + ); + assert_eq!( + models.models_for(&Category::Any), + [ModelId::from("executor/model")] + ); + assert_eq!( + models.models_for(&Category::Judge), + [ModelId::from("advisor/model")] + ); assert_eq!( state .models() diff --git a/crates/switchyard-runner/src/lib.rs b/crates/switchyard-runner/src/lib.rs index efb0441e4..175eb0204 100644 --- a/crates/switchyard-runner/src/lib.rs +++ b/crates/switchyard-runner/src/lib.rs @@ -10,10 +10,12 @@ 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}; +// Re-exported because `Route::new` takes it, so a host wiring routes does not need a libsy dep. +pub use libsy::RuntimeModels; pub use route::{ AuxiliaryTarget, CallerAuthKind, ModelCapabilities, Route, RunOutput, RunnerError, }; diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index 89f6e722c..2b594646f 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -6,7 +6,7 @@ use std::error::Error; use std::sync::Arc; -use libsy::{Algorithm, LibsyError, RoutingOutcome}; +use libsy::{Algorithm, LibsyError, RoutingOutcome, RuntimeModels}; use serde_json::Value; use switchyard_llm_client::{AuxiliaryOperation, ClientRouter, RunObserver, TranslatingLlmClient}; use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat}; @@ -125,6 +125,7 @@ pub struct Route { anthropic_auxiliary_target: Option, responses_auxiliary_target: Option, decision_targets: Vec, + models: Arc, } /// 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: RuntimeModels, ) -> Self { Self { algorithm, @@ -152,6 +155,7 @@ impl Route { anthropic_auxiliary_target, responses_auxiliary_target, decision_targets, + models: Arc::new(models), } } @@ -178,6 +182,11 @@ impl Route { .cloned() } + /// Returns the models grouped for one algorithm execution. + pub fn models(&self) -> &RuntimeModels { + &self.models + } + /// 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 @@ -198,6 +207,7 @@ impl Route { Arc::clone(&self.algorithm), self.clients.clone(), request, + Arc::clone(&self.models), observer, ) .await?; @@ -209,9 +219,14 @@ 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) + switchyard_llm_client::decide( + Arc::clone(&self.algorithm), + self.clients.clone(), + request, + Arc::clone(&self.models), + ) + .await + .map_err(Into::into) } /// Executes a model-bearing provider operation through a compatible target. diff --git a/crates/switchyard-runner/tests/route.rs b/crates/switchyard-runner/tests/route.rs index df03c147d..947030965 100644 --- a/crates/switchyard-runner/tests/route.rs +++ b/crates/switchyard-runner/tests/route.rs @@ -7,10 +7,11 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use futures_util::StreamExt; +use libsy::RuntimeModels; 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 +55,7 @@ fn plugin_route(client: Arc) -> Route { None, None, Vec::new(), + RuntimeModels::new([(Category::Any, vec![ModelId::from("semantic-target")])].into()), ) } @@ -69,7 +71,6 @@ 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 output = route .execute(request, Some(observer)) .await @@ -145,25 +146,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 69fc5e4ff..51c91f601 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -33,11 +33,11 @@ use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Extension, Json, Router}; use axum_server::tls_rustls::RustlsConfig; -use libsy::{Algorithm, LibsyError, RoutingOutcome}; +use libsy::{LibsyError, RoutingOutcome}; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; -use switchyard_llm_client::{AuxiliaryOperation, ClientRouter, RunObservation, RunObserver}; +use switchyard_llm_client::{AuxiliaryOperation, RunObservation, RunObserver}; use switchyard_protocol::{LlmClientError, Metadata, ModelId, Request, Usage}; use switchyard_runner::{ CallerAuthKind, DecisionTarget, ModelCapabilities, Route, RunOutput, Runner, RunnerError, @@ -179,29 +179,6 @@ impl SharedRoutingLog { } impl ServerState { - /// Creates server state from route model IDs, algorithms, and per-target clients. - pub fn new(routes: Vec<(ModelId, Arc, ClientRouter)>) -> ServerResult { - let routes = routes - .into_iter() - .map(|(model, algorithm, clients)| { - ( - model, - Route::new( - algorithm, - clients, - None, - ModelCapabilities::default(), - None, - None, - Vec::new(), - ), - ) - }) - .collect(); - let runner = Runner::new(routes); - Self::from_runner(runner) - } - /// Creates HTTP-server state around an already configured runner. pub fn from_runner(runner: Runner) -> ServerResult { let metrics = metrics::registry().map_err(ServerError::new)?; @@ -725,6 +702,7 @@ async fn decision( .as_deref() .map(ModelId::from) .unwrap_or_default(); + let mut outcome = match route.decide(request).await { Ok(outcome) => outcome, Err(error) => return runner_error(error), @@ -1031,6 +1009,7 @@ 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 { 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 19adfb1fb..133718317 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, RuntimeModels}; use switchyard_server::config::load_server_state; use switchyard_server::{ DEFAULT_MAX_REQUEST_BODY_BYTES, ServerState, build_llm_router, build_switchyard_router, @@ -301,23 +302,34 @@ async fn upstream_chat( .is_some_and(|content| content.contains("schema-invalid verdict")) }) }); + // A custom-mode task may name the group it wants the judge to pick, so one + // config can be driven through each of its groups in turn. + let requested_group = body["messages"].as_array().and_then(|messages| { + messages.iter().find_map(|message| { + message["content"] + .as_str()? + .split_once("route to ") + .map(|(_, group)| group.trim().to_string()) + }) + }); let content = if model == "model/classifier" && custom_target_schema { if requests_invalid_verdict { - r#"{"decision":{"target":"unknown"}}"# + r#"{"decision":{"target":"unknown"}}"#.to_string() } else { - r#"{"decision":{"target":"premium"}}"# + let group = requested_group.unwrap_or_else(|| "efficient".to_string()); + format!(r#"{{"decision":{{"target":"{group}"}}}}"#) } } else if body .pointer("/response_format/json_schema/schema/properties/escalate") .is_some() { - r#"{"escalate":false,"reason":"making progress"}"# + r#"{"escalate":false,"reason":"making progress"}"#.to_string() } else if model == "model/classifier" && requests_schema_invalid_verdict { - r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.1,"unexpected":true}"# + r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.1,"unexpected":true}"#.to_string() } else if model == "model/classifier" { - r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"# + r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string() } else { - "ok" + "ok".to_string() }; Json(json!({ "id": "chatcmpl-test", @@ -528,16 +540,39 @@ 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, + RuntimeModels::new( + [( + 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)> { @@ -1583,8 +1618,7 @@ new = ["send_message_to_user"] } #[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#" @@ -1618,9 +1652,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"], fast = ["weak"], balanced = ["middle"], reasoning = ["strong"], premium = ["premium"], any = ["weak", "middle", "strong", "premium"] }} +default_target = "premium" prompt = "CUSTOM MULTI TARGET" response_schema = ''' {{ @@ -1629,7 +1662,7 @@ response_schema = ''' "decision": {{ "type": "object", "properties": {{ - "target": {{"type": "string", "enum": ["weak", "middle", "strong", "premium"]}} + "target": {{"type": "string", "enum": ["fast", "balanced", "reasoning", "premium"]}} }}, "required": ["target"], "additionalProperties": false @@ -1648,9 +1681,14 @@ selector = "/decision/target" ))?; let app = build_switchyard_router(state); + // Each named group resolves to its own model, so the policy picks between + // four of them rather than between the two tier categories. for (task, selected) in [ - ("route this task", "model/premium"), - ("return an invalid verdict", "model/strong"), + ("route to fast", "model/weak"), + ("route to balanced", "model/middle"), + ("route to reasoning", "model/strong"), + ("route to premium", "model/premium"), + ("return an invalid verdict", "model/premium"), ] { let response = send( &app, @@ -1690,7 +1728,7 @@ selector = "/decision/target" assert_eq!( judge_call["response_format"]["json_schema"]["schema"]["properties"]["decision"]["properties"] ["target"]["enum"], - json!(["weak", "middle", "strong", "premium"]) + json!(["fast", "balanced", "reasoning", "premium"]) ); Ok(()) } diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index 92733bd57..7d4b8b914 100644 --- a/docs/reference/toml_schema.md +++ b/docs/reference/toml_schema.md @@ -181,12 +181,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. | @@ -219,12 +219,23 @@ 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 group. A verdict names a group and +the first model in it serves the turn; if that call fails the client falls through +the rest of that group, then through whatever `models.any` adds. + +The `[routes..models]` table takes any group name you choose. `any` and +`judge` are reserved and required; `capable` and `efficient` are reserved for the +tier meaning the other algorithms give them. Every other key is yours, which is +how one route chooses between more than two models. | 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 | — | Every selectable completion target, in last-resort fallback order. Every other group's targets must also appear here; one that does not is rejected at configuration load. | +| `models.judge` | Yes | — | One or more ordered judge candidates. Not a completion destination. | +| `models.capable` | No | — | Ordered capable-tier models. A `capable` verdict selects the first and falls through the rest in order. | +| `models.efficient` | No | — | Ordered efficient-tier models. An `efficient` verdict selects the first and falls through the rest in order. | +| `models.` | No | — | A group you name. A verdict naming it selects its first model and falls through the rest in order. | +| `default_target` | Yes | — | Group used when the judge fails or its verdict cannot be routed. Any group except `judge`, and it must contain at least one target. | | `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`. | @@ -232,6 +243,10 @@ 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 name a configured group. A label naming a target +rather than a group, or a group you did not configure, falls back to +`default_target`. `judge` is not routable. + 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..cbc893563 100644 --- a/docs/routing_algorithms/llm_classifier_routing.md +++ b/docs/routing_algorithms/llm_classifier_routing.md @@ -139,18 +139,17 @@ 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 the route's model groups, and you name those +groups yourself, so a route can choose between as many models as you like. ```toml [routes.smart] id = "smart" type = "llm_classifier" mode = "custom" -classifier_target = "classifier" -targets = ["fast", "balanced", "reasoning", "premium"] default_target = "premium" prompt = """ -Choose the best configured target for this request. +Choose the best group for this request. Return JSON matching the response schema supplied with the request. """ response_schema = ''' @@ -174,15 +173,36 @@ response_schema = ''' } ''' +[routes.smart.models] +judge = ["classifier"] +fast = ["fast"] +balanced = ["balanced"] +reasoning = ["reasoning", "premium"] +premium = ["premium"] +any = ["fast", "balanced", "reasoning", "premium"] + [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`. +missing, non-string, or unconfigured label falls back to `default_target`, and +`judge` is never routable. + +A verdict names a group, and the **first** model in that group serves the turn. +Later entries are that group's own fallbacks: if the serving call fails, the +client falls through the rest of the chosen group first — `reasoning` retries on +`premium` above — and then through whatever `models.any` adds. Every group's +targets must also appear in `models.any`; one that does not is rejected when the +configuration loads. `models.judge` supplies the judge call's own candidates in +order and is not a completion destination. + +`capable` and `efficient` are reserved names. Use them when you want a group to +carry the tier meaning the stage and composite routers give it; otherwise any +name works. 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 ab201f2a0..b3bdcd9f3 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 @@ -74,6 +73,11 @@ switchyard-server --config routes.toml --dry-run switchyard-server --config routes.toml --port 4000 ``` +The `subagents` table has its own model groups, separate from the parent +route's. A category name in the sub-agent table always means the sub-agent's own +models, even when the parent route uses that category too, and the parent never +falls back onto a model only the sub-agents were given. + The parent always uses `parent`. For a delegated request, the classifier sees the prompt supplied by the parent and selects one configured target. With `classify_trigger = "new_session"`, Switchyard reuses that decision for later 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..b2cd21001 100644 --- a/examples/litellm/src/switchyard_litellm/plugins/switchyard_routing_plugin.py +++ b/examples/litellm/src/switchyard_litellm/plugins/switchyard_routing_plugin.py @@ -157,18 +157,48 @@ class SwitchyardRoutingPlugin(LiteLLMRequestRewriter): Compatible algorithms must finish without intermediate model calls or an already-produced response. Supported request rewrites are carried to the selected deployment by the object's LiteLLM callback role. + + `models` maps a category name (`any`, `capable`, `efficient`, `judge`, or a + deployment-defined group) to its LiteLLM model names, ordered best-first. + The algorithm picks the first name it can use from a category, so the order + within each list is the preference order. Every name must be a LiteLLM + deployment that can appear in the routing candidate pool: the plugin keeps + only the selected candidate, and a selection outside the pool LiteLLM + offered for this request raises `ValueError`. + + When `models` is omitted, each of the four categories is filled with the + request's full candidate pool, so any selection is valid by construction. """ - 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 5acbac3d4..fd768ba71 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 groups. ``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,23 @@ 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 model groups using a schema-selected label. + + ``default_target`` names the group used when the judge fails or its + verdict cannot be routed. Alongside ``capable`` and ``efficient`` a + deployment may define its own group names, so one route can choose + between more than two models. + """ ... @final class LlmFallback: def __init__( self, - judge_target: str, *, config: TaskClassifierConfig, ) -> None: ... @@ -225,14 +222,14 @@ class Algorithm: def run_stream( self, request: Mapping[str, object], + models: Mapping[str, Sequence[str]], + subagent_models: Mapping[str, Sequence[str]] | None = None, 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 +239,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 4d1fb78e6..50c05665d 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,49 +415,58 @@ 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."}], + } + ] async def test_stage_router_applies_additive_tool_semantics() -> None: """Verify that configured mutation semantics reach stage-router scoring.""" algorithm = algorithms.stage_router( - "strong", - "fast", picker="capable_first", confidence_threshold=0.3, tool_semantics={ @@ -457,6 +480,11 @@ async def test_stage_router_applies_additive_tool_semantics() -> None: selected_model, _ = await run_algorithm( algorithm, {"strong": EchoClient("strong"), "fast": EchoClient("fast")}, + models={ + "efficient": ["fast"], + "capable": ["strong"], + "any": ["fast", "strong"], + }, request={ "model": "auto", "messages": [ @@ -496,8 +524,6 @@ async def test_stage_router_applies_additive_tool_semantics() -> None: def test_stage_router_rejects_unknown_tool_semantics_category() -> None: with pytest.raises(ValueError, match="unknown tool_semantics category"): algorithms.stage_router( - "strong", - "fast", picker="efficient_first", confidence_threshold=0.5, tool_semantics={"complete": ["end_conversation"]},