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    fn router_with_trigger(classify_trigger: ClassifyTrigger) -> Result<Arc<CompositeRouter>> {
268        Ok(Arc::new(CompositeRouter::new(CompositeRouterConfig {
269            judge: CompositeJudgeConfig::Capability(TaskClassifierConfig {
270                judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
271                    base_threshold: 0.5,
272                    ..LlmCapabilityConfig::default()
273                }),
274                classify_trigger,
275                ..Default::default()
276            }),
277            stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
278        })?))
279    }
280
281    #[test]
282    fn rejects_every_request_as_a_trigger() {
283        for judge in [
284            CompositeJudgeConfig::Capability(TaskClassifierConfig::default()),
285            CompositeJudgeConfig::Custom(CustomClassifierConfig::new(
286                "Choose a tier",
287                serde_json::json!({"type": "object"}),
288                crate::CustomClassifierPolicy::target_selector("/decision/target"),
289            )),
290        ] {
291            let error = CompositeRouter::new(CompositeRouterConfig {
292                judge,
293                stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
294            })
295            .err()
296            .expect("every_request must be rejected for both classifier modes");
297            assert!(matches!(error, LibsyError::AlgorithmError { message }
298                if message == "composite: classify_trigger must be user_turn or new_session"));
299        }
300    }
301
302    fn custom_judge(trigger: ClassifyTrigger) -> CompositeJudgeConfig {
303        let mut config = CustomClassifierConfig::new(
304            "Select a tier",
305            serde_json::json!({"type": "object"}),
306            crate::CustomClassifierPolicy::target_selector("/decision/target"),
307        );
308        config.classify_trigger = trigger;
309        CompositeJudgeConfig::Custom(config)
310    }
311
312    #[tokio::test]
313    async fn a_failed_session_judgment_retains_the_configured_stage_default() -> Result<()> {
314        for custom in [false, true] {
315            for (mode, expected) in [
316                (PickerMode::EfficientFirst, "weak"),
317                (PickerMode::CapableFirst, "strong"),
318            ] {
319                let judge = if custom {
320                    custom_judge(ClassifyTrigger::NewSession)
321                } else {
322                    CompositeJudgeConfig::Capability(TaskClassifierConfig {
323                        classify_trigger: ClassifyTrigger::NewSession,
324                        ..Default::default()
325                    })
326                };
327                let router = Arc::new(CompositeRouter::new(CompositeRouterConfig {
328                    judge,
329                    stage: StageRouterConfig::new(mode, 0.5),
330                })?);
331                let calls = Arc::new(Mutex::new(0));
332                let serve = {
333                    let calls = calls.clone();
334                    move |target: ModelId, _request: Request| {
335                        let completion = if target.as_str() == JUDGE {
336                            *calls.lock() += 1;
337                            "not JSON".to_string()
338                        } else {
339                            target.to_string()
340                        };
341                        async move { Ok(crate::core::testing::reply(completion)) }
342                    }
343                };
344                for request in [
345                    user_turn_request(),
346                    turn_request(false),
347                    user_turn_request(),
348                ] {
349                    let (selected, _) = test_drive_with_models(
350                        router.clone(),
351                        request,
352                        runtime_models(),
353                        serve.clone(),
354                    )
355                    .await?;
356                    assert_eq!(selected.as_str(), expected);
357                }
358                assert_eq!(
359                    *calls.lock(),
360                    1,
361                    "new_session must retain a failed judgment"
362                );
363            }
364        }
365        Ok(())
366    }
367
368    #[tokio::test]
369    async fn a_custom_judge_sets_the_tier_and_retains_it_on_failure() -> Result<()> {
370        let router = Arc::new(CompositeRouter::new(CompositeRouterConfig {
371            judge: custom_judge(ClassifyTrigger::UserTurn),
372            stage: StageRouterConfig::new(PickerMode::EfficientFirst, 0.5),
373        })?);
374        let replies = Arc::new(Mutex::new(
375            [
376                Some(r#"{"decision":{"target":"capable"}}"#),
377                None,
378                Some("not JSON"),
379                Some(r#"{"decision":{"target":"efficient"}}"#),
380            ]
381            .into_iter(),
382        ));
383        let serve = {
384            let replies = replies.clone();
385            move |target: ModelId, _request: Request| {
386                let completion = if target.as_str() == JUDGE {
387                    replies
388                        .lock()
389                        .next()
390                        .expect("unexpected judge call")
391                        .map(str::to_string)
392                } else {
393                    Some(target.to_string())
394                };
395                async move {
396                    completion
397                        .map(crate::core::testing::reply)
398                        .ok_or_else(|| LlmClientError::General("judge unavailable".into()))
399                }
400            }
401        };
402        for (request, expected) in [
403            (user_turn_request(), "strong"),
404            (turn_request(false), "strong"),
405            (user_turn_request(), "strong"),
406            (user_turn_request(), "strong"),
407            (user_turn_request(), "weak"),
408        ] {
409            let (selected, _) =
410                test_drive_with_models(router.clone(), request, runtime_models(), serve.clone())
411                    .await?;
412            assert_eq!(selected.as_str(), expected);
413        }
414        assert!(replies.lock().next().is_none());
415        Ok(())
416    }
417
418    /// A route may point its judge at the same target it serves capable turns on.
419    /// The tier then cannot be recovered by looking the served model up in the
420    /// runtime groups — one id, two categories — so it comes off the verdict itself.
421    #[tokio::test]
422    async fn a_judge_sharing_the_capable_model_still_latches_the_tier() -> Result<()> {
423        let models: HashMap<Category, Vec<ModelId>> = [
424            (Category::Judge, vec![ModelId::from("strong")]),
425            (Category::Capable, vec![ModelId::from("strong")]),
426            (Category::Efficient, vec![ModelId::from("weak")]),
427            (
428                Category::Any,
429                vec![ModelId::from("strong"), ModelId::from("weak")],
430            ),
431        ]
432        .into();
433        // The judge runs first as a request-side processor, so the opening call is its own.
434        let calls = Arc::new(Mutex::new(0u32));
435        let serve = {
436            let calls = Arc::clone(&calls);
437            move |target: ModelId, _request: Request| {
438                let calls = Arc::clone(&calls);
439                async move {
440                    let mut calls = calls.lock();
441                    *calls += 1;
442                    let completion = if *calls == 1 {
443                        // p_solve below the threshold: the judge does not trust the
444                        // efficient tier, so the verdict is capable.
445                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.1}"#.to_string()
446                    } else {
447                        target.to_string()
448                    };
449                    Ok(crate::core::testing::reply(completion))
450                }
451            }
452        };
453        let router = router()?;
454
455        test_drive_with_models(
456            router.clone(),
457            user_turn_request(),
458            models.clone(),
459            serve.clone(),
460        )
461        .await?;
462        // A tool step is not a user turn, so nothing re-judges and the latched tier decides.
463        let (selected, _) =
464            test_drive_with_models(router, turn_request(false), models, serve).await?;
465
466        assert_eq!(selected, ModelId::from("strong"));
467        Ok(())
468    }
469
470    #[tokio::test]
471    async fn a_session_without_an_id_keys_on_the_message_hash() -> Result<()> {
472        let recorder = Arc::new(Recorder::default());
473        *recorder.judge_p_solve.lock() = 0.1;
474        let router = hash_keyed_router()?;
475
476        test_drive_with_models(
477            router.clone(),
478            unkeyed(user_turn_request()),
479            runtime_models(),
480            recorder.serve(),
481        )
482        .await?;
483        test_drive_with_models(
484            router.clone(),
485            unkeyed(turn_request(false)),
486            runtime_models(),
487            recorder.serve(),
488        )
489        .await?;
490
491        assert_eq!(
492            recorder.judge_calls(),
493            1,
494            "a tool step is not a user turn, session id or not"
495        );
496        assert_eq!(
497            recorder.routed()[1].target,
498            "strong",
499            "and the tier survives the tool step"
500        );
501        Ok(())
502    }
503
504    #[tokio::test]
505    async fn the_judge_sets_the_tier_once_a_turn_and_the_signals_run_within_it() -> Result<()> {
506        let recorder = Arc::new(Recorder::default());
507        *recorder.judge_p_solve.lock() = 0.1;
508        let router = router()?;
509
510        test_drive_with_models(
511            router.clone(),
512            user_turn_request(),
513            runtime_models(),
514            recorder.serve(),
515        )
516        .await?;
517        test_drive_with_models(
518            router.clone(),
519            turn_request(false),
520            runtime_models(),
521            recorder.serve(),
522        )
523        .await?;
524
525        let routed = recorder.routed();
526        assert_eq!(
527            routed[0].target, "strong",
528            "a quiet turn falls open to the verdict"
529        );
530        assert_eq!(
531            routed[1].target, "strong",
532            "which holds across the tool steps after it"
533        );
534        assert_eq!(
535            recorder.judge_calls(),
536            1,
537            "a tool step is not a new user turn"
538        );
539        Ok(())
540    }
541
542    #[tokio::test]
543    async fn caller_state_carries_the_judged_tier_without_a_session_id() -> Result<()> {
544        for trigger in [ClassifyTrigger::UserTurn, ClassifyTrigger::NewSession] {
545            let recorder = Arc::new(Recorder::default());
546            *recorder.judge_p_solve.lock() = 0.1;
547            let router = router_with_trigger(trigger)?;
548            let run = |request, state| {
549                crate::core::testing::test_drive_with_state(
550                    router.clone(),
551                    request,
552                    runtime_models(),
553                    state,
554                    recorder.serve(),
555                )
556            };
557
558            let first = run(unkeyed(user_turn_request()), State::default()).await?;
559            assert_eq!(first.selected_model_id()?, &ModelId::from("strong"));
560
561            let continued = run(unkeyed(turn_request(false)), first.state).await?;
562            assert_eq!(continued.selected_model_id()?, &ModelId::from("strong"));
563            assert_eq!(
564                continued.metadata.and_then(|m| m.evidence),
565                Some(serde_json::json!({"source": "retained"}))
566            );
567            assert_eq!(recorder.judge_calls(), 1, "{trigger:?}");
568        }
569        Ok(())
570    }
571}
switchyard_libsy/algorithms/composite.rs

switchyard_libsy/algorithms/
composite.rs
