1use std::collections::HashSet;
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::{Deserialize, Deserializer};
11use serde_json::Value;
12use switchyard_protocol::{Category, ContentBlock, Message, Role};
13
14mod decision;
15pub use decision::{DecisionChoice, DecisionJudgeConfig, DirectClassifierConfig};
16use decision::{DecisionClassifier, DirectClassifier};
17
18use super::escalation;
19use super::fall_through::FallThrough;
20use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS;
21use super::util::affinity::{AffinityRouter, ClassifyTrigger, is_user_turn};
22use super::util::classifier_contract::{
23 ClassifierContract, ClassifierContractConfig, ClassifierResponseFormat,
24};
25use super::util::escalation::EscalationJudgeConfig;
26use super::util::llm_judge::{
27 ClassifierInput, JsonSchemaDecoder, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig,
28 SerdeDecoder, StructuredJudge,
29};
30use super::util::target_selector::TargetSelectorPolicy;
31use crate::core::algorithm::{Algorithm, Driver};
32use crate::core::classifier::{Classification, Classifier, Score};
33use crate::core::state::State;
34use crate::{LibsyError, Result};
35use switchyard_protocol::{Request, Response};
36
37const PROMPT_TEMPLATE: &str = include_str!("../prompts/capability-classifier/prompt.md");
38const SCHEMA_TEMPLATE: &str = include_str!("../prompts/capability-classifier/schema.json");
39const ALGORITHM_NAME: &str = "llm_task_classifier";
41
42const TRAILING_ROUTING_INSTRUCTION: &str =
51 "Route the conversation above. Output ONLY the routing JSON object, nothing else.";
52
53#[derive(Deserialize)]
54#[serde(deny_unknown_fields)]
55struct TaskClassifierVerdict {
56 crux: String,
57 primary_rule: String,
58 capability_boundary: String,
59 p_solve: f64,
60}
61
62impl TaskClassifierVerdict {
63 fn is_valid(&self) -> bool {
65 (0.0..=1.0).contains(&self.p_solve)
66 && !self.crux.trim().is_empty()
67 && matches!(
68 (
69 self.primary_rule.as_str(),
70 self.capability_boundary.as_str()
71 ),
72 ("SUP-1" | "SUP-2" | "SUP-3" | "SUP-4" | "SUP-5", "supported")
73 | ("UNC-1" | "UNC-2", "uncertain")
74 | ("LIM-1" | "LIM-2", "unsupported")
75 | ("none", "unmatched")
76 )
77 }
78
79 fn boundary_steps(&self, unmatched_steps: u8) -> Option<u8> {
85 match self.capability_boundary.as_str() {
86 "supported" => Some(0),
87 "uncertain" => Some(1),
88 "unmatched" => Some(unmatched_steps),
89 "unsupported" => Some(2),
90 _ => None,
91 }
92 }
93}
94
95fn trim_messages(messages: &[Message], recent_turn_window: usize) -> Vec<Message> {
105 let is_instruction = |message: &Message| matches!(message.role, Role::System | Role::Developer);
106 let mut kept: Vec<&Message> = messages.iter().filter(|m| is_instruction(m)).collect();
107 let Some(task) = messages.iter().position(is_user_turn) else {
108 return kept.into_iter().cloned().collect();
109 };
110 kept.push(&messages[task]);
111
112 let tail: Vec<&Message> = messages[task + 1..]
113 .iter()
114 .filter(|m| !is_instruction(m))
115 .collect();
116 kept.extend(&tail[window_start(&tail, recent_turn_window)..]);
117 kept.into_iter().cloned().collect()
118}
119
120fn window_start(tail: &[&Message], recent_turn_window: usize) -> usize {
133 let counted = tail.len().saturating_sub(recent_turn_window);
134 if counted == tail.len() {
136 return counted;
137 }
138 let mut unpaired: HashSet<&str> = HashSet::new();
139 for (start, message) in tail.iter().enumerate().rev() {
140 for block in message.content.iter().rev() {
143 match block {
144 ContentBlock::ToolResult(result) => {
145 unpaired.insert(result.tool_call_id.as_str());
146 }
147 ContentBlock::ToolCall(call) => {
148 unpaired.remove(call.id.as_str());
149 }
150 _ => {}
151 }
152 }
153 if start <= counted && unpaired.is_empty() {
154 return start;
155 }
156 }
157 counted
158}
159
160fn task_messages(messages: &[Message]) -> Vec<Message> {
162 let is_task_content = |block: &ContentBlock| {
165 !matches!(
166 block,
167 ContentBlock::ToolCall(_)
168 | ContentBlock::ToolResult(_)
169 | ContentBlock::Reasoning { .. }
170 )
171 };
172 let mut user_messages = messages
173 .iter()
174 .filter(|message| is_user_turn(message) && message.content.iter().any(is_task_content));
175 let Some(opening_task) = user_messages.next() else {
176 return Vec::new();
177 };
178 [Some(opening_task), user_messages.next_back()]
179 .into_iter()
180 .flatten()
181 .map(|message| {
182 Message::new(
183 Role::User,
184 message
185 .content
186 .iter()
187 .filter(|block| is_task_content(block))
188 .cloned()
189 .collect(),
190 )
191 })
192 .collect()
193}
194
195struct TaskInput {
197 recent_turn_window: Option<usize>,
198}
199
200impl TaskInput {
201 fn messages(&self, request: &Request) -> Vec<Message> {
203 let mut messages = match self.recent_turn_window {
206 Some(window) => trim_messages(&request.llm_request.messages, window),
207 None => task_messages(&request.llm_request.messages),
208 };
209 for message in &mut messages {
214 message
215 .content
216 .retain(|block| !matches!(block, ContentBlock::Reasoning { .. }));
217 }
218 messages.retain(|message| !message.content.is_empty());
219 messages
220 }
221}
222
223impl ClassifierInput for TaskInput {
224 fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
225 let mut messages = self.messages(request);
226 if self.recent_turn_window.is_some() {
230 messages.push(Message::text(
231 Role::User,
232 TRAILING_ROUTING_INSTRUCTION.to_string(),
233 ));
234 }
235 messages
236 }
237}
238
239struct TaskClassifierPolicy {
240 base_threshold: f64,
241 threshold_step: f64,
242 unmatched_steps: u8,
243}
244
245impl TaskClassifierPolicy {
246 fn new(config: &LlmCapabilityConfig) -> Self {
247 Self {
248 base_threshold: config.base_threshold,
249 threshold_step: config.threshold_step,
250 unmatched_steps: config.unmatched_steps,
251 }
252 }
253
254 fn threshold(&self, verdict: &TaskClassifierVerdict) -> Option<f64> {
256 let steps = verdict.boundary_steps(self.unmatched_steps)?;
257 Some(self.base_threshold + f64::from(steps) * self.threshold_step)
258 }
259}
260
261impl JudgePolicy for TaskClassifierPolicy {
262 type Verdict = TaskClassifierVerdict;
263
264 fn to_classification(
265 &self,
266 verdict: Option<&Self::Verdict>,
267 driver: &Driver,
268 ) -> Result<Classification> {
269 let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
272 return Ok(Classification::Ambiguous(vec![]));
273 };
274 let Some(threshold) = self.threshold(verdict) else {
277 return Ok(Classification::Ambiguous(vec![]));
278 };
279 let category = if verdict.p_solve >= threshold
280 || (threshold - verdict.p_solve).abs() <= f64::EPSILON
281 {
282 Category::Efficient
283 } else {
284 Category::Capable
285 };
286 let Some(target) = driver.models_for(&category).first().cloned() else {
289 return Ok(Classification::Ambiguous(vec![]));
290 };
291 Ok(Classification::Scores(vec![Score {
292 target,
293 confidence: 1.0,
294 category: Some(category),
295 }]))
296 }
297}
298
299fn capability_evidence(
301 policy: &TaskClassifierPolicy,
302 verdict: Option<&TaskClassifierVerdict>,
303) -> Option<Value> {
304 let verdict = verdict?;
305 let Some(threshold) = verdict
306 .is_valid()
307 .then(|| policy.threshold(verdict))
308 .flatten()
309 else {
310 return Some(serde_json::json!({
311 "source": "fail_open",
312 "reason_code": "invalid_verdict",
313 }));
314 };
315 Some(serde_json::json!({
316 "source": "llm-classifier",
317 "score": verdict.p_solve,
318 "threshold": threshold,
319 }))
320}
321
322#[derive(Clone, Debug)]
324pub struct TaskClassifierConfig {
325 pub judge: CapabilityJudgeConfig,
327 pub fail_open: bool,
329 pub classify_trigger: ClassifyTrigger,
331 pub message_hash_fallback: bool,
333 pub recent_turn_window: Option<usize>,
336}
337
338#[derive(Clone, Debug)]
340pub enum CapabilityJudgeConfig {
341 Llm(LlmCapabilityConfig),
343 Decision(DecisionJudgeConfig),
345}
346
347impl Default for CapabilityJudgeConfig {
348 fn default() -> Self {
349 Self::Llm(LlmCapabilityConfig::default())
350 }
351}
352
353#[derive(Clone, Debug)]
355pub struct LlmCapabilityConfig {
356 pub base_threshold: f64,
358 pub threshold_step: f64,
361 pub unmatched_steps: u8,
365 pub contract: ClassifierContractConfig,
367 pub max_output_tokens: u64,
369 pub judge_deadline_ms: Option<u64>,
373}
374
375impl Default for LlmCapabilityConfig {
376 fn default() -> Self {
377 Self {
378 base_threshold: 0.0,
379 threshold_step: 0.0,
380 unmatched_steps: DEFAULT_UNMATCHED_STEPS,
381 contract: ClassifierContractConfig::default(),
382 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
383 judge_deadline_ms: None,
384 }
385 }
386}
387
388#[derive(Deserialize)]
390#[serde(deny_unknown_fields)]
391struct TaskClassifierConfigWire {
392 base_threshold: f64,
393 #[serde(default = "default_fail_open")]
394 fail_open: bool,
395 #[serde(default)]
396 threshold_step: f64,
397 #[serde(default = "default_unmatched_steps")]
398 unmatched_steps: u8,
399 #[serde(default)]
400 classify_trigger: ClassifyTrigger,
401 #[serde(default)]
402 message_hash_fallback: bool,
403 #[serde(default)]
404 recent_turn_window: Option<usize>,
405 #[serde(default)]
406 prompt: Option<String>,
407 #[serde(default)]
408 response_format_type: ClassifierResponseFormat,
409 #[serde(default = "default_judge_max_output_tokens")]
410 max_output_tokens: u64,
411 #[serde(default)]
412 judge_deadline_ms: Option<u64>,
413}
414
415impl<'de> Deserialize<'de> for TaskClassifierConfig {
416 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
417 where
418 D: Deserializer<'de>,
419 {
420 let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
421 let mut contract = ClassifierContractConfig::default();
422 if let Some(prompt) = wire.prompt {
423 contract = contract.with_prompt(prompt);
424 }
425 contract = contract.with_response_format_type(wire.response_format_type);
426 Ok(Self {
427 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
428 base_threshold: wire.base_threshold,
429 threshold_step: wire.threshold_step,
430 unmatched_steps: wire.unmatched_steps,
431 contract,
432 max_output_tokens: wire.max_output_tokens,
433 judge_deadline_ms: wire.judge_deadline_ms,
434 }),
435 fail_open: wire.fail_open,
436 classify_trigger: wire.classify_trigger,
437 message_hash_fallback: wire.message_hash_fallback,
438 recent_turn_window: wire.recent_turn_window,
439 })
440 }
441}
442
443const fn default_fail_open() -> bool {
444 true
445}
446
447pub const DEFAULT_UNMATCHED_STEPS: u8 = 1;
449
450const fn default_unmatched_steps() -> u8 {
451 DEFAULT_UNMATCHED_STEPS
452}
453
454const fn default_judge_max_output_tokens() -> u64 {
455 DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
456}
457
458impl Default for TaskClassifierConfig {
459 fn default() -> Self {
460 Self {
461 judge: CapabilityJudgeConfig::default(),
462 fail_open: default_fail_open(),
463 classify_trigger: ClassifyTrigger::default(),
464 message_hash_fallback: false,
465 recent_turn_window: None,
466 }
467 }
468}
469
470impl LlmCapabilityConfig {
471 fn validate(&self) -> Result<()> {
472 if !(0.0..=1.0).contains(&self.base_threshold) {
473 return Err(LibsyError::AlgorithmError {
474 message: format!(
475 "base_threshold must be between 0 and 1, got {}",
476 self.base_threshold
477 ),
478 });
479 }
480 if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
481 return Err(LibsyError::AlgorithmError {
482 message: format!(
483 "threshold_step must be finite and greater than or equal to 0, got {}",
484 self.threshold_step
485 ),
486 });
487 }
488 if self.unmatched_steps > 2 {
489 return Err(LibsyError::AlgorithmError {
490 message: format!(
491 "unmatched_steps must be 0, 1, or 2, got {}",
492 self.unmatched_steps
493 ),
494 });
495 }
496 let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
497 if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
498 return Err(LibsyError::AlgorithmError {
499 message: format!(
500 "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
501 ),
502 });
503 }
504 if self.max_output_tokens == 0 {
505 return Err(LibsyError::AlgorithmError {
506 message: "max_output_tokens must be at least 1".to_string(),
507 });
508 }
509 Ok(())
510 }
511}
512
513impl TaskClassifierConfig {
514 fn validate(&self) -> Result<()> {
515 match &self.judge {
516 CapabilityJudgeConfig::Llm(config) => config.validate()?,
517 CapabilityJudgeConfig::Decision(config) => config.validate()?,
518 }
519 if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
521 return Err(LibsyError::AlgorithmError {
522 message:
523 "message_hash_fallback requires classify_trigger = new_session or user_turn"
524 .to_string(),
525 });
526 }
527 Ok(())
528 }
529}
530
531#[derive(Clone, Debug)]
533pub enum CustomClassifierPolicy {
534 TargetSelector {
536 selector: String,
538 },
539}
540
541impl CustomClassifierPolicy {
542 pub fn target_selector(selector: impl Into<String>) -> Self {
544 Self::TargetSelector {
545 selector: selector.into(),
546 }
547 }
548}
549
550#[derive(Clone, Debug)]
552pub struct CustomClassifierConfig {
553 pub prompt: String,
555 pub response_schema: Value,
557 pub policy: CustomClassifierPolicy,
559 pub classify_trigger: ClassifyTrigger,
561 pub message_hash_fallback: bool,
563 pub recent_turn_window: Option<usize>,
565 pub max_output_tokens: u64,
567 pub judge_deadline_ms: Option<u64>,
569}
570
571impl CustomClassifierConfig {
572 pub fn new(
574 prompt: impl Into<String>,
575 response_schema: Value,
576 policy: CustomClassifierPolicy,
577 ) -> Self {
578 Self {
579 prompt: prompt.into(),
580 response_schema,
581 policy,
582 classify_trigger: ClassifyTrigger::default(),
583 message_hash_fallback: false,
584 recent_turn_window: None,
585 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
586 judge_deadline_ms: None,
587 }
588 }
589
590 fn validate(&self) -> Result<()> {
591 if self.max_output_tokens == 0 {
592 return Err(LibsyError::AlgorithmError {
593 message: "max_output_tokens must be at least 1".to_string(),
594 });
595 }
596 if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
598 return Err(LibsyError::AlgorithmError {
599 message:
600 "message_hash_fallback requires classify_trigger = new_session or user_turn"
601 .to_string(),
602 });
603 }
604 Ok(())
605 }
606}
607
608enum CustomPolicyRuntime {
609 TargetSelector(TargetSelectorPolicy),
610}
611
612impl JudgePolicy for CustomPolicyRuntime {
613 type Verdict = Value;
614
615 fn to_classification(
616 &self,
617 verdict: Option<&Self::Verdict>,
618 driver: &Driver,
619 ) -> Result<Classification> {
620 match self {
621 Self::TargetSelector(policy) => policy.to_classification(verdict, driver),
622 }
623 }
624}
625
626fn affinity_router(
628 trigger: ClassifyTrigger,
629 message_hash_fallback: bool,
630) -> Option<Arc<AffinityRouter>> {
631 let router = match trigger {
632 ClassifyTrigger::EveryRequest => return None,
633 ClassifyTrigger::NewSession => AffinityRouter::new(),
634 ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
635 };
636 let router = if message_hash_fallback {
637 router.with_message_hash_fallback()
638 } else {
639 router
640 };
641 Some(Arc::new(router))
642}
643
644pub struct LlmTaskClassifier {
646 route: FallThrough,
647 inner: Arc<dyn Classifier>,
649}
650
651struct ClassifierRouteConfig {
652 default_target: Category,
653 classify_trigger: ClassifyTrigger,
654 message_hash_fallback: bool,
655}
656
657pub struct DefaultCategoryClassifier(pub Category);
660
661#[async_trait]
662impl Classifier for DefaultCategoryClassifier {
663 async fn score(
664 &self,
665 _state: &mut State,
666 _request: &mut Request,
667 driver: &Driver,
668 ) -> Result<(Classification, Option<Response>)> {
669 let target = driver.first_model_for(&self.0)?;
670 driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"}));
671 Ok((
672 Classification::Scores(vec![Score {
673 target: target.clone(),
674 confidence: 0.0,
675 category: Some(self.0.clone()),
676 }]),
677 None,
678 ))
679 }
680}
681
682#[derive(Clone)]
684#[non_exhaustive]
685pub enum LlmClassifierConfig {
686 Capability {
688 config: TaskClassifierConfig,
690 },
691 Escalation {
693 contract: ClassifierContractConfig,
695 config: EscalationJudgeConfig,
697 max_output_tokens: u64,
699 judge_deadline_ms: Option<u64>,
701 },
702 Custom {
704 default_target: Category,
706 config: CustomClassifierConfig,
708 },
709 Direct {
711 default_target: Category,
713 config: DirectClassifierConfig,
715 },
716}
717
718impl LlmTaskClassifier {
719 pub fn new(config: LlmClassifierConfig) -> Result<Self> {
726 match config {
727 LlmClassifierConfig::Capability { config } => Self::build_capability(config),
728 LlmClassifierConfig::Escalation {
729 contract,
730 config,
731 max_output_tokens,
732 judge_deadline_ms,
733 } => Self::build_escalation(contract, config, max_output_tokens, judge_deadline_ms),
734 LlmClassifierConfig::Custom {
735 default_target,
736 config,
737 } => Self::build_custom(default_target, config),
738 LlmClassifierConfig::Direct {
739 default_target,
740 config,
741 } => Self::build_direct(default_target, config),
742 }
743 }
744
745 fn build_capability(config: TaskClassifierConfig) -> Result<Self> {
746 config.validate()?;
747 let classify_trigger = config.classify_trigger;
748 let message_hash_fallback = config.message_hash_fallback;
749 let input = TaskInput {
750 recent_turn_window: config.recent_turn_window,
751 };
752 let classifier: Arc<dyn Classifier> = match config.judge {
753 CapabilityJudgeConfig::Llm(judge) => Arc::new(
754 JudgeClassifier::new(
755 StructuredJudge::new(
756 input,
757 Self::load_capability_contract(&judge.contract)?,
758 SerdeDecoder::new(),
759 JudgeRuntimeConfig::new(judge.max_output_tokens)?
760 .with_deadline_ms(judge.judge_deadline_ms)?,
761 ),
762 TaskClassifierPolicy::new(&judge),
763 )
764 .with_error_recovery(config.fail_open)
765 .with_evidence(capability_evidence),
766 ),
767 CapabilityJudgeConfig::Decision(judge) => {
768 Arc::new(DecisionClassifier::new(judge, input, config.fail_open)?)
769 }
770 };
771 Self::from_classifier(
772 classifier,
773 ClassifierRouteConfig {
774 default_target: Category::Capable,
775 classify_trigger,
776 message_hash_fallback,
777 },
778 )
779 }
780
781 fn build_custom(default_target: Category, config: CustomClassifierConfig) -> Result<Self> {
782 Self::from_classifier(
783 Self::custom_classifier(&config, false)?,
784 ClassifierRouteConfig {
785 default_target,
786 classify_trigger: config.classify_trigger,
787 message_hash_fallback: config.message_hash_fallback,
788 },
789 )
790 }
791
792 fn build_direct(default_target: Category, config: DirectClassifierConfig) -> Result<Self> {
793 if !config.options.contains(&default_target) {
794 return Err(LibsyError::AlgorithmError {
795 message: "direct default_target must be one of the options".into(),
796 });
797 }
798 let route = ClassifierRouteConfig {
799 default_target,
800 classify_trigger: config.classify_trigger,
801 message_hash_fallback: config.message_hash_fallback,
802 };
803 Self::from_classifier(Arc::new(DirectClassifier::new(config)?), route)
804 }
805
806 pub(crate) fn custom_classifier(
808 config: &CustomClassifierConfig,
809 fail_open: bool,
810 ) -> Result<Arc<dyn Classifier>> {
811 config.validate()?;
812 let contract =
813 ClassifierContract::from_inner_schema(&config.prompt, config.response_schema.clone())?;
814 let policy = match &config.policy {
815 CustomClassifierPolicy::TargetSelector { selector } => {
816 CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(selector)?)
817 }
818 };
819 let classifier = JudgeClassifier::new(
820 StructuredJudge::new(
821 TaskInput {
822 recent_turn_window: config.recent_turn_window,
823 },
824 contract,
825 JsonSchemaDecoder::new(),
826 JudgeRuntimeConfig::new(config.max_output_tokens)?
827 .with_deadline_ms(config.judge_deadline_ms)?,
828 ),
829 policy,
830 )
831 .with_error_recovery(fail_open);
832 let classifier = if fail_open {
833 classifier.with_evidence(|_, verdict| {
834 verdict.map(|_| serde_json::json!({"source": "llm-classifier"}))
835 })
836 } else {
837 classifier
838 };
839 Ok(Arc::new(classifier))
840 }
841
842 fn build_escalation(
843 contract_config: ClassifierContractConfig,
844 config: EscalationJudgeConfig,
845 max_output_tokens: u64,
846 judge_deadline_ms: Option<u64>,
847 ) -> Result<Self> {
848 let inner = escalation::build_classifier(
849 contract_config,
850 config,
851 max_output_tokens,
852 judge_deadline_ms,
853 )?;
854 Ok(Self {
855 route: FallThrough::new_with_state()
856 .with_name(ALGORITHM_NAME)
857 .with_classifier(Arc::clone(&inner)),
858 inner,
859 })
860 }
861
862 fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
864 ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
865 }
866
867 fn from_classifier(inner: Arc<dyn Classifier>, config: ClassifierRouteConfig) -> Result<Self> {
869 if config.message_hash_fallback && config.classify_trigger == ClassifyTrigger::EveryRequest
871 {
872 return Err(LibsyError::AlgorithmError {
873 message:
874 "message_hash_fallback requires classify_trigger = new_session or user_turn"
875 .to_string(),
876 });
877 }
878 let mut route = FallThrough::new_with_state().with_name(ALGORITHM_NAME);
880 if let Some(affinity) =
881 affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
882 {
883 route = route
885 .with_processor(affinity.clone())
886 .with_classifier(affinity.clone());
887 }
888 let fallback = DefaultCategoryClassifier(config.default_target);
889 Ok(Self {
890 route: route
891 .with_classifier(inner.clone())
892 .with_classifier(Arc::new(fallback)),
893 inner,
894 })
895 }
896}
897
898#[async_trait]
899impl Classifier for LlmTaskClassifier {
900 async fn score(
901 &self,
902 state: &mut State,
903 request: &mut Request,
904 driver: &Driver,
905 ) -> Result<(Classification, Option<Response>)> {
906 self.inner.score(state, request, driver).await
907 }
908}
909
910#[async_trait]
911impl Algorithm for LlmTaskClassifier {
912 fn name(&self) -> &str {
913 "llm_task_classifier"
914 }
915
916 async fn route(
917 self: Arc<Self>,
918 driver: Driver,
919 request: Request,
920 ) -> Result<crate::RoutingOutcome> {
921 self.route.execute(driver, request).await
922 }
923}
924
925#[cfg(test)]
926mod tests {
927 use std::collections::HashMap;
928 use std::sync::Arc;
929
930 use parking_lot::Mutex;
931 use serde_json::Value;
932
933 use super::*;
934 use switchyard_protocol::{
935 ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ModelId, ToolCall,
936 ToolResult, completion_text, text_request, text_response,
937 };
938
939 use crate::algorithms::util::llm_judge::Judge;
940 use crate::core::testing::{Serve, test_drive_with_models};
941 use switchyard_protocol::{LlmResponse, Response};
942
943 const TEST_THRESHOLD: f64 = 0.5;
944
945 type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;
946
947 fn llm_config(base_threshold: f64) -> LlmCapabilityConfig {
948 LlmCapabilityConfig {
949 base_threshold,
950 ..LlmCapabilityConfig::default()
951 }
952 }
953
954 fn test_config(base_threshold: f64) -> TaskClassifierConfig {
955 TaskClassifierConfig {
956 judge: CapabilityJudgeConfig::Llm(llm_config(base_threshold)),
957 ..TaskClassifierConfig::default()
958 }
959 }
960
961 fn policy() -> TaskClassifierPolicy {
962 TaskClassifierPolicy::new(&llm_config(TEST_THRESHOLD))
963 }
964
965 fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
966 [
967 (Category::Judge, vec![ModelId::from("judge")]),
968 (Category::Efficient, vec![ModelId::from("efficient")]),
969 (Category::Capable, vec![ModelId::from("capable")]),
970 (
971 Category::Any,
972 vec![ModelId::from("efficient"), ModelId::from("capable")],
973 ),
974 ]
975 .into()
976 }
977
978 fn policy_driver() -> Driver {
979 Driver::new("test", Arc::new(runtime_models().into())).0
980 }
981
982 fn verdict(
983 p_solve: f64,
984 capability_boundary: &str,
985 primary_rule: &str,
986 ) -> TaskClassifierVerdict {
987 TaskClassifierVerdict {
988 crux: "test crux".to_string(),
989 primary_rule: primary_rule.to_string(),
990 capability_boundary: capability_boundary.to_string(),
991 p_solve,
992 }
993 }
994
995 fn selected(
996 policy: &TaskClassifierPolicy,
997 verdict: Option<&TaskClassifierVerdict>,
998 ) -> Result<ModelId> {
999 policy
1000 .to_classification(verdict, &policy_driver())?
1001 .argmax(false)?
1002 .map(|score| score.target)
1003 .ok_or_else(|| LibsyError::AlgorithmError {
1004 message: "policy abstained".to_string(),
1005 })
1006 }
1007
1008 #[derive(Default)]
1011 struct Recorder {
1012 calls: Mutex<Vec<String>>,
1013 call_roles: Mutex<Vec<(String, bool)>>,
1014 judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
1015 judge_system_prompts: Mutex<Vec<String>>,
1016 }
1017
1018 impl Recorder {
1019 fn calls(&self) -> Vec<String> {
1020 self.calls.lock().clone()
1021 }
1022
1023 fn call_roles(&self) -> Vec<(String, bool)> {
1024 self.call_roles.lock().clone()
1025 }
1026
1027 fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
1028 self.judge_max_output_tokens.lock().clone()
1029 }
1030
1031 fn judge_system_prompts(&self) -> Vec<String> {
1032 self.judge_system_prompts.lock().clone()
1033 }
1034
1035 fn serve(self: &Arc<Self>) -> impl Serve {
1036 let recorder = Arc::clone(self);
1037 move |model: ModelId, request: Request| {
1038 let recorder = Arc::clone(&recorder);
1039 async move {
1040 let model = model.to_string();
1041 recorder.calls.lock().push(model.clone());
1042 recorder
1043 .call_roles
1044 .lock()
1045 .push((model.clone(), model != "judge"));
1046 let completion = if model == "judge" {
1047 recorder
1048 .judge_max_output_tokens
1049 .lock()
1050 .push(request.llm_request.output.max_output_tokens);
1051 recorder.judge_system_prompts.lock().extend(
1052 request
1053 .llm_request
1054 .instructions
1055 .first()
1056 .and_then(|instruction| {
1057 instruction.content.iter().find_map(|b| {
1058 if let ContentBlock::Text { text } = b {
1059 Some(text.clone())
1060 } else {
1061 None
1062 }
1063 })
1064 }),
1065 );
1066 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1067 } else {
1068 format!("answer from {model}")
1069 };
1070 Ok(Response {
1071 llm_response: LlmResponse::Agg(text_response(None, completion)),
1072 metadata: request.metadata,
1073 upstream_headers: http::HeaderMap::new(),
1074 })
1075 }
1076 }
1077 }
1078 }
1079
1080 fn unreachable_judge() -> impl Serve {
1082 |model: ModelId, request: Request| async move {
1083 let model = model.to_string();
1084 if model == "judge" {
1085 return Err(LlmClientError::Timeout {
1086 source: Box::new(std::io::Error::other("judge unreachable")),
1087 });
1088 }
1089 Ok(Response {
1090 llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
1091 metadata: request.metadata,
1092 upstream_headers: http::HeaderMap::new(),
1093 })
1094 }
1095 }
1096
1097 fn slow_judge(delay: std::time::Duration) -> impl Serve {
1099 move |model: ModelId, request: Request| {
1100 let model = model.to_string();
1101 let slow = model == "judge";
1102 async move {
1103 if slow {
1104 tokio::time::sleep(delay).await;
1105 }
1106 Ok(Response {
1107 llm_response: LlmResponse::Agg(text_response(
1108 None,
1109 format!("answer from {model}"),
1110 )),
1111 metadata: request.metadata,
1112 upstream_headers: http::HeaderMap::new(),
1113 })
1114 }
1115 }
1116 }
1117
1118 fn stalled_judge_stream() -> impl Serve {
1120 use futures::StreamExt;
1121 |model: ModelId, request: Request| async move {
1122 let model = model.to_string();
1123 let llm_response = if model == "judge" {
1124 LlmResponse::Stream(futures::stream::pending().boxed())
1125 } else {
1126 LlmResponse::Agg(text_response(None, format!("answer from {model}")))
1127 };
1128 Ok(Response {
1129 llm_response,
1130 metadata: request.metadata,
1131 upstream_headers: http::HeaderMap::new(),
1132 })
1133 }
1134 }
1135
1136 fn deadline_router(fail_open: bool) -> Result<Arc<LlmTaskClassifier>> {
1137 Ok(Arc::new(LlmTaskClassifier::new(
1138 LlmClassifierConfig::Capability {
1139 config: TaskClassifierConfig {
1140 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1141 judge_deadline_ms: Some(10),
1142 ..llm_config(TEST_THRESHOLD)
1143 }),
1144 fail_open,
1145 ..TaskClassifierConfig::default()
1146 },
1147 },
1148 )?))
1149 }
1150
1151 fn router() -> Result<Arc<LlmTaskClassifier>> {
1152 Ok(Arc::new(LlmTaskClassifier::new(
1153 LlmClassifierConfig::Capability {
1154 config: test_config(TEST_THRESHOLD),
1155 },
1156 )?))
1157 }
1158
1159 fn classify_request() -> Request {
1160 Request {
1161 llm_request: text_request(Some("auto".to_string()), "classify this task"),
1162 raw_request: None,
1163 metadata: None,
1164 }
1165 }
1166
1167 fn classify_session_request() -> Request {
1168 Request {
1169 metadata: Some(Metadata {
1170 session_id: Some("session-1".to_string()),
1171 ..Metadata::default()
1172 }),
1173 ..classify_request()
1174 }
1175 }
1176
1177 fn classify_follow_up_request() -> Request {
1178 let mut request = classify_request();
1179 request
1180 .llm_request
1181 .messages
1182 .push(Message::text(Role::Assistant, "I will add the test."));
1183 request.llm_request.messages.push(Message::text(
1184 Role::User,
1185 "Now run the test suite and report the result.",
1186 ));
1187 request
1188 }
1189
1190 #[tokio::test]
1191 async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
1192 let router = router()?;
1193
1194 let (selected_model, response) = test_drive_with_models(
1195 router,
1196 classify_request(),
1197 runtime_models(),
1198 unreachable_judge(),
1199 )
1200 .await?;
1201
1202 assert_eq!(selected_model, "capable");
1203 assert_eq!(
1204 response.llm_response.as_agg().map(completion_text),
1205 Some("answer from capable".to_string())
1206 );
1207 Ok(())
1208 }
1209
1210 #[tokio::test]
1211 async fn a_judge_past_its_deadline_routes_capable() -> Result<()> {
1212 let (selected_model, response) = test_drive_with_models(
1213 deadline_router(true)?,
1214 classify_request(),
1215 runtime_models(),
1216 slow_judge(std::time::Duration::from_millis(500)),
1217 )
1218 .await?;
1219
1220 assert_eq!(selected_model, "capable");
1221 assert_eq!(
1222 response.llm_response.as_agg().map(completion_text),
1223 Some("answer from capable".to_string())
1224 );
1225 Ok(())
1226 }
1227
1228 #[tokio::test]
1229 async fn a_stalled_judge_stream_is_cut_by_the_deadline() -> Result<()> {
1230 let (selected_model, _) = test_drive_with_models(
1233 deadline_router(true)?,
1234 classify_request(),
1235 runtime_models(),
1236 stalled_judge_stream(),
1237 )
1238 .await?;
1239
1240 assert_eq!(selected_model, "capable");
1241 Ok(())
1242 }
1243
1244 #[tokio::test]
1245 async fn a_judge_past_its_deadline_stops_when_not_failing_open() -> Result<()> {
1246 let outcome = test_drive_with_models(
1247 deadline_router(false)?,
1248 classify_request(),
1249 runtime_models(),
1250 slow_judge(std::time::Duration::from_millis(500)),
1251 )
1252 .await;
1253 let error = outcome
1254 .err()
1255 .expect("a deadline expiry with fail_open = false must stop the request");
1256
1257 assert!(
1258 matches!(
1259 error,
1260 LibsyError::ClientCall {
1261 source: LlmClientError::Timeout { .. },
1262 ..
1263 }
1264 ),
1265 "expected a client timeout, got {error:?}"
1266 );
1267 Ok(())
1268 }
1269
1270 #[test]
1271 fn a_zero_judge_deadline_is_rejected_in_every_mode() {
1272 let capability = LlmClassifierConfig::Capability {
1273 config: TaskClassifierConfig {
1274 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1275 judge_deadline_ms: Some(0),
1276 ..llm_config(TEST_THRESHOLD)
1277 }),
1278 ..TaskClassifierConfig::default()
1279 },
1280 };
1281 assert!(
1282 LlmTaskClassifier::new(capability).is_err(),
1283 "capability mode must reject judge_deadline_ms = 0"
1284 );
1285
1286 let mut custom = CustomClassifierConfig::new(
1287 "Route by topic.",
1288 serde_json::json!({"type": "object"}),
1289 CustomClassifierPolicy::target_selector("/target"),
1290 );
1291 custom.judge_deadline_ms = Some(0);
1292 assert!(
1293 LlmTaskClassifier::new(LlmClassifierConfig::Custom {
1294 default_target: Category::Capable,
1295 config: custom,
1296 })
1297 .is_err(),
1298 "custom mode must reject judge_deadline_ms = 0"
1299 );
1300
1301 let escalation = LlmClassifierConfig::Escalation {
1302 contract: ClassifierContractConfig::default(),
1303 config: EscalationJudgeConfig::default(),
1304 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
1305 judge_deadline_ms: Some(0),
1306 };
1307 assert!(
1308 LlmTaskClassifier::new(escalation).is_err(),
1309 "escalation mode must reject judge_deadline_ms = 0"
1310 );
1311 }
1312
1313 #[test]
1314 fn classifier_config_parses_a_judge_deadline() {
1315 let config: TaskClassifierConfig = serde_json::from_value(serde_json::json!({
1316 "base_threshold": 0.5,
1317 "judge_deadline_ms": 250,
1318 }))
1319 .expect("a configured judge deadline should parse");
1320 let CapabilityJudgeConfig::Llm(judge) = config.judge else {
1321 panic!("expected the LLM judge variant");
1322 };
1323 assert_eq!(judge.judge_deadline_ms, Some(250));
1324
1325 let config: TaskClassifierConfig = serde_json::from_value(serde_json::json!({
1326 "base_threshold": 0.5,
1327 }))
1328 .expect("an unset judge deadline should parse");
1329 let CapabilityJudgeConfig::Llm(judge) = config.judge else {
1330 panic!("expected the LLM judge variant");
1331 };
1332 assert_eq!(judge.judge_deadline_ms, None);
1333 }
1334
1335 #[tokio::test]
1336 async fn classifier_judges_each_request_without_affinity() -> Result<()> {
1337 let recorder = Arc::new(Recorder::default());
1338 let router = router()?;
1339 let request = classify_request();
1340 let models = runtime_models();
1341
1342 test_drive_with_models(
1343 router.clone(),
1344 request.clone(),
1345 models.clone(),
1346 recorder.serve(),
1347 )
1348 .await?;
1349 test_drive_with_models(router, request, models, recorder.serve()).await?;
1350
1351 assert_eq!(
1352 recorder.calls(),
1353 vec!["judge", "efficient", "judge", "efficient"]
1354 );
1355 assert_eq!(
1356 recorder.call_roles(),
1357 vec![
1358 ("judge".to_string(), false),
1359 ("efficient".to_string(), true),
1360 ("judge".to_string(), false),
1361 ("efficient".to_string(), true),
1362 ]
1363 );
1364 Ok(())
1365 }
1366
1367 #[tokio::test]
1368 async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1369 let recorder = Arc::new(Recorder::default());
1370 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1371 config: TaskClassifierConfig {
1372 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1373 max_output_tokens: 512,
1374 ..llm_config(TEST_THRESHOLD)
1375 }),
1376 ..test_config(TEST_THRESHOLD)
1377 },
1378 })?);
1379
1380 test_drive_with_models(
1381 router,
1382 classify_request(),
1383 runtime_models(),
1384 recorder.serve(),
1385 )
1386 .await?;
1387
1388 assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1389 Ok(())
1390 }
1391
1392 #[tokio::test]
1393 async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1394 let recorder = Arc::new(Recorder::default());
1395 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1396 config: TaskClassifierConfig {
1397 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1398 contract: ClassifierContractConfig::default()
1399 .with_prompt("Custom capability rubric."),
1400 ..llm_config(TEST_THRESHOLD)
1401 }),
1402 ..test_config(TEST_THRESHOLD)
1403 },
1404 })?);
1405
1406 test_drive_with_models(
1407 router,
1408 classify_request(),
1409 runtime_models(),
1410 recorder.serve(),
1411 )
1412 .await?;
1413
1414 let prompts = recorder.judge_system_prompts();
1415 assert_eq!(prompts.len(), 1);
1416 assert_eq!(prompts[0], "Custom capability rubric.");
1417 Ok(())
1418 }
1419
1420 #[tokio::test]
1421 async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1422 let recorder = Arc::new(Recorder::default());
1423 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1424 config: TaskClassifierConfig {
1425 classify_trigger: ClassifyTrigger::NewSession,
1426 ..test_config(TEST_THRESHOLD)
1427 },
1428 })?);
1429
1430 let request = classify_session_request();
1431 let models = runtime_models();
1432 test_drive_with_models(
1433 router.clone(),
1434 request.clone(),
1435 models.clone(),
1436 recorder.serve(),
1437 )
1438 .await?;
1439 test_drive_with_models(router, request, models, recorder.serve()).await?;
1440
1441 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1442 Ok(())
1443 }
1444
1445 #[tokio::test]
1446 async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1447 let recorder = Arc::new(Recorder::default());
1448 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1449 config: TaskClassifierConfig {
1450 classify_trigger: ClassifyTrigger::NewSession,
1451 message_hash_fallback: true,
1452 recent_turn_window: None,
1453 ..test_config(TEST_THRESHOLD)
1454 },
1455 })?);
1456
1457 let models = runtime_models();
1458 test_drive_with_models(
1459 router.clone(),
1460 classify_request(),
1461 models.clone(),
1462 recorder.serve(),
1463 )
1464 .await?;
1465 test_drive_with_models(
1466 router,
1467 classify_follow_up_request(),
1468 models,
1469 recorder.serve(),
1470 )
1471 .await?;
1472
1473 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1474 Ok(())
1475 }
1476
1477 #[tokio::test]
1478 async fn one_classifier_uses_each_requests_runtime_models() -> Result<()> {
1479 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1480 config: TaskClassifierConfig {
1481 classify_trigger: ClassifyTrigger::NewSession,
1482 ..test_config(TEST_THRESHOLD)
1483 },
1484 })?);
1485 let calls = Arc::new(Mutex::new(Vec::new()));
1486 let serve = |calls: Arc<Mutex<Vec<String>>>| {
1487 move |model: ModelId, _request: Request| {
1488 let calls = Arc::clone(&calls);
1489 async move {
1490 calls.lock().push(model.to_string());
1491 let text = if model.as_str().starts_with("judge-") {
1492 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1493 } else {
1494 model.to_string()
1495 };
1496 Ok(Response {
1497 llm_response: LlmResponse::Agg(text_response(None, text)),
1498 metadata: None,
1499 upstream_headers: Default::default(),
1500 })
1501 }
1502 }
1503 };
1504 let models = |suffix: &str| -> HashMap<Category, Vec<ModelId>> {
1505 [
1506 (
1507 Category::Judge,
1508 vec![ModelId::from(format!("judge-{suffix}"))],
1509 ),
1510 (
1511 Category::Efficient,
1512 vec![ModelId::from(format!("efficient-{suffix}"))],
1513 ),
1514 (
1515 Category::Capable,
1516 vec![ModelId::from(format!("capable-{suffix}"))],
1517 ),
1518 (
1519 Category::Any,
1520 vec![
1521 ModelId::from(format!("efficient-{suffix}")),
1522 ModelId::from(format!("capable-{suffix}")),
1523 ],
1524 ),
1525 ]
1526 .into()
1527 };
1528 let request = classify_session_request();
1529
1530 let (first, _) = test_drive_with_models(
1531 router.clone(),
1532 request.clone(),
1533 models("a"),
1534 serve(Arc::clone(&calls)),
1535 )
1536 .await?;
1537 let (second, _) =
1538 test_drive_with_models(router, request, models("b"), serve(Arc::clone(&calls))).await?;
1539
1540 assert_eq!(first, "efficient-a");
1541 assert_eq!(second, "efficient-b");
1542 assert_eq!(
1543 &*calls.lock(),
1544 &["judge-a", "efficient-a", "judge-b", "efficient-b"]
1545 );
1546 Ok(())
1547 }
1548
1549 #[tokio::test]
1550 async fn decision_judge_routes_on_capable_probability() -> Result<()> {
1551 use crate::{Call, RuntimeModels, drive};
1552 use serde_json::json;
1553 use std::collections::BTreeMap;
1554 use std::sync::atomic::{AtomicUsize, Ordering};
1555 use switchyard_protocol::{
1556 BooleanEstimate, DecisionAnswer, DecisionResponse, DecisionValue, Probability,
1557 ProviderConfidence,
1558 };
1559
1560 let settings = DecisionJudgeConfig {
1561 cutoff: Some(0.4),
1562 ..DecisionJudgeConfig::default()
1563 };
1564 let mut request = classify_session_request();
1565 request.llm_request.messages.push(Message::new(
1566 Role::Assistant,
1567 vec![
1568 ContentBlock::Reasoning {
1569 text: "private reasoning".into(),
1570 signature: None,
1571 details: vec![],
1572 },
1573 ContentBlock::Text {
1574 text: "visible progress".into(),
1575 },
1576 ],
1577 ));
1578 request
1579 .llm_request
1580 .messages
1581 .push(Message::text(Role::User, "latest follow-up"));
1582 let original_messages = request.llm_request.messages.clone();
1583 let expected_task = json!([
1584 original_messages[0].clone(),
1585 Message::text(Role::Assistant, "visible progress"),
1586 Message::text(Role::User, "latest follow-up"),
1587 ]);
1588
1589 for (name, score, fail_open, expected) in [
1590 ("above", 0.45, true, Some("capable")),
1591 ("equal", 0.4, true, Some("efficient")),
1592 ("below", 0.39, false, Some("efficient")),
1593 ("missing answer", 0.0, true, Some("capable")),
1594 ("no distribution", 0.0, true, Some("capable")),
1595 ("wrong type", 0.0, true, Some("capable")),
1596 ("out of range", 1.1, true, Some("capable")),
1597 ("nonfinite", f64::NAN, true, Some("capable")),
1598 ("missing probability", 0.0, true, Some("capable")),
1599 ("missing answer", 0.0, false, None),
1600 ("no distribution", 0.0, false, None),
1601 ("wrong type", 0.0, false, None),
1602 ("out of range", 1.1, false, None),
1603 ("negative", -0.1, false, None),
1604 ("nonfinite", f64::NAN, false, None),
1605 ("infinite", f64::INFINITY, false, None),
1606 ("missing probability", 0.0, false, None),
1607 ("provider error", 0.0, true, Some("capable")),
1608 ("provider error", 0.0, false, None),
1609 ("dropped reply", 0.0, true, Some("capable")),
1610 ("custom choices", 0.45, false, Some("capable")),
1611 ] {
1612 let mut config = settings.clone();
1613 if name == "below" {
1614 config.instructions = Some("Choose capable when the task needs it.".into());
1615 }
1616 if name == "custom choices" {
1617 config.choices[0].id = "complex".into();
1618 config.choices[1].id = "simple".into();
1619 config.apply_cutoff_to = "complex".into();
1620 }
1621 let override_instructions = config.instructions.clone();
1622 let choice_ids = [config.choices[0].id.clone(), config.choices[1].id.clone()];
1623 let router: Arc<dyn Algorithm> =
1624 Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1625 config: TaskClassifierConfig {
1626 judge: CapabilityJudgeConfig::Decision(config),
1627 fail_open,
1628 recent_turn_window: None,
1629 classify_trigger: ClassifyTrigger::NewSession,
1630 ..TaskClassifierConfig::default()
1631 },
1632 })?);
1633 let calls = AtomicUsize::new(0);
1634 let serve = |call| {
1635 let calls = &calls;
1636 let expected_task = &expected_task;
1637 let override_instructions = &override_instructions;
1638 let [capable_id, efficient_id] = choice_ids.clone();
1639 async move {
1640 let Call::Decision(call) = call else {
1641 panic!("judge must use the decision step")
1642 };
1643 calls.fetch_add(1, Ordering::SeqCst);
1644 assert_eq!(call.model, "judge");
1645 assert_eq!(call.request.model, Some("judge".into()));
1646 assert_eq!(
1647 serde_json::to_value(&call.request.input).unwrap(),
1648 *expected_task
1649 );
1650 assert_eq!(call.request.questions.len(), 1);
1651 let question = &call.request.questions["route"];
1652 let switchyard_protocol::DecisionKind::Choice { options } = &question.kind
1653 else {
1654 panic!("expected a choice")
1655 };
1656 assert_eq!(
1657 options
1658 .iter()
1659 .map(|option| option.id.as_str())
1660 .collect::<Vec<_>>(),
1661 [capable_id.as_str(), efficient_id.as_str()]
1662 );
1663 if let Some(instructions) = override_instructions {
1664 assert_eq!(&question.instructions, instructions);
1665 } else {
1666 assert!(
1667 question
1668 .instructions
1669 .as_str()
1670 .is_some_and(|text| !text.is_empty())
1671 );
1672 }
1673 if name == "provider error" {
1674 return call
1675 .respond(std::future::ready(Err(LibsyError::client_call(
1676 "judge",
1677 LlmClientError::General("private provider body".into()),
1678 ))))
1679 .await;
1680 }
1681 if name == "dropped reply" {
1682 drop(call);
1683 return Ok(());
1684 }
1685 let value = if name == "wrong type" {
1686 DecisionValue::Boolean(BooleanEstimate::Value(true))
1687 } else {
1688 DecisionValue::Choice {
1689 selected: efficient_id.clone(),
1690 probabilities: (name != "no distribution").then(|| {
1691 BTreeMap::from([
1692 (
1693 if name == "missing probability" {
1694 "other".into()
1695 } else {
1696 capable_id.clone()
1697 },
1698 Probability(score),
1699 ),
1700 (efficient_id.clone(), Probability(1.0 - score)),
1701 ])
1702 }),
1703 }
1704 };
1705 let answers = if name == "missing answer" {
1706 BTreeMap::new()
1707 } else {
1708 BTreeMap::from([(
1709 "route".into(),
1710 DecisionAnswer {
1711 value,
1712 provider_confidence: Some(ProviderConfidence(0.99)),
1713 },
1714 )])
1715 };
1716 call.respond(std::future::ready(Ok(DecisionResponse {
1717 id: None,
1718 model: Some("provider-judge".into()),
1719 answers,
1720 usage: Default::default(),
1721 })))
1722 .await
1723 }
1724 };
1725 let models = Arc::new(RuntimeModels::new(runtime_models()));
1726 let result = drive(router.clone(), request.clone(), models.clone(), &serve).await;
1727 assert_eq!(calls.load(Ordering::SeqCst), 1, "{name}");
1728 let Some(expected) = expected else {
1729 if name == "provider error" {
1730 assert!(
1731 matches!(result, Err(LibsyError::ClientCall { .. })),
1732 "{name}"
1733 );
1734 } else {
1735 assert!(
1736 matches!(result, Err(LibsyError::InvalidDecision { target, .. }) if target == "judge"),
1737 "expected invalid decision error: {name}"
1738 );
1739 }
1740 continue;
1741 };
1742 let outcome = result?;
1743 assert_eq!(outcome.selected_model_id()?, expected, "{name}");
1744 assert!(outcome.response.is_none());
1745 assert_eq!(outcome.request.llm_request.messages, original_messages);
1746 let metadata = outcome.metadata.expect("routing metadata");
1747 assert_eq!(metadata.judgments.len(), 1);
1748 let evidence = metadata.evidence.expect("routing evidence");
1749 if matches!(name, "above" | "equal" | "below" | "custom choices") {
1750 assert_eq!(metadata.judgments[0].target.as_deref(), Some(expected));
1751 assert_eq!(evidence["source"], "decision_classifier");
1752 assert_eq!(evidence["verdict"], expected);
1753 assert_eq!(evidence["capable_threshold"], 0.4);
1754 assert_eq!(evidence["capable_score"], score);
1755 } else {
1756 let reason = match name {
1757 "provider error" => "client_error",
1758 "dropped reply" => "call_error",
1759 _ => "invalid_verdict",
1760 };
1761 assert_eq!(
1762 evidence,
1763 json!({"source": "fail_open", "reason_code": reason})
1764 );
1765 }
1766 let retained = drive(router, request.clone(), models, &serve).await?;
1767 assert_eq!(retained.selected_model_id()?, expected);
1768 assert!(
1769 retained
1770 .metadata
1771 .expect("retained metadata")
1772 .judgments
1773 .is_empty()
1774 );
1775 assert_eq!(
1776 calls.load(Ordering::SeqCst),
1777 1,
1778 "affinity should skip the judge: {name}"
1779 );
1780 }
1781 for cutoff in [-0.1, 1.1, f64::NAN] {
1782 let mut config = settings.clone();
1783 config.cutoff = Some(cutoff);
1784 assert!(
1785 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1786 config: TaskClassifierConfig {
1787 judge: CapabilityJudgeConfig::Decision(config),
1788 ..TaskClassifierConfig::default()
1789 },
1790 })
1791 .is_err()
1792 );
1793 }
1794 Ok(())
1795 }
1796
1797 #[tokio::test]
1798 async fn direct_judge_falls_back_to_the_default_target() -> Result<()> {
1799 use crate::{Call, RuntimeModels, drive};
1800 use std::collections::BTreeMap;
1801 use switchyard_protocol::{DecisionAnswer, DecisionResponse, DecisionValue};
1802
1803 for (reply, fail_open, expected) in [
1804 (Some("unknown"), true, Some("efficient")),
1805 (Some("unknown"), false, None),
1806 (None, true, Some("efficient")),
1807 (None, false, None),
1808 ] {
1809 let router: Arc<dyn Algorithm> =
1810 Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Direct {
1811 default_target: Category::Efficient,
1812 config: DirectClassifierConfig {
1813 options: vec![Category::Efficient, Category::Capable],
1814 fail_open,
1815 classify_trigger: ClassifyTrigger::EveryRequest,
1816 message_hash_fallback: false,
1817 recent_turn_window: None,
1818 },
1819 })?);
1820 let serve = |call| async move {
1821 let Call::Decision(call) = call else {
1822 panic!("judge must use the decision step")
1823 };
1824 let response = match reply {
1825 Some(selected) => Ok(DecisionResponse {
1826 id: None,
1827 model: None,
1828 answers: BTreeMap::from([(
1829 "route".into(),
1830 DecisionAnswer {
1831 value: DecisionValue::Choice {
1832 selected: selected.into(),
1833 probabilities: None,
1834 },
1835 provider_confidence: None,
1836 },
1837 )]),
1838 usage: Default::default(),
1839 }),
1840 None => Err(LibsyError::client_call(
1841 "judge",
1842 LlmClientError::General("unavailable".into()),
1843 )),
1844 };
1845 call.respond(std::future::ready(response)).await
1846 };
1847 let models = Arc::new(RuntimeModels::new(runtime_models()));
1848 let result = drive(router, classify_request(), models, &serve).await;
1849 match expected {
1850 Some(expected) => assert_eq!(result?.selected_model_id()?, expected),
1851 None => assert!(result.is_err()),
1852 }
1853 }
1854 Ok(())
1855 }
1856
1857 #[test]
1858 fn the_threshold_boundary_is_inclusive() -> Result<()> {
1859 let policy = policy();
1860 let at_threshold = verdict(0.5, "supported", "SUP-1");
1861 let below_threshold = verdict(0.49, "supported", "SUP-1");
1862 assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1863 assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1864 Ok(())
1865 }
1866
1867 #[test]
1868 fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1869 let borderline = verdict(0.5, "supported", "SUP-1");
1870 let strict = TaskClassifierPolicy::new(&llm_config(0.9));
1871 let lenient = TaskClassifierPolicy::new(&llm_config(0.1));
1872 assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1873 assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1874 Ok(())
1875 }
1876
1877 #[test]
1878 fn classifier_config_rejects_unknown_fields() {
1879 let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1880 "base_threshold": 0.5,
1881 "classifier_magic": true,
1882 }))
1883 .expect_err("unknown classifier fields must be rejected");
1884
1885 assert!(
1886 error
1887 .to_string()
1888 .contains("unknown field `classifier_magic`"),
1889 "{error}"
1890 );
1891 }
1892
1893 #[test]
1894 fn invalid_classifier_config_is_rejected() -> Result<()> {
1895 for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1896 assert!(
1897 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1898 config: test_config(bad),
1899 })
1900 .is_err(),
1901 "base threshold {bad} should be rejected"
1902 );
1903 }
1904 for config in [
1905 TaskClassifierConfig {
1906 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1907 base_threshold: 0.5,
1908 threshold_step: -0.1,
1909 ..LlmCapabilityConfig::default()
1910 }),
1911 ..TaskClassifierConfig::default()
1912 },
1913 TaskClassifierConfig {
1914 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1915 base_threshold: 0.8,
1916 threshold_step: 0.11,
1917 ..LlmCapabilityConfig::default()
1918 }),
1919 ..TaskClassifierConfig::default()
1920 },
1921 TaskClassifierConfig {
1922 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1923 base_threshold: 0.5,
1924 ..LlmCapabilityConfig::default()
1925 }),
1926 message_hash_fallback: true,
1927 ..TaskClassifierConfig::default()
1928 },
1929 TaskClassifierConfig {
1930 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1931 base_threshold: 0.5,
1932 max_output_tokens: 0,
1933 ..LlmCapabilityConfig::default()
1934 }),
1935 ..TaskClassifierConfig::default()
1936 },
1937 ] {
1938 assert!(LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).is_err());
1939 }
1940 for base_threshold in [0.0, 1.0] {
1941 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1942 config: test_config(base_threshold),
1943 })?;
1944 }
1945 Ok(())
1946 }
1947
1948 #[test]
1949 fn message_hash_fallback_accepts_retaining_triggers() -> Result<()> {
1950 for trigger in [ClassifyTrigger::NewSession, ClassifyTrigger::UserTurn] {
1954 let config = TaskClassifierConfig {
1955 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1956 base_threshold: 0.5,
1957 ..LlmCapabilityConfig::default()
1958 }),
1959 classify_trigger: trigger,
1960 message_hash_fallback: true,
1961 ..TaskClassifierConfig::default()
1962 };
1963 LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).map_err(
1964 |error| LibsyError::AlgorithmError {
1965 message: format!("{trigger:?} with message_hash_fallback rejected: {error}"),
1966 },
1967 )?;
1968 }
1969 let every_request = TaskClassifierConfig {
1970 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1971 base_threshold: 0.5,
1972 ..LlmCapabilityConfig::default()
1973 }),
1974 classify_trigger: ClassifyTrigger::EveryRequest,
1975 message_hash_fallback: true,
1976 ..TaskClassifierConfig::default()
1977 };
1978 assert!(
1979 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1980 config: every_request
1981 })
1982 .is_err(),
1983 "every_request with message_hash_fallback should stay rejected"
1984 );
1985 Ok(())
1986 }
1987
1988 #[test]
1989 fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1990 let policy = policy();
1991 let inconsistent_rule = TaskClassifierVerdict {
1992 capability_boundary: "uncertain".to_string(),
1993 ..verdict(1.0, "supported", "SUP-1")
1994 };
1995 let empty_crux = TaskClassifierVerdict {
1996 crux: " ".to_string(),
1997 ..verdict(1.0, "supported", "SUP-1")
1998 };
1999 let unusable = [
2000 Some(verdict(1.1, "supported", "SUP-1")),
2001 Some(inconsistent_rule),
2002 Some(empty_crux),
2003 None,
2004 ];
2005 for verdict in unusable {
2006 let classification = policy.to_classification(verdict.as_ref(), &policy_driver())?;
2007 assert!(matches!(classification, Classification::Ambiguous(_)));
2008 assert!(classification.argmax(false)?.is_none());
2009 assert!(classification.argmax(true)?.is_none());
2010 }
2011 Ok(())
2012 }
2013
2014 #[test]
2015 fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
2016 let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
2017 threshold_step: 0.1,
2018 ..llm_config(0.4)
2019 });
2020
2021 assert_eq!(
2022 selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
2023 "efficient"
2024 );
2025 assert_eq!(
2026 selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
2027 "capable"
2028 );
2029 assert_eq!(
2030 selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
2031 "efficient"
2032 );
2033 assert_eq!(
2034 selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
2035 "efficient"
2036 );
2037 assert_eq!(
2038 selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
2039 "capable"
2040 );
2041 assert_eq!(
2042 selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
2043 "efficient"
2044 );
2045 Ok(())
2046 }
2047
2048 #[test]
2049 fn unmatched_steps_raise_only_the_unmatched_threshold() -> Result<()> {
2050 let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
2051 threshold_step: 0.1,
2052 unmatched_steps: 2,
2053 ..llm_config(0.75)
2054 });
2055
2056 assert_eq!(
2058 selected(&policy, Some(&verdict(0.85, "unmatched", "none")))?,
2059 "capable"
2060 );
2061 assert_eq!(
2062 selected(&policy, Some(&verdict(0.95, "unmatched", "none")))?,
2063 "efficient"
2064 );
2065 assert_eq!(
2067 selected(&policy, Some(&verdict(0.85, "uncertain", "UNC-1")))?,
2068 "efficient"
2069 );
2070 assert_eq!(
2071 selected(&policy, Some(&verdict(0.75, "supported", "SUP-2")))?,
2072 "efficient"
2073 );
2074 Ok(())
2075 }
2076
2077 #[test]
2078 fn unmatched_steps_default_to_one_and_reject_values_above_two() {
2079 let parsed: TaskClassifierConfig =
2080 serde_json::from_value(serde_json::json!({ "base_threshold": 0.5 }))
2081 .expect("config without unmatched_steps parses");
2082 let CapabilityJudgeConfig::Llm(judge) = &parsed.judge else {
2083 panic!("a flat capability config builds an LLM judge");
2084 };
2085 assert_eq!(judge.unmatched_steps, 1);
2086
2087 let invalid = TaskClassifierConfig {
2088 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
2089 unmatched_steps: 3,
2090 ..llm_config(0.5)
2091 }),
2092 ..TaskClassifierConfig::default()
2093 };
2094 assert!(invalid.validate().is_err());
2095 }
2096
2097 fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
2100 Ok(StructuredJudge::new(
2101 TaskInput { recent_turn_window },
2102 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
2103 SerdeDecoder::new(),
2104 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
2105 ))
2106 }
2107
2108 fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
2109 let judge = capability_judge(Some(recent_turn_window))?;
2110 let request = Request {
2111 llm_request: LlmRequest {
2112 messages: vec![
2113 Message::text(Role::System, "client instructions"),
2114 Message::text(
2115 Role::User,
2116 "<environment_context>setup</environment_context>",
2117 ),
2118 Message::text(Role::User, "initial task"),
2119 Message::text(Role::Assistant, "old response"),
2120 Message::text(Role::User, "old follow-up"),
2121 Message::text(Role::Assistant, "recent 1"),
2122 Message::text(Role::User, "recent 2"),
2123 ],
2124 ..LlmRequest::default()
2125 },
2126 raw_request: None,
2127 metadata: None,
2128 };
2129 Ok(judge
2130 .build_request(&State::default(), &request)
2131 .llm_request
2132 .messages
2133 .iter()
2134 .filter_map(|message| message.text_content("\n"))
2135 .collect())
2136 }
2137
2138 #[test]
2139 fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
2140 let contents = judged_contents(2)?;
2142 assert!(contents.contains(&"client instructions".to_string()));
2143 assert!(contents.contains(&"initial task".to_string()));
2144 assert!(contents.contains(&"recent 1".to_string()));
2145 assert!(contents.contains(&"recent 2".to_string()));
2146 assert!(!contents.contains(&"old response".to_string()));
2147 Ok(())
2148 }
2149
2150 #[test]
2151 fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
2152 let contents = judged_contents(0)?;
2153 assert!(contents.contains(&"client instructions".to_string()));
2154 assert!(contents.contains(&"initial task".to_string()));
2155 assert!(!contents.contains(&"recent 2".to_string()));
2156 Ok(())
2157 }
2158
2159 fn tool_call(id: &str) -> Message {
2160 Message::new(
2161 Role::Assistant,
2162 vec![ContentBlock::ToolCall(ToolCall {
2163 id: id.to_string(),
2164 name: "search".to_string(),
2165 arguments: Value::Null,
2166 })],
2167 )
2168 }
2169
2170 fn tool_result(id: &str) -> Message {
2171 Message::new(
2172 Role::Tool,
2173 vec![ContentBlock::ToolResult(ToolResult {
2174 tool_call_id: id.to_string(),
2175 content: vec![ContentBlock::Text {
2176 text: "tool output".to_string(),
2177 }],
2178 is_error: None,
2179 })],
2180 )
2181 }
2182
2183 #[test]
2184 fn default_task_input_keeps_user_content_around_tool_results() {
2185 let context = Message::text(
2186 Role::User,
2187 "<environment_context><current_date>2026-10-31</current_date></environment_context>",
2188 );
2189 let mut result = tool_result("call-1");
2190 result.role = Role::User;
2191 let mut mixed = result.clone();
2192 mixed.content.push(ContentBlock::Text {
2193 text: "latest follow-up".to_string(),
2194 });
2195 let input = TaskInput {
2196 recent_turn_window: None,
2197 };
2198 let mut request = Request {
2199 llm_request: LlmRequest {
2200 messages: vec![
2201 context.clone(),
2202 result.clone(),
2203 Message::text(Role::User, "initial task"),
2204 tool_call("call-1"),
2205 mixed,
2206 result.clone(),
2207 context,
2208 ],
2209 ..LlmRequest::default()
2210 },
2211 ..Request::default()
2212 };
2213 assert_eq!(
2214 input.build_messages(&State::default(), &request),
2215 vec![
2216 Message::text(Role::User, "initial task"),
2217 Message::text(Role::User, "latest follow-up"),
2218 ]
2219 );
2220 request.llm_request.messages = vec![result];
2221 assert!(input.build_messages(&State::default(), &request).is_empty());
2222 }
2223
2224 #[test]
2227 fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
2228 let messages = vec![
2229 Message::text(Role::System, "client instructions"),
2230 Message::text(Role::User, "initial task"),
2231 Message::text(Role::Assistant, "old response"),
2232 tool_call("call-1"),
2233 tool_result("call-1"),
2234 Message::text(Role::Assistant, "recent 1"),
2235 Message::text(Role::User, "recent 2"),
2236 Message::text(Role::Assistant, "recent 3"),
2237 Message::text(Role::User, "recent 4"),
2238 ];
2239
2240 let kept = trim_messages(&messages, 5);
2242
2243 assert_eq!(
2244 kept,
2245 vec![
2246 Message::text(Role::System, "client instructions"),
2247 Message::text(Role::User, "initial task"),
2248 tool_call("call-1"),
2249 tool_result("call-1"),
2250 Message::text(Role::Assistant, "recent 1"),
2251 Message::text(Role::User, "recent 2"),
2252 Message::text(Role::Assistant, "recent 3"),
2253 Message::text(Role::User, "recent 4"),
2254 ]
2255 );
2256 }
2257
2258 #[test]
2261 fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
2262 let messages = vec![
2263 Message::text(Role::System, "client instructions"),
2264 Message::text(Role::User, "initial task"),
2265 tool_call("x"),
2266 tool_result("x"),
2267 Message::text(Role::Assistant, "later"),
2268 tool_call("x"),
2269 tool_result("x"),
2270 ];
2271
2272 let kept = trim_messages(&messages, 4);
2274
2275 assert_eq!(
2276 kept,
2277 vec![
2278 Message::text(Role::System, "client instructions"),
2279 Message::text(Role::User, "initial task"),
2280 tool_call("x"),
2281 tool_result("x"),
2282 Message::text(Role::Assistant, "later"),
2283 tool_call("x"),
2284 tool_result("x"),
2285 ]
2286 );
2287 }
2288
2289 #[test]
2292 fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
2293 let messages = vec![
2294 Message::text(Role::System, "client instructions"),
2295 tool_call("orphan"),
2296 Message::text(Role::User, "initial task"),
2297 Message::text(Role::Assistant, "old response"),
2298 tool_result("orphan"),
2299 Message::text(Role::Assistant, "recent 1"),
2300 Message::text(Role::User, "recent 2"),
2301 ];
2302
2303 let kept = trim_messages(&messages, 3);
2304
2305 assert_eq!(
2306 kept,
2307 vec![
2308 Message::text(Role::System, "client instructions"),
2309 Message::text(Role::User, "initial task"),
2310 tool_result("orphan"),
2311 Message::text(Role::Assistant, "recent 1"),
2312 Message::text(Role::User, "recent 2"),
2313 ]
2314 );
2315 }
2316
2317 #[test]
2318 fn a_window_restates_the_routing_instruction_last() -> Result<()> {
2319 let contents = judged_contents(2)?;
2320
2321 assert_eq!(
2324 contents.last().map(String::as_str),
2325 Some(TRAILING_ROUTING_INSTRUCTION)
2326 );
2327 Ok(())
2328 }
2329
2330 #[test]
2335 fn a_window_drops_reasoning_but_keeps_visible_text_and_tool_pairs() {
2336 let messages = vec![
2337 Message::text(Role::System, "client instructions"),
2338 Message::text(Role::User, "initial task"),
2339 Message::new(
2340 Role::Assistant,
2341 vec![
2342 ContentBlock::Reasoning {
2343 text: "private chain of thought".to_string(),
2344 signature: None,
2345 details: Vec::new(),
2346 },
2347 ContentBlock::Text {
2348 text: "visible answer".to_string(),
2349 },
2350 ],
2351 ),
2352 tool_call("call-1"),
2353 tool_result("call-1"),
2354 Message::new(
2355 Role::Assistant,
2356 vec![ContentBlock::Reasoning {
2357 text: "reasoning-only turn".to_string(),
2358 signature: None,
2359 details: Vec::new(),
2360 }],
2361 ),
2362 Message::text(Role::User, "follow-up"),
2363 ];
2364 let request = Request {
2365 llm_request: LlmRequest {
2366 messages,
2367 ..LlmRequest::default()
2368 },
2369 raw_request: None,
2370 metadata: None,
2371 };
2372
2373 let built = TaskInput {
2374 recent_turn_window: Some(10),
2375 }
2376 .build_messages(&State::default(), &request);
2377
2378 assert!(
2379 !built
2380 .iter()
2381 .flat_map(|message| &message.content)
2382 .any(|block| matches!(block, ContentBlock::Reasoning { .. })),
2383 "{built:?}"
2384 );
2385 assert!(
2386 built
2387 .iter()
2388 .any(|message| message.text_content("\n").as_deref() == Some("visible answer"))
2389 );
2390 assert!(built.contains(&tool_call("call-1")));
2391 assert!(built.contains(&tool_result("call-1")));
2392 assert_eq!(built.len(), 7);
2395 assert!(built.iter().all(|message| !message.content.is_empty()));
2396 }
2397
2398 #[test]
2399 fn the_default_path_is_left_unchanged() -> Result<()> {
2400 let judge = capability_judge(None)?;
2403 let request = Request {
2404 llm_request: LlmRequest {
2405 messages: vec![Message::text(Role::User, "the task")],
2406 ..LlmRequest::default()
2407 },
2408 raw_request: None,
2409 metadata: None,
2410 };
2411
2412 let built = judge.build_request(&State::default(), &request);
2413
2414 assert_eq!(built.llm_request.messages.len(), 1);
2417 assert!(!built.llm_request.instructions.is_empty());
2418 assert!(
2419 !built
2420 .llm_request
2421 .messages
2422 .iter()
2423 .filter_map(|message| message.text_content("\n"))
2424 .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
2425 );
2426 Ok(())
2427 }
2428
2429 #[test]
2430 fn capability_judge_builds_a_structured_request() -> Result<()> {
2431 let judge = capability_judge(None)?;
2432 let request = Request {
2433 llm_request: LlmRequest {
2434 model: Some("inbound".to_string()),
2435 messages: vec![
2436 Message::text(Role::System, "client instructions"),
2437 Message::text(Role::Developer, "client developer instructions"),
2438 Message::text(Role::User, "initial task"),
2439 Message::text(Role::Assistant, "old response"),
2440 Message::text(Role::User, "old follow-up"),
2441 Message::text(Role::Assistant, "recent 1"),
2442 Message::text(Role::User, "recent 2"),
2443 Message::text(Role::Assistant, "recent 3"),
2444 Message::text(Role::User, "recent 4"),
2445 Message::text(Role::Assistant, "recent 5"),
2446 ],
2447 ..LlmRequest::default()
2448 },
2449 raw_request: None,
2450 metadata: None,
2451 };
2452 let judge_request = judge.build_request(&State::default(), &request);
2453
2454 assert_eq!(judge_request.llm_request.model, request.llm_request.model);
2455 assert_eq!(judge_request.llm_request.instructions.len(), 1);
2456 assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
2457 assert_eq!(
2458 judge_request.llm_request.instructions[0].content,
2459 InstructionBlock {
2460 role: Role::System,
2461 content: Message::text(Role::System, judge.contract().system_prompt()).content,
2462 }
2463 .content,
2464 );
2465 assert_eq!(judge_request.llm_request.messages.len(), 2);
2466 let contents = judge_request
2467 .llm_request
2468 .messages
2469 .iter()
2470 .filter_map(|message| message.text_content("\n"))
2471 .collect::<Vec<_>>();
2472 assert!(contents.contains(&"recent 4".to_string()));
2473 assert!(contents.contains(&"initial task".to_string()));
2474 assert!(!contents.contains(&"recent 5".to_string()));
2475 assert!(!contents.contains(&"client instructions".to_string()));
2476 assert_eq!(
2477 judge_request.llm_request.output.response_format,
2478 Some(judge.contract().response_format().clone())
2479 );
2480 assert_eq!(
2481 judge_request.llm_request.output.max_output_tokens,
2482 Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
2483 );
2484 Ok(())
2485 }
2486
2487 fn sample_value(spec: &Value) -> Value {
2488 if let Some(first) = spec
2489 .get("enum")
2490 .and_then(Value::as_array)
2491 .and_then(|values| values.first())
2492 {
2493 return first.clone();
2494 }
2495 match spec.get("type").and_then(Value::as_str) {
2496 Some("number") => serde_json::json!(0.5),
2497 Some("boolean") => serde_json::json!(false),
2498 _ => serde_json::json!("sample"),
2499 }
2500 }
2501
2502 fn schema_shaped_verdict(schema: &Value) -> Result<String> {
2503 let properties = schema
2504 .pointer("/json_schema/schema/properties")
2505 .and_then(Value::as_object)
2506 .ok_or_else(|| LibsyError::AlgorithmError {
2507 message: "packaged schema declares no properties".to_string(),
2508 })?;
2509 Ok(Value::Object(
2510 properties
2511 .iter()
2512 .map(|(name, spec)| (name.clone(), sample_value(spec)))
2513 .collect(),
2514 )
2515 .to_string())
2516 }
2517
2518 #[test]
2521 fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
2522 let contract =
2523 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
2524 let schema = contract.response_format();
2525 let reply = schema_shaped_verdict(schema)?;
2526 let judge: CapabilityJudge = StructuredJudge::new(
2527 TaskInput {
2528 recent_turn_window: None,
2529 },
2530 contract,
2531 SerdeDecoder::new(),
2532 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
2533 );
2534
2535 let verdict = judge.parse(&text_response(None, reply))?;
2536
2537 assert!(verdict.is_valid());
2538 assert!((0.0..=1.0).contains(&verdict.p_solve));
2539 Ok(())
2540 }
2541
2542 #[test]
2543 fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
2544 let contract =
2545 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
2546 let prompt = contract.system_prompt();
2547 let schema_name = contract
2548 .response_format()
2549 .pointer("/json_schema/name")
2550 .and_then(Value::as_str)
2551 .ok_or_else(|| LibsyError::AlgorithmError {
2552 message: "packaged response schema has no name".to_string(),
2553 })?;
2554 assert_eq!(schema_name, "CapabilityClassifierDecision");
2555 assert!(prompt.contains("SUP-1 [supported]"));
2556 assert!(prompt.contains("SUP-5 [supported]"));
2557 assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
2558 assert!(!prompt.contains("\"type\": \"object\""));
2559 assert!(!prompt.contains("\"json_schema\""));
2560 assert!(!prompt.contains(schema_name));
2561 let rule_values = contract
2562 .response_format()
2563 .pointer("/json_schema/schema/properties/primary_rule/enum")
2564 .and_then(Value::as_array)
2565 .ok_or_else(|| LibsyError::AlgorithmError {
2566 message: "rendered response schema has no primary rule enum".to_string(),
2567 })?;
2568 assert!(
2569 rule_values
2570 .iter()
2571 .any(|value| value.as_str() == Some("SUP-1"))
2572 );
2573 assert!(
2574 rule_values
2575 .iter()
2576 .any(|value| value.as_str() == Some("none"))
2577 );
2578 Ok(())
2579 }
2580}