switchyard_libsy/algorithms/
composite.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Routing that stacks a judge over a stage router.
5//!
6//! The judge runs as a [`Processor`]: it sets configuration the stage router reads,
7//! and picks no target itself.
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use parking_lot::Mutex;
14
15use super::fall_through::FallThrough;
16use super::llm_class::{
17    CustomClassifierConfig, LlmClassifierConfig, LlmTaskClassifier, TaskClassifierConfig,
18};
19use super::stage::{StageRouterConfig, build_stage_route};
20use super::util::affinity::{ClassifyTrigger, evict_if_full, has_new_user_turn, retention_key};
21use super::util::stage::{Tier, fall_open_tier, set_fall_open};
22use crate::core::algorithm::{Algorithm, Driver, RoutingIdentity};
23use crate::core::classifier::Classifier;
24use crate::core::processor::{Event, Processor};
25use crate::core::state::State;
26use crate::{LibsyError, Result};
27use switchyard_protocol::{Category, Request};
28
29const COMPOSITE: &str = "composite";
30
31/// Sets the stage router's fall-open tier from a judge verdict.
32///
33/// Retains the tier per routing identity, so it survives requests that carry no
34/// session ID when `message_hash_fallback` is on. The retained tier is replayed
35/// into state on every request so the cascade below reads it.
36struct TierSetter {
37    judge: Arc<dyn Classifier>,
38    trigger: ClassifyTrigger,
39    message_hash_fallback: bool,
40    default_tier: Tier,
41    tiers: Mutex<HashMap<RoutingIdentity, Tier>>,
42}
43
44impl TierSetter {
45    fn tier_for(category: Option<Category>) -> Option<Tier> {
46        match category {
47            Some(Category::Capable) => Some(Tier::Capable),
48            Some(Category::Efficient) => Some(Tier::Efficient),
49            _ => None,
50        }
51    }
52
53    /// Two requests for one identity can both pass this and both judge, since a
54    /// judge call sits between here and [`retain`](Self::retain). The later wins.
55    fn is_due(&self, identity: Option<&RoutingIdentity>, state: &State, request: &Request) -> bool {
56        match self.trigger {
57            ClassifyTrigger::UserTurn => has_new_user_turn(&request.llm_request.messages),
58            // Unkeyed requests with no carried tier cannot be told apart, so every one
59            // is a new session.
60            ClassifyTrigger::NewSession => {
61                identity.is_none_or(|identity| !self.tiers.lock().contains_key(identity))
62                    && fall_open_tier(state).is_none()
63            }
64            // Rejected by the constructor, and only in the enum for the standalone route.
65            ClassifyTrigger::EveryRequest => true,
66        }
67    }
68
69    fn retain(&self, identity: RoutingIdentity, tier: Tier) {
70        let mut tiers = self.tiers.lock();
71        // A user turn re-decides, so it overwrites. A session keeps its first verdict,
72        // matching how affinity retains an assignment.
73        let writable = self.trigger == ClassifyTrigger::UserTurn || !tiers.contains_key(&identity);
74        if writable {
75            evict_if_full(&mut tiers);
76            tiers.insert(identity, tier);
77        }
78    }
79}
80
81#[async_trait]
82impl Processor for TierSetter {
83    async fn process(&self, state: &mut State, event: Event<'_>) -> Result<()> {
84        let Event::Request { request, driver } = event else {
85            return Ok(());
86        };
87        let identity = retention_key(request, self.message_hash_fallback);
88        if self.is_due(identity.as_ref(), state, request) {
89            let (classification, _) = self.judge.score(state, request, driver).await?;
90            if let Some(winner) = classification.argmax(false)?
91                && let Some(tier) = Self::tier_for(winner.category)
92            {
93                set_fall_open(state, tier);
94                if let Some(identity) = identity {
95                    self.retain(identity, tier);
96                }
97                return Ok(());
98            }
99            if self.trigger == ClassifyTrigger::NewSession
100                && let Some(identity) = identity.as_ref()
101            {
102                // A failed opening judgment is still the session's one judgment.
103                self.retain(identity.clone(), self.default_tier);
104            }
105        }
106        // Either not this request's turn or the judge had no verdict, so the last
107        // tier stands rather than dropping back to the picker default.
108        if let Some(tier) = identity
109            .and_then(|identity| self.tiers.lock().get(&identity).copied())
110            .or_else(|| fall_open_tier(state))
111        {
112            set_fall_open(state, tier);
113            driver.set_evidence_if_empty(serde_json::json!({"source": "retained"}));
114        }
115        Ok(())
116    }
117}
118
119/// Classifier settings for the composite router's retained tier.
120pub enum CompositeJudgeConfig {
121    /// Selects a tier from a capability verdict.
122    Capability(TaskClassifierConfig),
123    /// Selects a tier using a custom schema and policy, retaining the fallback on failure.
124    Custom(CustomClassifierConfig),
125}
126
127/// A judge stacked over a stage router.
128pub struct CompositeRouterConfig {
129    /// Judge settings, including how often `classify_trigger` runs it.
130    pub judge: CompositeJudgeConfig,
131    /// Serves the turns, with the tier the judge picked as its fall-open default.
132    pub stage: StageRouterConfig,
133}
134
135/// Runs a stage router with a tier the judge picks.
136pub struct CompositeRouter {
137    route: FallThrough,
138}
139
140impl CompositeRouter {
141    /// Stacks the judge over a stage router across the same tier pair.
142    ///
143    /// Errors on a configuration either algorithm rejects and on `every_request`.
144    ///
145    /// A stage router carrying its own judge is allowed, but that judge sits ahead
146    /// of the fall-open tier and so answers most of the turns this one set a tier for.
147    pub fn new(config: CompositeRouterConfig) -> Result<Self> {
148        let (trigger, message_hash_fallback) = match &config.judge {
149            CompositeJudgeConfig::Capability(judge) => {
150                (judge.classify_trigger, judge.message_hash_fallback)
151            }
152            CompositeJudgeConfig::Custom(config) => {
153                (config.classify_trigger, config.message_hash_fallback)
154            }
155        };
156        if trigger == ClassifyTrigger::EveryRequest {
157            return Err(LibsyError::AlgorithmError {
158                message: "composite: classify_trigger must be user_turn or new_session".to_string(),
159            });
160        }
161        let judge: Arc<dyn Classifier> = match config.judge {
162            CompositeJudgeConfig::Capability(judge) => {
163                // Composite owns the trigger and retention; standalone affinity never runs.
164                Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
165                    config: TaskClassifierConfig {
166                        classify_trigger: ClassifyTrigger::EveryRequest,
167                        message_hash_fallback: false,
168                        ..judge
169                    },
170                })?)
171            }
172            CompositeJudgeConfig::Custom(config) => {
173                LlmTaskClassifier::custom_classifier(&config, true)?
174            }
175        };
176        let setter = TierSetter {
177            judge,
178            trigger,
179            message_hash_fallback,
180            default_tier: config.stage.mode.default_tier(),
181            tiers: Mutex::new(HashMap::new()),
182        };
183        let route = build_stage_route(config.stage)?
184            .with_name(COMPOSITE)
185            .with_processor(Arc::new(setter));
186        Ok(Self { route })
187    }
188}
189
190#[async_trait]
191impl Algorithm for CompositeRouter {
192    fn name(&self) -> &str {
193        COMPOSITE
194    }
195
196    async fn route(
197        self: Arc<Self>,
198        driver: Driver,
199        request: Request,
200    ) -> Result<crate::RoutingOutcome> {
201        self.route.execute(driver, request).await
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use crate::{CapabilityJudgeConfig, LlmCapabilityConfig};
208    use std::collections::HashMap;
209    use std::sync::Arc;
210
211    use switchyard_protocol::{Category, LlmClientError, Message, ModelId, Role};
212
213    use super::*;
214    use crate::algorithms::util::stage::PickerMode;
215    use crate::algorithms::util::tier_fixtures::{JUDGE, Recorder, turn_request};
216    use crate::core::testing::test_drive_with_models;
217
218    fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
219        [
220            (Category::Judge, vec![ModelId::from(JUDGE)]),
221            (Category::Efficient, vec![ModelId::from("weak")]),
222            (Category::Capable, vec![ModelId::from("strong")]),
223            (
224                Category::Any,
225                vec![ModelId::from("strong"), ModelId::from("weak")],
226            ),
227        ]
228        .into()
229    }
230
231    fn user_turn_request() -> Request {
232        let mut request = turn_request(false);
233        request
234            .llm_request
235            .messages
236            .push(Message::text(Role::User, "now rewrite the parser"));
237        request
238    }
239
240    /// The same request shape with no session ID, so only the hash can key it.
241    fn unkeyed(mut request: Request) -> Request {
242        if let Some(metadata) = request.metadata.as_mut() {
243            metadata.session_id = None;
244        }
245        request
246    }
247
248    fn hash_keyed_router() -> Result<Arc<CompositeRouter>> {
249        Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig {
250            judge: CompositeJudgeConfig::Capability(TaskClassifierConfig {
251                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
252                    base_threshold: 0.5,
253                    ..LlmCapabilityConfig::default()
254                }),
255                classify_trigger: ClassifyTrigger::UserTurn,
256                message_hash_fallback: true,
257                ..Default::default()
258            }),
259            stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
260        })?))
261    }
262
263    fn router() -> Result<Arc<CompositeRouter>> {
264        router_with_trigger(ClassifyTrigger::UserTurn)
265    }
266
267    #[tokio::test]
268    async fn capable_first_defaults_to_pinning_efficient_after_a_switch() -> Result<()> {
269        for capable_first in [false, true] {
270            for turns in [None, Some(0), Some(1), Some(-1)] {
271                let mut stage = StageRouterConfig::new(PickerMode::EfficientFirst, 0.5);
272                stage.efficient_hold_turns = turns;
273                let router = router_with_stage(ClassifyTrigger::UserTurn, stage)?;
274                let recorder = Arc::new(Recorder::default());
275                *recorder.judge_p_solve.lock() = if capable_first { 0.0 } else { 1.0 };
276                let mut initial = turn_request(false);
277                initial.llm_request.messages.truncate(1);
278                let mut write = turn_request(false);
279                if let switchyard_protocol::ContentBlock::ToolCall(call) =
280                    &mut write.llm_request.messages[1].content[0]
281                {
282                    call.name = "Write".into();
283                }
284                for request in [
285                    initial,
286                    turn_request(true),
287                    write.clone(),
288                    write.clone(),
289                    write,
290                    turn_request(true),
291                    turn_request(true),
292                    turn_request(true),
293                ] {
294                    test_drive_with_models(
295                        router.clone(),
296                        request,
297                        runtime_models(),
298                        recorder.serve(),
299                    )
300                    .await?;
301                }
302                let actual: Vec<_> = recorder
303                    .routed()
304                    .into_iter()
305                    .map(|call| call.target)
306                    .collect();
307                let initial = if capable_first { "strong" } else { "weak" };
308                let mut expected = vec![initial, "strong", "strong", "strong", "weak"];
309                let count = turns.unwrap_or(if capable_first { -1 } else { 0 });
310                expected.extend((0..3).map(|i| {
311                    if count == -1 || i < count {
312                        "weak"
313                    } else {
314                        "strong"
315                    }
316                }));
317                assert_eq!(
318                    actual, expected,
319                    "capable_first={capable_first}, turns={turns:?}"
320                );
321                assert_eq!(recorder.judge_calls(), 1);
322            }
323        }
324        Ok(())
325    }
326
327    fn router_with_trigger(classify_trigger: ClassifyTrigger) -> Result<Arc<CompositeRouter>> {
328        router_with_stage(
329            classify_trigger,
330            StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
331        )
332    }
333
334    fn router_with_stage(
335        classify_trigger: ClassifyTrigger,
336        stage: StageRouterConfig,
337    ) -> Result<Arc<CompositeRouter>> {
338        Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig {
339            judge: CompositeJudgeConfig::Capability(TaskClassifierConfig {
340                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
341                    base_threshold: 0.5,
342                    ..LlmCapabilityConfig::default()
343                }),
344                classify_trigger,
345                ..Default::default()
346            }),
347            stage,
348        })?))
349    }
350
351    #[test]
352    fn rejects_every_request_as_a_trigger() {
353        for judge in [
354            CompositeJudgeConfig::Capability(TaskClassifierConfig::default()),
355            CompositeJudgeConfig::Custom(CustomClassifierConfig::new(
356                "Choose a tier",
357                serde_json::json!({"type": "object"}),
358                crate::CustomClassifierPolicy::target_selector("/decision/target"),
359            )),
360        ] {
361            let error = CompositeRouter::new(CompositeRouterConfig {
362                judge,
363                stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
364            })
365            .err()
366            .expect("every_request must be rejected for both classifier modes");
367            assert!(matches!(error, LibsyError::AlgorithmError { message }
368                if message == "composite: classify_trigger must be user_turn or new_session"));
369        }
370    }
371
372    fn custom_judge(trigger: ClassifyTrigger) -> CompositeJudgeConfig {
373        let mut config = CustomClassifierConfig::new(
374            "Select a tier",
375            serde_json::json!({"type": "object"}),
376            crate::CustomClassifierPolicy::target_selector("/decision/target"),
377        );
378        config.classify_trigger = trigger;
379        CompositeJudgeConfig::Custom(config)
380    }
381
382    #[tokio::test]
383    async fn a_failed_session_judgment_retains_the_configured_stage_default() -> Result<()> {
384        for custom in [false, true] {
385            for (mode, expected) in [
386                (PickerMode::EfficientFirst, "weak"),
387                (PickerMode::CapableFirst, "strong"),
388            ] {
389                let judge = if custom {
390                    custom_judge(ClassifyTrigger::NewSession)
391                } else {
392                    CompositeJudgeConfig::Capability(TaskClassifierConfig {
393                        classify_trigger: ClassifyTrigger::NewSession,
394                        ..Default::default()
395                    })
396                };
397                let router = Arc::new(CompositeRouter::new(CompositeRouterConfig {
398                    judge,
399                    stage: StageRouterConfig::new(mode, 0.5),
400                })?);
401                let calls = Arc::new(Mutex::new(0));
402                let serve = {
403                    let calls = calls.clone();
404                    move |target: ModelId, _request: Request| {
405                        let completion = if target.as_str() == JUDGE {
406                            *calls.lock() += 1;
407                            "not JSON".to_string()
408                        } else {
409                            target.to_string()
410                        };
411                        async move { Ok(crate::core::testing::reply(completion)) }
412                    }
413                };
414                for request in [
415                    user_turn_request(),
416                    turn_request(false),
417                    user_turn_request(),
418                ] {
419                    let (selected, _) = test_drive_with_models(
420                        router.clone(),
421                        request,
422                        runtime_models(),
423                        serve.clone(),
424                    )
425                    .await?;
426                    assert_eq!(selected.as_str(), expected);
427                }
428                assert_eq!(
429                    *calls.lock(),
430                    1,
431                    "new_session must retain a failed judgment"
432                );
433            }
434        }
435        Ok(())
436    }
437
438    #[tokio::test]
439    async fn a_custom_judge_sets_the_tier_and_retains_it_on_failure() -> Result<()> {
440        let router = Arc::new(CompositeRouter::new(CompositeRouterConfig {
441            judge: custom_judge(ClassifyTrigger::UserTurn),
442            stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
443        })?);
444        let replies = Arc::new(Mutex::new(
445            [
446                Some(r#"{"decision":{"target":"capable"}}"#),
447                None,
448                Some("not JSON"),
449                Some(r#"{"decision":{"target":"efficient"}}"#),
450            ]
451            .into_iter(),
452        ));
453        let serve = {
454            let replies = replies.clone();
455            move |target: ModelId, _request: Request| {
456                let completion = if target.as_str() == JUDGE {
457                    replies
458                        .lock()
459                        .next()
460                        .expect("unexpected judge call")
461                        .map(str::to_string)
462                } else {
463                    Some(target.to_string())
464                };
465                async move {
466                    completion
467                        .map(crate::core::testing::reply)
468                        .ok_or_else(|| LlmClientError::General("judge unavailable".into()))
469                }
470            }
471        };
472        for (request, expected) in [
473            (user_turn_request(), "strong"),
474            (turn_request(false), "strong"),
475            (user_turn_request(), "strong"),
476            (user_turn_request(), "strong"),
477            (user_turn_request(), "weak"),
478        ] {
479            let (selected, _) =
480                test_drive_with_models(router.clone(), request, runtime_models(), serve.clone())
481                    .await?;
482            assert_eq!(selected.as_str(), expected);
483        }
484        assert!(replies.lock().next().is_none());
485        Ok(())
486    }
487
488    /// A route may point its judge at the same target it serves capable turns on.
489    /// The tier then cannot be recovered by looking the served model up in the
490    /// runtime groups — one id, two categories — so it comes off the verdict itself.
491    #[tokio::test]
492    async fn a_judge_sharing_the_capable_model_still_latches_the_tier() -> Result<()> {
493        let models: HashMap<Category, Vec<ModelId>> = [
494            (Category::Judge, vec![ModelId::from("strong")]),
495            (Category::Capable, vec![ModelId::from("strong")]),
496            (Category::Efficient, vec![ModelId::from("weak")]),
497            (
498                Category::Any,
499                vec![ModelId::from("strong"), ModelId::from("weak")],
500            ),
501        ]
502        .into();
503        // The judge runs first as a request-side processor, so the opening call is its own.
504        let calls = Arc::new(Mutex::new(0u32));
505        let serve = {
506            let calls = Arc::clone(&calls);
507            move |target: ModelId, _request: Request| {
508                let calls = Arc::clone(&calls);
509                async move {
510                    let mut calls = calls.lock();
511                    *calls += 1;
512                    let completion = if *calls == 1 {
513                        // p_solve below the threshold: the judge does not trust the
514                        // efficient tier, so the verdict is capable.
515                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.1}"#.to_string()
516                    } else {
517                        target.to_string()
518                    };
519                    Ok(crate::core::testing::reply(completion))
520                }
521            }
522        };
523        let router = router()?;
524
525        test_drive_with_models(
526            router.clone(),
527            user_turn_request(),
528            models.clone(),
529            serve.clone(),
530        )
531        .await?;
532        // A tool step is not a user turn, so nothing re-judges and the latched tier decides.
533        let (selected, _) =
534            test_drive_with_models(router, turn_request(false), models, serve).await?;
535
536        assert_eq!(selected, ModelId::from("strong"));
537        Ok(())
538    }
539
540    #[tokio::test]
541    async fn a_session_without_an_id_keys_on_the_message_hash() -> Result<()> {
542        let recorder = Arc::new(Recorder::default());
543        *recorder.judge_p_solve.lock() = 0.1;
544        let router = hash_keyed_router()?;
545
546        test_drive_with_models(
547            router.clone(),
548            unkeyed(user_turn_request()),
549            runtime_models(),
550            recorder.serve(),
551        )
552        .await?;
553        test_drive_with_models(
554            router.clone(),
555            unkeyed(turn_request(false)),
556            runtime_models(),
557            recorder.serve(),
558        )
559        .await?;
560
561        assert_eq!(
562            recorder.judge_calls(),
563            1,
564            "a tool step is not a user turn, session id or not"
565        );
566        assert_eq!(
567            recorder.routed()[1].target,
568            "strong",
569            "and the tier survives the tool step"
570        );
571        Ok(())
572    }
573
574    #[tokio::test]
575    async fn the_judge_sets_the_tier_once_a_turn_and_the_signals_run_within_it() -> Result<()> {
576        let recorder = Arc::new(Recorder::default());
577        *recorder.judge_p_solve.lock() = 0.1;
578        let router = router()?;
579
580        test_drive_with_models(
581            router.clone(),
582            user_turn_request(),
583            runtime_models(),
584            recorder.serve(),
585        )
586        .await?;
587        test_drive_with_models(
588            router.clone(),
589            turn_request(false),
590            runtime_models(),
591            recorder.serve(),
592        )
593        .await?;
594
595        let routed = recorder.routed();
596        assert_eq!(
597            routed[0].target, "strong",
598            "a quiet turn falls open to the verdict"
599        );
600        assert_eq!(
601            routed[1].target, "strong",
602            "which holds across the tool steps after it"
603        );
604        assert_eq!(
605            recorder.judge_calls(),
606            1,
607            "a tool step is not a new user turn"
608        );
609        Ok(())
610    }
611
612    #[tokio::test]
613    async fn caller_state_carries_the_judged_tier_without_a_session_id() -> Result<()> {
614        for trigger in [ClassifyTrigger::UserTurn, ClassifyTrigger::NewSession] {
615            let recorder = Arc::new(Recorder::default());
616            *recorder.judge_p_solve.lock() = 0.1;
617            let router = router_with_trigger(trigger)?;
618            let run = |request, state| {
619                crate::core::testing::test_drive_with_state(
620                    router.clone(),
621                    request,
622                    runtime_models(),
623                    state,
624                    recorder.serve(),
625                )
626            };
627
628            let first = run(unkeyed(user_turn_request()), State::default()).await?;
629            assert_eq!(first.selected_model_id()?, &ModelId::from("strong"));
630
631            let continued = run(unkeyed(turn_request(false)), first.state).await?;
632            assert_eq!(continued.selected_model_id()?, &ModelId::from("strong"));
633            assert_eq!(
634                continued.metadata.and_then(|m| m.evidence),
635                Some(serde_json::json!({"source": "retained"}))
636            );
637            assert_eq!(recorder.judge_calls(), 1, "{trigger:?}");
638        }
639        Ok(())
640    }
641}
switchyard_libsy/algorithms/composite.rs

switchyard_libsy/algorithms/
composite.rs
