switchyard_libsy/algorithms/
subagent.rs1use std::sync::Arc;
7
8use switchyard_protocol::{Category, Metadata, Request};
9
10use super::fall_through::FallThrough;
11use super::util::affinity::{AffinityRouter, ClassifyTrigger};
12use super::util::subagent::SubagentGate;
13use crate::algorithms::llm_class::DefaultCategoryClassifier;
14use crate::core::algorithm::{Algorithm, Driver};
15use crate::core::classifier::Classifier;
16use crate::core::state::State;
17use crate::{LibsyError, Result, RoutingOutcome};
18
19pub struct SubagentRouterConfig {
21 pub classifier: Arc<dyn Classifier<State>>,
23 pub default_target: Category,
25 pub classify_trigger: ClassifyTrigger,
27 pub message_hash_fallback: bool,
29}
30
31impl SubagentRouterConfig {
32 pub fn fixed_target() -> Self {
34 Self {
35 classifier: Arc::new(DefaultCategoryClassifier(Category::Any)),
36 default_target: Category::Any,
37 classify_trigger: ClassifyTrigger::EveryRequest,
38 message_hash_fallback: false,
39 }
40 }
41}
42
43pub struct SubagentRouter {
45 parent: Arc<dyn Algorithm>,
46 subagent: FallThrough<State>,
47}
48
49impl SubagentRouter {
50 pub fn new(parent: Arc<dyn Algorithm>, config: SubagentRouterConfig) -> Result<Self> {
56 if config.message_hash_fallback {
57 return Err(LibsyError::AlgorithmError {
58 message: "sub-agent routing cannot use message_hash_fallback".to_string(),
59 });
60 }
61
62 let mut subagent = match config.classify_trigger {
63 ClassifyTrigger::EveryRequest => FallThrough::new_with_state().with_name("subagent"),
64 ClassifyTrigger::NewSession => {
65 let affinity = Arc::new(AffinityRouter::for_subagents());
66 FallThrough::new_with_state()
67 .with_name("subagent")
68 .with_processor(affinity.clone())
69 .with_classifier(affinity)
70 }
71 ClassifyTrigger::UserTurn => {
72 return Err(LibsyError::AlgorithmError {
73 message: "sub-agent routing cannot use classify_trigger = user_turn"
74 .to_string(),
75 });
76 }
77 };
78 subagent = subagent
79 .with_classifier(Arc::new(SubagentGate::new(config.classifier)))
80 .with_classifier(Arc::new(DefaultCategoryClassifier(config.default_target)));
81
82 Ok(Self { parent, subagent })
83 }
84}
85
86#[async_trait::async_trait]
87impl Algorithm for SubagentRouter {
88 fn name(&self) -> &str {
89 self.parent.name()
90 }
91
92 async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<RoutingOutcome> {
93 if request
94 .metadata
95 .as_ref()
96 .is_some_and(Metadata::is_subagent_work)
97 {
98 self.subagent.execute(driver.for_subagent()?, request).await
100 } else {
101 self.parent.clone().route(driver, request).await
102 }
103 }
104}
105
106#[cfg(test)]
107mod tests {
108 use std::sync::Arc;
109 use std::sync::atomic::{AtomicUsize, Ordering};
110
111 use async_trait::async_trait;
112 use switchyard_protocol::{Category, Metadata, ModelId, Request, Response, text_request};
113
114 use super::{SubagentRouter, SubagentRouterConfig};
115 use crate::algorithms::passthrough::Passthrough;
116 use crate::core::classifier::{Classification, Classifier, Score};
117 use crate::core::testing::{echo, test_drive_with_models};
118 use crate::{ClassifyTrigger, Driver, RuntimeModels, State};
119
120 struct ScriptedClassifier {
121 calls: AtomicUsize,
122 }
123
124 #[async_trait]
125 impl Classifier<State> for ScriptedClassifier {
126 async fn score(
127 &self,
128 _state: &mut State,
129 _request: &mut Request,
130 driver: &Driver,
131 ) -> crate::Result<(Classification, Option<Response>)> {
132 let category = match self.calls.fetch_add(1, Ordering::Relaxed) {
133 0 => Some(Category::Capable),
134 1 => Some(Category::Efficient),
135 _ => None,
136 };
137 let scores = match category {
138 Some(category) => vec![Score {
139 confidence: 1.0,
140 target: driver.first_model_for(&category)?.clone(),
141 category: Some(category),
142 }],
143 None => Vec::new(),
144 };
145 Ok((Classification::Scores(scores), None))
146 }
147 }
148
149 fn request(metadata: Option<Metadata>) -> Request {
150 Request {
151 llm_request: text_request(Some("auto".to_string()), "hi"),
152 raw_request: None,
153 metadata,
154 }
155 }
156
157 fn child(agent_id: &str) -> Request {
158 request(Some(Metadata {
159 session_id: Some("session-1".to_string()),
160 agent_id: Some(agent_id.to_string()),
161 is_subagent: true,
162 is_delegated_work: true,
163 ..Metadata::default()
164 }))
165 }
166
167 fn configured(classifier: Arc<dyn Classifier<State>>) -> crate::Result<Arc<SubagentRouter>> {
168 Ok(Arc::new(SubagentRouter::new(
169 Arc::new(Passthrough),
170 SubagentRouterConfig {
171 classifier,
172 default_target: Category::Capable,
173 classify_trigger: ClassifyTrigger::NewSession,
174 message_hash_fallback: false,
175 },
176 )?))
177 }
178
179 #[tokio::test]
180 async fn routes_parent_and_children_with_affinity_and_default() -> crate::Result<()> {
181 let classifier = Arc::new(ScriptedClassifier {
182 calls: AtomicUsize::new(0),
183 });
184 let router = configured(classifier.clone())?;
185
186 let models = RuntimeModels::new([(Category::Any, vec![ModelId::from("parent")])].into())
188 .with_subagent(
189 [
190 (
191 Category::Any,
192 vec![ModelId::from("worker"), ModelId::from("reviewer")],
193 ),
194 (Category::Capable, vec![ModelId::from("worker")]),
195 (Category::Efficient, vec![ModelId::from("reviewer")]),
196 ]
197 .into(),
198 );
199 let (selected_parent, _) =
200 test_drive_with_models(router.clone(), request(None), models.clone(), echo()).await?;
201 let (first, _) =
202 test_drive_with_models(router.clone(), child("child-1"), models.clone(), echo())
203 .await?;
204 let (same_child, _) =
205 test_drive_with_models(router.clone(), child("child-1"), models.clone(), echo())
206 .await?;
207 let (sibling, _) =
208 test_drive_with_models(router.clone(), child("child-2"), models.clone(), echo())
209 .await?;
210 let (defaulted, _) =
211 test_drive_with_models(router.clone(), child("child-3"), models.clone(), echo())
212 .await?;
213 let maintenance = request(Some(Metadata {
214 session_id: Some("session-1".to_string()),
215 agent_id: Some("child-1".to_string()),
216 is_subagent: true,
217 is_delegated_work: false,
218 ..Metadata::default()
219 }));
220 let (maintenance, _) =
221 test_drive_with_models(router, maintenance, models.clone(), echo()).await?;
222
223 assert_eq!(selected_parent, "parent");
224 assert_eq!(first, "worker");
225 assert_eq!(same_child, "worker");
226 assert_eq!(sibling, "reviewer");
227 assert_eq!(defaulted, "worker");
228 assert_eq!(maintenance, "parent");
229 assert_eq!(classifier.calls.load(Ordering::Relaxed), 3);
230
231 let fixed = Arc::new(SubagentRouter::new(
232 Arc::new(Passthrough),
233 SubagentRouterConfig::fixed_target(),
234 )?);
235 let (fixed, _) = test_drive_with_models(fixed, child("fixed"), models, echo()).await?;
236 assert_eq!(fixed, "worker");
237 Ok(())
238 }
239}