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;
15use decision::DecisionClassifier;
16pub use decision::DecisionJudgeConfig;
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| Message {
182 role: Role::User,
183 content: message
184 .content
185 .iter()
186 .filter(|block| is_task_content(block))
187 .cloned()
188 .collect(),
189 })
190 .collect()
191}
192
193struct TaskInput {
195 recent_turn_window: Option<usize>,
196}
197
198impl TaskInput {
199 fn messages(&self, request: &Request) -> Vec<Message> {
201 let mut messages = match self.recent_turn_window {
204 Some(window) => trim_messages(&request.llm_request.messages, window),
205 None => task_messages(&request.llm_request.messages),
206 };
207 for message in &mut messages {
212 message
213 .content
214 .retain(|block| !matches!(block, ContentBlock::Reasoning { .. }));
215 }
216 messages.retain(|message| !message.content.is_empty());
217 messages
218 }
219}
220
221impl ClassifierInput for TaskInput {
222 fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
223 let mut messages = self.messages(request);
224 if self.recent_turn_window.is_some() {
228 messages.push(Message::text(
229 Role::User,
230 TRAILING_ROUTING_INSTRUCTION.to_string(),
231 ));
232 }
233 messages
234 }
235}
236
237struct TaskClassifierPolicy {
238 base_threshold: f64,
239 threshold_step: f64,
240 unmatched_steps: u8,
241}
242
243impl TaskClassifierPolicy {
244 fn new(config: &LlmCapabilityConfig) -> Self {
245 Self {
246 base_threshold: config.base_threshold,
247 threshold_step: config.threshold_step,
248 unmatched_steps: config.unmatched_steps,
249 }
250 }
251
252 fn threshold(&self, verdict: &TaskClassifierVerdict) -> Option<f64> {
254 let steps = verdict.boundary_steps(self.unmatched_steps)?;
255 Some(self.base_threshold + f64::from(steps) * self.threshold_step)
256 }
257}
258
259impl JudgePolicy for TaskClassifierPolicy {
260 type Verdict = TaskClassifierVerdict;
261
262 fn to_classification(
263 &self,
264 verdict: Option<&Self::Verdict>,
265 driver: &Driver,
266 ) -> Result<Classification> {
267 let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
270 return Ok(Classification::Ambiguous(vec![]));
271 };
272 let Some(threshold) = self.threshold(verdict) else {
275 return Ok(Classification::Ambiguous(vec![]));
276 };
277 let category = if verdict.p_solve >= threshold
278 || (threshold - verdict.p_solve).abs() <= f64::EPSILON
279 {
280 Category::Efficient
281 } else {
282 Category::Capable
283 };
284 let Some(target) = driver.models_for(&category).first().cloned() else {
287 return Ok(Classification::Ambiguous(vec![]));
288 };
289 Ok(Classification::Scores(vec![Score {
290 target,
291 confidence: 1.0,
292 category: Some(category),
293 }]))
294 }
295}
296
297fn capability_evidence(
299 policy: &TaskClassifierPolicy,
300 verdict: Option<&TaskClassifierVerdict>,
301) -> Option<Value> {
302 let verdict = verdict?;
303 let Some(threshold) = verdict
304 .is_valid()
305 .then(|| policy.threshold(verdict))
306 .flatten()
307 else {
308 return Some(serde_json::json!({
309 "source": "fail_open",
310 "reason_code": "invalid_verdict",
311 }));
312 };
313 Some(serde_json::json!({
314 "source": "llm-classifier",
315 "score": verdict.p_solve,
316 "threshold": threshold,
317 }))
318}
319
320#[derive(Clone, Debug)]
322pub struct TaskClassifierConfig {
323 pub judge: CapabilityJudgeConfig,
325 pub fail_open: bool,
327 pub classify_trigger: ClassifyTrigger,
329 pub message_hash_fallback: bool,
331 pub recent_turn_window: Option<usize>,
334}
335
336#[derive(Clone, Debug)]
338pub enum CapabilityJudgeConfig {
339 Llm(LlmCapabilityConfig),
341 Decision(DecisionJudgeConfig),
343}
344
345impl Default for CapabilityJudgeConfig {
346 fn default() -> Self {
347 Self::Llm(LlmCapabilityConfig::default())
348 }
349}
350
351#[derive(Clone, Debug)]
353pub struct LlmCapabilityConfig {
354 pub base_threshold: f64,
356 pub threshold_step: f64,
359 pub unmatched_steps: u8,
363 pub contract: ClassifierContractConfig,
365 pub max_output_tokens: u64,
367}
368
369impl Default for LlmCapabilityConfig {
370 fn default() -> Self {
371 Self {
372 base_threshold: 0.0,
373 threshold_step: 0.0,
374 unmatched_steps: DEFAULT_UNMATCHED_STEPS,
375 contract: ClassifierContractConfig::default(),
376 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
377 }
378 }
379}
380
381#[derive(Deserialize)]
383#[serde(deny_unknown_fields)]
384struct TaskClassifierConfigWire {
385 base_threshold: f64,
386 #[serde(default = "default_fail_open")]
387 fail_open: bool,
388 #[serde(default)]
389 threshold_step: f64,
390 #[serde(default = "default_unmatched_steps")]
391 unmatched_steps: u8,
392 #[serde(default)]
393 classify_trigger: ClassifyTrigger,
394 #[serde(default)]
395 message_hash_fallback: bool,
396 #[serde(default)]
397 recent_turn_window: Option<usize>,
398 #[serde(default)]
399 prompt: Option<String>,
400 #[serde(default)]
401 response_format_type: ClassifierResponseFormat,
402 #[serde(default = "default_judge_max_output_tokens")]
403 max_output_tokens: u64,
404}
405
406impl<'de> Deserialize<'de> for TaskClassifierConfig {
407 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
408 where
409 D: Deserializer<'de>,
410 {
411 let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
412 let mut contract = ClassifierContractConfig::default();
413 if let Some(prompt) = wire.prompt {
414 contract = contract.with_prompt(prompt);
415 }
416 contract = contract.with_response_format_type(wire.response_format_type);
417 Ok(Self {
418 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
419 base_threshold: wire.base_threshold,
420 threshold_step: wire.threshold_step,
421 unmatched_steps: wire.unmatched_steps,
422 contract,
423 max_output_tokens: wire.max_output_tokens,
424 }),
425 fail_open: wire.fail_open,
426 classify_trigger: wire.classify_trigger,
427 message_hash_fallback: wire.message_hash_fallback,
428 recent_turn_window: wire.recent_turn_window,
429 })
430 }
431}
432
433const fn default_fail_open() -> bool {
434 true
435}
436
437pub const DEFAULT_UNMATCHED_STEPS: u8 = 1;
439
440const fn default_unmatched_steps() -> u8 {
441 DEFAULT_UNMATCHED_STEPS
442}
443
444const fn default_judge_max_output_tokens() -> u64 {
445 DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
446}
447
448impl Default for TaskClassifierConfig {
449 fn default() -> Self {
450 Self {
451 judge: CapabilityJudgeConfig::default(),
452 fail_open: default_fail_open(),
453 classify_trigger: ClassifyTrigger::default(),
454 message_hash_fallback: false,
455 recent_turn_window: None,
456 }
457 }
458}
459
460impl LlmCapabilityConfig {
461 fn validate(&self) -> Result<()> {
462 if !(0.0..=1.0).contains(&self.base_threshold) {
463 return Err(LibsyError::AlgorithmError {
464 message: format!(
465 "base_threshold must be between 0 and 1, got {}",
466 self.base_threshold
467 ),
468 });
469 }
470 if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
471 return Err(LibsyError::AlgorithmError {
472 message: format!(
473 "threshold_step must be finite and greater than or equal to 0, got {}",
474 self.threshold_step
475 ),
476 });
477 }
478 if self.unmatched_steps > 2 {
479 return Err(LibsyError::AlgorithmError {
480 message: format!(
481 "unmatched_steps must be 0, 1, or 2, got {}",
482 self.unmatched_steps
483 ),
484 });
485 }
486 let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
487 if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
488 return Err(LibsyError::AlgorithmError {
489 message: format!(
490 "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
491 ),
492 });
493 }
494 if self.max_output_tokens == 0 {
495 return Err(LibsyError::AlgorithmError {
496 message: "max_output_tokens must be at least 1".to_string(),
497 });
498 }
499 Ok(())
500 }
501}
502
503impl TaskClassifierConfig {
504 fn validate(&self) -> Result<()> {
505 match &self.judge {
506 CapabilityJudgeConfig::Llm(config) => config.validate()?,
507 CapabilityJudgeConfig::Decision(config) => config.validate()?,
508 }
509 if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
511 return Err(LibsyError::AlgorithmError {
512 message:
513 "message_hash_fallback requires classify_trigger = new_session or user_turn"
514 .to_string(),
515 });
516 }
517 Ok(())
518 }
519}
520
521#[derive(Clone, Debug)]
523pub enum CustomClassifierPolicy {
524 TargetSelector {
526 selector: String,
528 },
529}
530
531impl CustomClassifierPolicy {
532 pub fn target_selector(selector: impl Into<String>) -> Self {
534 Self::TargetSelector {
535 selector: selector.into(),
536 }
537 }
538}
539
540#[derive(Clone, Debug)]
542pub struct CustomClassifierConfig {
543 pub prompt: String,
545 pub response_schema: Value,
547 pub policy: CustomClassifierPolicy,
549 pub classify_trigger: ClassifyTrigger,
551 pub message_hash_fallback: bool,
553 pub recent_turn_window: Option<usize>,
555 pub max_output_tokens: u64,
557}
558
559impl CustomClassifierConfig {
560 pub fn new(
562 prompt: impl Into<String>,
563 response_schema: Value,
564 policy: CustomClassifierPolicy,
565 ) -> Self {
566 Self {
567 prompt: prompt.into(),
568 response_schema,
569 policy,
570 classify_trigger: ClassifyTrigger::default(),
571 message_hash_fallback: false,
572 recent_turn_window: None,
573 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
574 }
575 }
576
577 fn validate(&self) -> Result<()> {
578 if self.max_output_tokens == 0 {
579 return Err(LibsyError::AlgorithmError {
580 message: "max_output_tokens must be at least 1".to_string(),
581 });
582 }
583 if self.message_hash_fallback && self.classify_trigger == ClassifyTrigger::EveryRequest {
585 return Err(LibsyError::AlgorithmError {
586 message:
587 "message_hash_fallback requires classify_trigger = new_session or user_turn"
588 .to_string(),
589 });
590 }
591 Ok(())
592 }
593}
594
595enum CustomPolicyRuntime {
596 TargetSelector(TargetSelectorPolicy),
597}
598
599impl JudgePolicy for CustomPolicyRuntime {
600 type Verdict = Value;
601
602 fn to_classification(
603 &self,
604 verdict: Option<&Self::Verdict>,
605 driver: &Driver,
606 ) -> Result<Classification> {
607 match self {
608 Self::TargetSelector(policy) => policy.to_classification(verdict, driver),
609 }
610 }
611}
612
613fn affinity_router(
615 trigger: ClassifyTrigger,
616 message_hash_fallback: bool,
617) -> Option<Arc<AffinityRouter>> {
618 let router = match trigger {
619 ClassifyTrigger::EveryRequest => return None,
620 ClassifyTrigger::NewSession => AffinityRouter::new(),
621 ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
622 };
623 let router = if message_hash_fallback {
624 router.with_message_hash_fallback()
625 } else {
626 router
627 };
628 Some(Arc::new(router))
629}
630
631pub struct LlmTaskClassifier {
633 route: FallThrough,
634 inner: Arc<dyn Classifier>,
636}
637
638struct ClassifierRouteConfig {
639 default_target: Category,
640 classify_trigger: ClassifyTrigger,
641 message_hash_fallback: bool,
642}
643
644pub struct DefaultCategoryClassifier(pub Category);
647
648#[async_trait]
649impl Classifier for DefaultCategoryClassifier {
650 async fn score(
651 &self,
652 _state: &mut State,
653 _request: &mut Request,
654 driver: &Driver,
655 ) -> Result<(Classification, Option<Response>)> {
656 let target = driver.first_model_for(&self.0)?;
657 driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"}));
658 Ok((
659 Classification::Scores(vec![Score {
660 target: target.clone(),
661 confidence: 0.0,
662 category: Some(self.0.clone()),
663 }]),
664 None,
665 ))
666 }
667}
668
669#[derive(Clone)]
671#[non_exhaustive]
672pub enum LlmClassifierConfig {
673 Capability {
675 config: TaskClassifierConfig,
677 },
678 Escalation {
680 contract: ClassifierContractConfig,
682 config: EscalationJudgeConfig,
684 max_output_tokens: u64,
686 },
687 Custom {
689 default_target: Category,
691 config: CustomClassifierConfig,
693 },
694}
695
696impl LlmTaskClassifier {
697 pub fn new(config: LlmClassifierConfig) -> Result<Self> {
704 match config {
705 LlmClassifierConfig::Capability { config } => Self::build_capability(config),
706 LlmClassifierConfig::Escalation {
707 contract,
708 config,
709 max_output_tokens,
710 } => Self::build_escalation(contract, config, max_output_tokens),
711 LlmClassifierConfig::Custom {
712 default_target,
713 config,
714 } => Self::build_custom(default_target, config),
715 }
716 }
717
718 fn build_capability(config: TaskClassifierConfig) -> Result<Self> {
719 config.validate()?;
720 let classify_trigger = config.classify_trigger;
721 let message_hash_fallback = config.message_hash_fallback;
722 let input = TaskInput {
723 recent_turn_window: config.recent_turn_window,
724 };
725 let classifier: Arc<dyn Classifier> = match config.judge {
726 CapabilityJudgeConfig::Llm(judge) => Arc::new(
727 JudgeClassifier::new(
728 StructuredJudge::new(
729 input,
730 Self::load_capability_contract(&judge.contract)?,
731 SerdeDecoder::new(),
732 JudgeRuntimeConfig::new(judge.max_output_tokens)?,
733 ),
734 TaskClassifierPolicy::new(&judge),
735 )
736 .with_error_recovery(config.fail_open)
737 .with_evidence(capability_evidence),
738 ),
739 CapabilityJudgeConfig::Decision(judge) => {
740 Arc::new(DecisionClassifier::new(judge, input, config.fail_open)?)
741 }
742 };
743 Self::from_classifier(
744 classifier,
745 ClassifierRouteConfig {
746 default_target: Category::Capable,
747 classify_trigger,
748 message_hash_fallback,
749 },
750 )
751 }
752
753 fn build_custom(default_target: Category, config: CustomClassifierConfig) -> Result<Self> {
754 Self::from_classifier(
755 Self::custom_classifier(&config, false)?,
756 ClassifierRouteConfig {
757 default_target,
758 classify_trigger: config.classify_trigger,
759 message_hash_fallback: config.message_hash_fallback,
760 },
761 )
762 }
763
764 pub(crate) fn custom_classifier(
766 config: &CustomClassifierConfig,
767 fail_open: bool,
768 ) -> Result<Arc<dyn Classifier>> {
769 config.validate()?;
770 let contract =
771 ClassifierContract::from_inner_schema(&config.prompt, config.response_schema.clone())?;
772 let policy = match &config.policy {
773 CustomClassifierPolicy::TargetSelector { selector } => {
774 CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(selector)?)
775 }
776 };
777 let classifier = JudgeClassifier::new(
778 StructuredJudge::new(
779 TaskInput {
780 recent_turn_window: config.recent_turn_window,
781 },
782 contract,
783 JsonSchemaDecoder::new(),
784 JudgeRuntimeConfig::new(config.max_output_tokens)?,
785 ),
786 policy,
787 )
788 .with_error_recovery(fail_open);
789 let classifier = if fail_open {
790 classifier.with_evidence(|_, verdict| {
791 verdict.map(|_| serde_json::json!({"source": "llm-classifier"}))
792 })
793 } else {
794 classifier
795 };
796 Ok(Arc::new(classifier))
797 }
798
799 fn build_escalation(
800 contract_config: ClassifierContractConfig,
801 config: EscalationJudgeConfig,
802 max_output_tokens: u64,
803 ) -> Result<Self> {
804 let inner = escalation::build_classifier(contract_config, config, max_output_tokens)?;
805 Ok(Self {
806 route: FallThrough::new_with_state()
807 .with_name(ALGORITHM_NAME)
808 .with_classifier(Arc::clone(&inner)),
809 inner,
810 })
811 }
812
813 fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
815 ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
816 }
817
818 fn from_classifier(inner: Arc<dyn Classifier>, config: ClassifierRouteConfig) -> Result<Self> {
820 if config.message_hash_fallback && config.classify_trigger == ClassifyTrigger::EveryRequest
822 {
823 return Err(LibsyError::AlgorithmError {
824 message:
825 "message_hash_fallback requires classify_trigger = new_session or user_turn"
826 .to_string(),
827 });
828 }
829 let mut route = FallThrough::new_with_state().with_name(ALGORITHM_NAME);
831 if let Some(affinity) =
832 affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
833 {
834 route = route
836 .with_processor(affinity.clone())
837 .with_classifier(affinity.clone());
838 }
839 let fallback = DefaultCategoryClassifier(config.default_target);
840 Ok(Self {
841 route: route
842 .with_classifier(inner.clone())
843 .with_classifier(Arc::new(fallback)),
844 inner,
845 })
846 }
847}
848
849#[async_trait]
850impl Classifier for LlmTaskClassifier {
851 async fn score(
852 &self,
853 state: &mut State,
854 request: &mut Request,
855 driver: &Driver,
856 ) -> Result<(Classification, Option<Response>)> {
857 self.inner.score(state, request, driver).await
858 }
859}
860
861#[async_trait]
862impl Algorithm for LlmTaskClassifier {
863 fn name(&self) -> &str {
864 "llm_task_classifier"
865 }
866
867 async fn route(
868 self: Arc<Self>,
869 driver: Driver,
870 request: Request,
871 ) -> Result<crate::RoutingOutcome> {
872 self.route.execute(driver, request).await
873 }
874}
875
876#[cfg(test)]
877mod tests {
878 use std::collections::HashMap;
879 use std::sync::Arc;
880
881 use parking_lot::Mutex;
882 use serde_json::Value;
883
884 use super::*;
885 use switchyard_protocol::{
886 ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ModelId, ToolCall,
887 ToolResult, completion_text, text_request, text_response,
888 };
889
890 use crate::algorithms::util::llm_judge::Judge;
891 use crate::core::testing::{Serve, test_drive_with_models};
892 use switchyard_protocol::{LlmResponse, Response};
893
894 const TEST_THRESHOLD: f64 = 0.5;
895
896 type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;
897
898 fn llm_config(base_threshold: f64) -> LlmCapabilityConfig {
899 LlmCapabilityConfig {
900 base_threshold,
901 ..LlmCapabilityConfig::default()
902 }
903 }
904
905 fn test_config(base_threshold: f64) -> TaskClassifierConfig {
906 TaskClassifierConfig {
907 judge: CapabilityJudgeConfig::Llm(llm_config(base_threshold)),
908 ..TaskClassifierConfig::default()
909 }
910 }
911
912 fn policy() -> TaskClassifierPolicy {
913 TaskClassifierPolicy::new(&llm_config(TEST_THRESHOLD))
914 }
915
916 fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
917 [
918 (Category::Judge, vec![ModelId::from("judge")]),
919 (Category::Efficient, vec![ModelId::from("efficient")]),
920 (Category::Capable, vec![ModelId::from("capable")]),
921 (
922 Category::Any,
923 vec![ModelId::from("efficient"), ModelId::from("capable")],
924 ),
925 ]
926 .into()
927 }
928
929 fn policy_driver() -> Driver {
930 Driver::new("test", Arc::new(runtime_models().into())).0
931 }
932
933 fn verdict(
934 p_solve: f64,
935 capability_boundary: &str,
936 primary_rule: &str,
937 ) -> TaskClassifierVerdict {
938 TaskClassifierVerdict {
939 crux: "test crux".to_string(),
940 primary_rule: primary_rule.to_string(),
941 capability_boundary: capability_boundary.to_string(),
942 p_solve,
943 }
944 }
945
946 fn selected(
947 policy: &TaskClassifierPolicy,
948 verdict: Option<&TaskClassifierVerdict>,
949 ) -> Result<ModelId> {
950 policy
951 .to_classification(verdict, &policy_driver())?
952 .argmax(false)?
953 .map(|score| score.target)
954 .ok_or_else(|| LibsyError::AlgorithmError {
955 message: "policy abstained".to_string(),
956 })
957 }
958
959 #[derive(Default)]
962 struct Recorder {
963 calls: Mutex<Vec<String>>,
964 call_roles: Mutex<Vec<(String, bool)>>,
965 judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
966 judge_system_prompts: Mutex<Vec<String>>,
967 }
968
969 impl Recorder {
970 fn calls(&self) -> Vec<String> {
971 self.calls.lock().clone()
972 }
973
974 fn call_roles(&self) -> Vec<(String, bool)> {
975 self.call_roles.lock().clone()
976 }
977
978 fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
979 self.judge_max_output_tokens.lock().clone()
980 }
981
982 fn judge_system_prompts(&self) -> Vec<String> {
983 self.judge_system_prompts.lock().clone()
984 }
985
986 fn serve(self: &Arc<Self>) -> impl Serve {
987 let recorder = Arc::clone(self);
988 move |model: ModelId, request: Request| {
989 let recorder = Arc::clone(&recorder);
990 async move {
991 let model = model.to_string();
992 recorder.calls.lock().push(model.clone());
993 recorder
994 .call_roles
995 .lock()
996 .push((model.clone(), model != "judge"));
997 let completion = if model == "judge" {
998 recorder
999 .judge_max_output_tokens
1000 .lock()
1001 .push(request.llm_request.output.max_output_tokens);
1002 recorder.judge_system_prompts.lock().extend(
1003 request
1004 .llm_request
1005 .instructions
1006 .first()
1007 .and_then(|instruction| {
1008 instruction.content.iter().find_map(|b| {
1009 if let ContentBlock::Text { text } = b {
1010 Some(text.clone())
1011 } else {
1012 None
1013 }
1014 })
1015 }),
1016 );
1017 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1018 } else {
1019 format!("answer from {model}")
1020 };
1021 Ok(Response {
1022 llm_response: LlmResponse::Agg(text_response(None, completion)),
1023 metadata: request.metadata,
1024 upstream_headers: http::HeaderMap::new(),
1025 })
1026 }
1027 }
1028 }
1029 }
1030
1031 fn unreachable_judge() -> impl Serve {
1033 |model: ModelId, request: Request| async move {
1034 let model = model.to_string();
1035 if model == "judge" {
1036 return Err(LlmClientError::Timeout {
1037 source: Box::new(std::io::Error::other("judge unreachable")),
1038 });
1039 }
1040 Ok(Response {
1041 llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
1042 metadata: request.metadata,
1043 upstream_headers: http::HeaderMap::new(),
1044 })
1045 }
1046 }
1047
1048 fn router() -> Result<Arc<LlmTaskClassifier>> {
1049 Ok(Arc::new(LlmTaskClassifier::new(
1050 LlmClassifierConfig::Capability {
1051 config: test_config(TEST_THRESHOLD),
1052 },
1053 )?))
1054 }
1055
1056 fn classify_request() -> Request {
1057 Request {
1058 llm_request: text_request(Some("auto".to_string()), "classify this task"),
1059 raw_request: None,
1060 metadata: None,
1061 }
1062 }
1063
1064 fn classify_session_request() -> Request {
1065 Request {
1066 metadata: Some(Metadata {
1067 session_id: Some("session-1".to_string()),
1068 ..Metadata::default()
1069 }),
1070 ..classify_request()
1071 }
1072 }
1073
1074 fn classify_follow_up_request() -> Request {
1075 let mut request = classify_request();
1076 request
1077 .llm_request
1078 .messages
1079 .push(Message::text(Role::Assistant, "I will add the test."));
1080 request.llm_request.messages.push(Message::text(
1081 Role::User,
1082 "Now run the test suite and report the result.",
1083 ));
1084 request
1085 }
1086
1087 #[tokio::test]
1088 async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
1089 let router = router()?;
1090
1091 let (selected_model, response) = test_drive_with_models(
1092 router,
1093 classify_request(),
1094 runtime_models(),
1095 unreachable_judge(),
1096 )
1097 .await?;
1098
1099 assert_eq!(selected_model, "capable");
1100 assert_eq!(
1101 response.llm_response.as_agg().map(completion_text),
1102 Some("answer from capable".to_string())
1103 );
1104 Ok(())
1105 }
1106
1107 #[tokio::test]
1108 async fn classifier_judges_each_request_without_affinity() -> Result<()> {
1109 let recorder = Arc::new(Recorder::default());
1110 let router = router()?;
1111 let request = classify_request();
1112 let models = runtime_models();
1113
1114 test_drive_with_models(
1115 router.clone(),
1116 request.clone(),
1117 models.clone(),
1118 recorder.serve(),
1119 )
1120 .await?;
1121 test_drive_with_models(router, request, models, recorder.serve()).await?;
1122
1123 assert_eq!(
1124 recorder.calls(),
1125 vec!["judge", "efficient", "judge", "efficient"]
1126 );
1127 assert_eq!(
1128 recorder.call_roles(),
1129 vec![
1130 ("judge".to_string(), false),
1131 ("efficient".to_string(), true),
1132 ("judge".to_string(), false),
1133 ("efficient".to_string(), true),
1134 ]
1135 );
1136 Ok(())
1137 }
1138
1139 #[tokio::test]
1140 async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1141 let recorder = Arc::new(Recorder::default());
1142 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1143 config: TaskClassifierConfig {
1144 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1145 max_output_tokens: 512,
1146 ..llm_config(TEST_THRESHOLD)
1147 }),
1148 ..test_config(TEST_THRESHOLD)
1149 },
1150 })?);
1151
1152 test_drive_with_models(
1153 router,
1154 classify_request(),
1155 runtime_models(),
1156 recorder.serve(),
1157 )
1158 .await?;
1159
1160 assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1161 Ok(())
1162 }
1163
1164 #[tokio::test]
1165 async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1166 let recorder = Arc::new(Recorder::default());
1167 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1168 config: TaskClassifierConfig {
1169 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1170 contract: ClassifierContractConfig::default()
1171 .with_prompt("Custom capability rubric."),
1172 ..llm_config(TEST_THRESHOLD)
1173 }),
1174 ..test_config(TEST_THRESHOLD)
1175 },
1176 })?);
1177
1178 test_drive_with_models(
1179 router,
1180 classify_request(),
1181 runtime_models(),
1182 recorder.serve(),
1183 )
1184 .await?;
1185
1186 let prompts = recorder.judge_system_prompts();
1187 assert_eq!(prompts.len(), 1);
1188 assert_eq!(prompts[0], "Custom capability rubric.");
1189 Ok(())
1190 }
1191
1192 #[tokio::test]
1193 async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1194 let recorder = Arc::new(Recorder::default());
1195 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1196 config: TaskClassifierConfig {
1197 classify_trigger: ClassifyTrigger::NewSession,
1198 ..test_config(TEST_THRESHOLD)
1199 },
1200 })?);
1201
1202 let request = classify_session_request();
1203 let models = runtime_models();
1204 test_drive_with_models(
1205 router.clone(),
1206 request.clone(),
1207 models.clone(),
1208 recorder.serve(),
1209 )
1210 .await?;
1211 test_drive_with_models(router, request, models, recorder.serve()).await?;
1212
1213 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1214 Ok(())
1215 }
1216
1217 #[tokio::test]
1218 async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1219 let recorder = Arc::new(Recorder::default());
1220 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1221 config: TaskClassifierConfig {
1222 classify_trigger: ClassifyTrigger::NewSession,
1223 message_hash_fallback: true,
1224 recent_turn_window: None,
1225 ..test_config(TEST_THRESHOLD)
1226 },
1227 })?);
1228
1229 let models = runtime_models();
1230 test_drive_with_models(
1231 router.clone(),
1232 classify_request(),
1233 models.clone(),
1234 recorder.serve(),
1235 )
1236 .await?;
1237 test_drive_with_models(
1238 router,
1239 classify_follow_up_request(),
1240 models,
1241 recorder.serve(),
1242 )
1243 .await?;
1244
1245 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1246 Ok(())
1247 }
1248
1249 #[tokio::test]
1250 async fn one_classifier_uses_each_requests_runtime_models() -> Result<()> {
1251 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1252 config: TaskClassifierConfig {
1253 classify_trigger: ClassifyTrigger::NewSession,
1254 ..test_config(TEST_THRESHOLD)
1255 },
1256 })?);
1257 let calls = Arc::new(Mutex::new(Vec::new()));
1258 let serve = |calls: Arc<Mutex<Vec<String>>>| {
1259 move |model: ModelId, _request: Request| {
1260 let calls = Arc::clone(&calls);
1261 async move {
1262 calls.lock().push(model.to_string());
1263 let text = if model.as_str().starts_with("judge-") {
1264 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1265 } else {
1266 model.to_string()
1267 };
1268 Ok(Response {
1269 llm_response: LlmResponse::Agg(text_response(None, text)),
1270 metadata: None,
1271 upstream_headers: Default::default(),
1272 })
1273 }
1274 }
1275 };
1276 let models = |suffix: &str| -> HashMap<Category, Vec<ModelId>> {
1277 [
1278 (
1279 Category::Judge,
1280 vec![ModelId::from(format!("judge-{suffix}"))],
1281 ),
1282 (
1283 Category::Efficient,
1284 vec![ModelId::from(format!("efficient-{suffix}"))],
1285 ),
1286 (
1287 Category::Capable,
1288 vec![ModelId::from(format!("capable-{suffix}"))],
1289 ),
1290 (
1291 Category::Any,
1292 vec![
1293 ModelId::from(format!("efficient-{suffix}")),
1294 ModelId::from(format!("capable-{suffix}")),
1295 ],
1296 ),
1297 ]
1298 .into()
1299 };
1300 let request = classify_session_request();
1301
1302 let (first, _) = test_drive_with_models(
1303 router.clone(),
1304 request.clone(),
1305 models("a"),
1306 serve(Arc::clone(&calls)),
1307 )
1308 .await?;
1309 let (second, _) =
1310 test_drive_with_models(router, request, models("b"), serve(Arc::clone(&calls))).await?;
1311
1312 assert_eq!(first, "efficient-a");
1313 assert_eq!(second, "efficient-b");
1314 assert_eq!(
1315 &*calls.lock(),
1316 &["judge-a", "efficient-a", "judge-b", "efficient-b"]
1317 );
1318 Ok(())
1319 }
1320
1321 #[tokio::test]
1322 async fn decision_judge_routes_on_relative_advantage() -> Result<()> {
1323 use crate::{Call, RuntimeModels, drive};
1324 use serde_json::json;
1325 use std::collections::BTreeMap;
1326 use std::sync::atomic::{AtomicUsize, Ordering};
1327 use switchyard_protocol::{
1328 BooleanEstimate, DecisionAnswer, DecisionResponse, DecisionValue, Probability,
1329 ProviderConfidence,
1330 };
1331
1332 let settings = DecisionJudgeConfig {
1333 cutoff: 0.4,
1334 instructions: None,
1335 criteria: None,
1336 candidates: BTreeMap::from([
1337 ("a".into(), "capable".into()),
1338 ("b".into(), "efficient".into()),
1339 ("c".into(), "third".into()),
1340 ]),
1341 evidence: json!({
1342 "candidate_descriptions": {"a": "Profile A", "b": "Profile B", "c": "Profile C"},
1343 "reference_cases": [{"task": "Independent example", "observed_solved": {"a": true, "b": null, "c": false}}],
1344 "summaries": [],
1345 }),
1346 };
1347 let mut request = classify_session_request();
1348 request.llm_request.messages.push(Message {
1349 role: Role::Assistant,
1350 content: vec![
1351 ContentBlock::Reasoning {
1352 text: "private reasoning".into(),
1353 signature: None,
1354 details: vec![],
1355 },
1356 ContentBlock::Text {
1357 text: "visible progress".into(),
1358 },
1359 ],
1360 });
1361 request
1362 .llm_request
1363 .messages
1364 .push(Message::text(Role::User, "latest follow-up"));
1365 let original_messages = request.llm_request.messages.clone();
1366 let expected_task = json!([
1367 original_messages[0].clone(),
1368 Message::text(Role::Assistant, "visible progress"),
1369 Message::text(Role::User, "latest follow-up"),
1370 ]);
1371
1372 for (name, score, fail_open, expected) in [
1373 ("above", 0.45, true, Some("capable")),
1374 ("equal", 0.4, true, Some("efficient")),
1375 ("below", 0.39, false, Some("efficient")),
1376 ("missing answer", 0.0, true, Some("capable")),
1377 ("no distribution", 0.0, true, Some("capable")),
1378 ("wrong type", 0.0, true, Some("capable")),
1379 ("out of range", 1.1, true, Some("capable")),
1380 ("nonfinite", f64::NAN, false, Some("capable")),
1381 ("provider error", 0.0, true, Some("capable")),
1382 ("provider error", 0.0, false, None),
1383 ("dropped reply", 0.0, true, Some("capable")),
1384 ("missing candidate", 0.0, true, None),
1385 ] {
1386 let mut config = settings.clone();
1387 if name == "missing candidate" {
1388 config.candidates.remove("a");
1389 }
1390 if name == "below" {
1391 config.instructions = Some(
1392 json!({"policy": "Judge capable-only success using the supplied evidence."}),
1393 );
1394 }
1395 let override_instructions = config.instructions.clone();
1396 let router: Arc<dyn Algorithm> =
1397 Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1398 config: TaskClassifierConfig {
1399 judge: CapabilityJudgeConfig::Decision(config),
1400 fail_open,
1401 recent_turn_window: None,
1402 classify_trigger: ClassifyTrigger::NewSession,
1403 ..TaskClassifierConfig::default()
1404 },
1405 })?);
1406 let calls = AtomicUsize::new(0);
1407 let serve = |call| {
1408 let calls = &calls;
1409 let expected_task = &expected_task;
1410 let settings = &settings;
1411 let override_instructions = &override_instructions;
1412 async move {
1413 let Call::Decision(call) = call else {
1414 panic!("judge must use the decision step")
1415 };
1416 calls.fetch_add(1, Ordering::SeqCst);
1417 assert_eq!(call.model, "judge");
1418 assert_eq!(call.request.model, Some("judge".into()));
1419 assert_eq!(
1420 call.request.context,
1421 json!({
1422 "task": expected_task,
1423 "session": "session-1",
1424 "candidates": ["a", "b", "c"],
1425 "comparison": {"capable": "a", "efficient": "b"},
1426 "evidence": settings.evidence,
1427 })
1428 );
1429 assert_eq!(call.request.questions.len(), 1);
1430 let question = &call.request.questions["route"];
1431 let switchyard_protocol::DecisionKind::Choice { options } = &question.kind
1432 else {
1433 panic!("expected a choice")
1434 };
1435 assert_eq!(
1436 options
1437 .iter()
1438 .map(|option| option.id.as_str())
1439 .collect::<Vec<_>>(),
1440 ["advantage", "no_advantage"]
1441 );
1442 if let Some(instructions) = override_instructions {
1443 assert_eq!(&question.instructions, instructions);
1444 } else {
1445 for key in ["setting", "evidence", "comparison", "boundary", "policy"] {
1446 assert!(question.instructions[key].as_str().is_some());
1447 }
1448 }
1449 if name == "provider error" {
1450 return call
1451 .respond(std::future::ready(Err(LibsyError::client_call(
1452 "judge",
1453 LlmClientError::General("private provider body".into()),
1454 ))))
1455 .await;
1456 }
1457 if name == "dropped reply" {
1458 drop(call);
1459 return Ok(());
1460 }
1461 let value = if name == "wrong type" {
1462 DecisionValue::Boolean(BooleanEstimate::Value(true))
1463 } else {
1464 DecisionValue::Choice {
1465 selected: "no_advantage".into(),
1466 probabilities: (name != "no distribution").then(|| {
1467 BTreeMap::from([
1468 ("advantage".into(), Probability(score)),
1469 ("no_advantage".into(), Probability(1.0 - score)),
1470 ])
1471 }),
1472 }
1473 };
1474 let answers = if name == "missing answer" {
1475 BTreeMap::new()
1476 } else {
1477 BTreeMap::from([(
1478 "route".into(),
1479 DecisionAnswer {
1480 value,
1481 provider_confidence: Some(ProviderConfidence(0.99)),
1482 },
1483 )])
1484 };
1485 call.respond(std::future::ready(Ok(DecisionResponse {
1486 id: None,
1487 model: Some("provider-judge".into()),
1488 answers,
1489 usage: Default::default(),
1490 })))
1491 .await
1492 }
1493 };
1494 let models = Arc::new(RuntimeModels::new(runtime_models()));
1495 let result = drive(router.clone(), request.clone(), models.clone(), &serve).await;
1496 assert_eq!(
1497 calls.load(Ordering::SeqCst),
1498 usize::from(name != "missing candidate"),
1499 "{name}"
1500 );
1501 let Some(expected) = expected else {
1502 if name == "missing candidate" {
1503 assert!(
1504 matches!(result, Err(LibsyError::AlgorithmError { message }) if message.contains("candidate is missing"))
1505 );
1506 } else {
1507 assert!(
1508 matches!(result, Err(LibsyError::ClientCall { .. })),
1509 "{name}"
1510 );
1511 }
1512 continue;
1513 };
1514 let outcome = result?;
1515 assert_eq!(outcome.selected_model_id()?, expected, "{name}");
1516 assert!(outcome.response.is_none());
1517 assert_eq!(outcome.request.llm_request.messages, original_messages);
1518 let evidence = outcome
1519 .metadata
1520 .and_then(|metadata| metadata.evidence)
1521 .expect("routing evidence");
1522 if matches!(name, "above" | "equal" | "below") {
1523 assert_eq!(evidence["source"], "decision_classifier");
1524 assert_eq!(evidence["verdict"], "relative_advantage");
1525 assert_eq!(evidence["threshold"], settings.cutoff);
1526 assert_eq!(evidence["score"], score);
1527 } else {
1528 let reason = match name {
1529 "provider error" => "client_error",
1530 "dropped reply" => "call_error",
1531 _ => "invalid_verdict",
1532 };
1533 assert_eq!(
1534 evidence,
1535 json!({"source": "fail_open", "reason_code": reason})
1536 );
1537 }
1538 let retained = drive(router, request.clone(), models, &serve).await?;
1539 assert_eq!(retained.selected_model_id()?, expected);
1540 assert_eq!(
1541 calls.load(Ordering::SeqCst),
1542 1,
1543 "affinity should skip the judge: {name}"
1544 );
1545 }
1546 for cutoff in [-0.1, 1.1, f64::NAN] {
1547 let mut config = settings.clone();
1548 config.cutoff = cutoff;
1549 assert!(
1550 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1551 config: TaskClassifierConfig {
1552 judge: CapabilityJudgeConfig::Decision(config),
1553 ..TaskClassifierConfig::default()
1554 },
1555 })
1556 .is_err()
1557 );
1558 }
1559 let mut duplicate = settings.clone();
1560 duplicate
1561 .candidates
1562 .insert("duplicate-a".into(), "capable".into());
1563 assert!(duplicate.validate().is_err());
1564 let mut custom = settings.clone();
1565 custom.candidates.remove("c");
1566 custom.criteria = Some(BTreeMap::from([
1567 ("a".into(), json!("Start capable")),
1568 ("b".into(), json!("Start efficient")),
1569 ]));
1570 assert!(custom.validate().is_err());
1571 custom.instructions = Some(json!("Choose a profile."));
1572 assert!(custom.validate().is_ok());
1573 custom.candidates.remove("b");
1574 assert!(custom.validate().is_err());
1575 Ok(())
1576 }
1577
1578 #[test]
1579 fn the_threshold_boundary_is_inclusive() -> Result<()> {
1580 let policy = policy();
1581 let at_threshold = verdict(0.5, "supported", "SUP-1");
1582 let below_threshold = verdict(0.49, "supported", "SUP-1");
1583 assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1584 assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1585 Ok(())
1586 }
1587
1588 #[test]
1589 fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1590 let borderline = verdict(0.5, "supported", "SUP-1");
1591 let strict = TaskClassifierPolicy::new(&llm_config(0.9));
1592 let lenient = TaskClassifierPolicy::new(&llm_config(0.1));
1593 assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1594 assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1595 Ok(())
1596 }
1597
1598 #[test]
1599 fn classifier_config_rejects_unknown_fields() {
1600 let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1601 "base_threshold": 0.5,
1602 "classifier_magic": true,
1603 }))
1604 .expect_err("unknown classifier fields must be rejected");
1605
1606 assert!(
1607 error
1608 .to_string()
1609 .contains("unknown field `classifier_magic`"),
1610 "{error}"
1611 );
1612 }
1613
1614 #[test]
1615 fn invalid_classifier_config_is_rejected() -> Result<()> {
1616 for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1617 assert!(
1618 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1619 config: test_config(bad),
1620 })
1621 .is_err(),
1622 "base threshold {bad} should be rejected"
1623 );
1624 }
1625 for config in [
1626 TaskClassifierConfig {
1627 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1628 base_threshold: 0.5,
1629 threshold_step: -0.1,
1630 ..LlmCapabilityConfig::default()
1631 }),
1632 ..TaskClassifierConfig::default()
1633 },
1634 TaskClassifierConfig {
1635 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1636 base_threshold: 0.8,
1637 threshold_step: 0.11,
1638 ..LlmCapabilityConfig::default()
1639 }),
1640 ..TaskClassifierConfig::default()
1641 },
1642 TaskClassifierConfig {
1643 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1644 base_threshold: 0.5,
1645 ..LlmCapabilityConfig::default()
1646 }),
1647 message_hash_fallback: true,
1648 ..TaskClassifierConfig::default()
1649 },
1650 TaskClassifierConfig {
1651 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1652 base_threshold: 0.5,
1653 max_output_tokens: 0,
1654 ..LlmCapabilityConfig::default()
1655 }),
1656 ..TaskClassifierConfig::default()
1657 },
1658 ] {
1659 assert!(LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).is_err());
1660 }
1661 for base_threshold in [0.0, 1.0] {
1662 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1663 config: test_config(base_threshold),
1664 })?;
1665 }
1666 Ok(())
1667 }
1668
1669 #[test]
1670 fn message_hash_fallback_accepts_retaining_triggers() -> Result<()> {
1671 for trigger in [ClassifyTrigger::NewSession, ClassifyTrigger::UserTurn] {
1675 let config = TaskClassifierConfig {
1676 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1677 base_threshold: 0.5,
1678 ..LlmCapabilityConfig::default()
1679 }),
1680 classify_trigger: trigger,
1681 message_hash_fallback: true,
1682 ..TaskClassifierConfig::default()
1683 };
1684 LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).map_err(
1685 |error| LibsyError::AlgorithmError {
1686 message: format!("{trigger:?} with message_hash_fallback rejected: {error}"),
1687 },
1688 )?;
1689 }
1690 let every_request = TaskClassifierConfig {
1691 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1692 base_threshold: 0.5,
1693 ..LlmCapabilityConfig::default()
1694 }),
1695 classify_trigger: ClassifyTrigger::EveryRequest,
1696 message_hash_fallback: true,
1697 ..TaskClassifierConfig::default()
1698 };
1699 assert!(
1700 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1701 config: every_request
1702 })
1703 .is_err(),
1704 "every_request with message_hash_fallback should stay rejected"
1705 );
1706 Ok(())
1707 }
1708
1709 #[test]
1710 fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1711 let policy = policy();
1712 let inconsistent_rule = TaskClassifierVerdict {
1713 capability_boundary: "uncertain".to_string(),
1714 ..verdict(1.0, "supported", "SUP-1")
1715 };
1716 let empty_crux = TaskClassifierVerdict {
1717 crux: " ".to_string(),
1718 ..verdict(1.0, "supported", "SUP-1")
1719 };
1720 let unusable = [
1721 Some(verdict(1.1, "supported", "SUP-1")),
1722 Some(inconsistent_rule),
1723 Some(empty_crux),
1724 None,
1725 ];
1726 for verdict in unusable {
1727 let classification = policy.to_classification(verdict.as_ref(), &policy_driver())?;
1728 assert!(matches!(classification, Classification::Ambiguous(_)));
1729 assert!(classification.argmax(false)?.is_none());
1730 assert!(classification.argmax(true)?.is_none());
1731 }
1732 Ok(())
1733 }
1734
1735 #[test]
1736 fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
1737 let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
1738 threshold_step: 0.1,
1739 ..llm_config(0.4)
1740 });
1741
1742 assert_eq!(
1743 selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
1744 "efficient"
1745 );
1746 assert_eq!(
1747 selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
1748 "capable"
1749 );
1750 assert_eq!(
1751 selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
1752 "efficient"
1753 );
1754 assert_eq!(
1755 selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
1756 "efficient"
1757 );
1758 assert_eq!(
1759 selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
1760 "capable"
1761 );
1762 assert_eq!(
1763 selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
1764 "efficient"
1765 );
1766 Ok(())
1767 }
1768
1769 #[test]
1770 fn unmatched_steps_raise_only_the_unmatched_threshold() -> Result<()> {
1771 let policy = TaskClassifierPolicy::new(&LlmCapabilityConfig {
1772 threshold_step: 0.1,
1773 unmatched_steps: 2,
1774 ..llm_config(0.75)
1775 });
1776
1777 assert_eq!(
1779 selected(&policy, Some(&verdict(0.85, "unmatched", "none")))?,
1780 "capable"
1781 );
1782 assert_eq!(
1783 selected(&policy, Some(&verdict(0.95, "unmatched", "none")))?,
1784 "efficient"
1785 );
1786 assert_eq!(
1788 selected(&policy, Some(&verdict(0.85, "uncertain", "UNC-1")))?,
1789 "efficient"
1790 );
1791 assert_eq!(
1792 selected(&policy, Some(&verdict(0.75, "supported", "SUP-2")))?,
1793 "efficient"
1794 );
1795 Ok(())
1796 }
1797
1798 #[test]
1799 fn unmatched_steps_default_to_one_and_reject_values_above_two() {
1800 let parsed: TaskClassifierConfig =
1801 serde_json::from_value(serde_json::json!({ "base_threshold": 0.5 }))
1802 .expect("config without unmatched_steps parses");
1803 let CapabilityJudgeConfig::Llm(judge) = &parsed.judge else {
1804 panic!("a flat capability config builds an LLM judge");
1805 };
1806 assert_eq!(judge.unmatched_steps, 1);
1807
1808 let invalid = TaskClassifierConfig {
1809 judge: CapabilityJudgeConfig::Llm(LlmCapabilityConfig {
1810 unmatched_steps: 3,
1811 ..llm_config(0.5)
1812 }),
1813 ..TaskClassifierConfig::default()
1814 };
1815 assert!(invalid.validate().is_err());
1816 }
1817
1818 fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
1821 Ok(StructuredJudge::new(
1822 TaskInput { recent_turn_window },
1823 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
1824 SerdeDecoder::new(),
1825 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1826 ))
1827 }
1828
1829 fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
1830 let judge = capability_judge(Some(recent_turn_window))?;
1831 let request = Request {
1832 llm_request: LlmRequest {
1833 messages: vec![
1834 Message::text(Role::System, "client instructions"),
1835 Message::text(
1836 Role::User,
1837 "<environment_context>setup</environment_context>",
1838 ),
1839 Message::text(Role::User, "initial task"),
1840 Message::text(Role::Assistant, "old response"),
1841 Message::text(Role::User, "old follow-up"),
1842 Message::text(Role::Assistant, "recent 1"),
1843 Message::text(Role::User, "recent 2"),
1844 ],
1845 ..LlmRequest::default()
1846 },
1847 raw_request: None,
1848 metadata: None,
1849 };
1850 Ok(judge
1851 .build_request(&State::default(), &request)
1852 .llm_request
1853 .messages
1854 .iter()
1855 .filter_map(|message| message.text_content("\n"))
1856 .collect())
1857 }
1858
1859 #[test]
1860 fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
1861 let contents = judged_contents(2)?;
1863 assert!(contents.contains(&"client instructions".to_string()));
1864 assert!(contents.contains(&"initial task".to_string()));
1865 assert!(contents.contains(&"recent 1".to_string()));
1866 assert!(contents.contains(&"recent 2".to_string()));
1867 assert!(!contents.contains(&"old response".to_string()));
1868 Ok(())
1869 }
1870
1871 #[test]
1872 fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
1873 let contents = judged_contents(0)?;
1874 assert!(contents.contains(&"client instructions".to_string()));
1875 assert!(contents.contains(&"initial task".to_string()));
1876 assert!(!contents.contains(&"recent 2".to_string()));
1877 Ok(())
1878 }
1879
1880 fn tool_call(id: &str) -> Message {
1881 Message {
1882 role: Role::Assistant,
1883 content: vec![ContentBlock::ToolCall(ToolCall {
1884 id: id.to_string(),
1885 name: "search".to_string(),
1886 arguments: Value::Null,
1887 })],
1888 }
1889 }
1890
1891 fn tool_result(id: &str) -> Message {
1892 Message {
1893 role: Role::Tool,
1894 content: vec![ContentBlock::ToolResult(ToolResult {
1895 tool_call_id: id.to_string(),
1896 content: vec![ContentBlock::Text {
1897 text: "tool output".to_string(),
1898 }],
1899 is_error: None,
1900 })],
1901 }
1902 }
1903
1904 #[test]
1905 fn default_task_input_keeps_user_content_around_tool_results() {
1906 let context = Message::text(
1907 Role::User,
1908 "<environment_context><current_date>2026-10-31</current_date></environment_context>",
1909 );
1910 let mut result = tool_result("call-1");
1911 result.role = Role::User;
1912 let mut mixed = result.clone();
1913 mixed.content.push(ContentBlock::Text {
1914 text: "latest follow-up".to_string(),
1915 });
1916 let input = TaskInput {
1917 recent_turn_window: None,
1918 };
1919 let mut request = Request {
1920 llm_request: LlmRequest {
1921 messages: vec![
1922 context.clone(),
1923 result.clone(),
1924 Message::text(Role::User, "initial task"),
1925 tool_call("call-1"),
1926 mixed,
1927 result.clone(),
1928 context,
1929 ],
1930 ..LlmRequest::default()
1931 },
1932 ..Request::default()
1933 };
1934 assert_eq!(
1935 input.build_messages(&State::default(), &request),
1936 vec![
1937 Message::text(Role::User, "initial task"),
1938 Message::text(Role::User, "latest follow-up"),
1939 ]
1940 );
1941 request.llm_request.messages = vec![result];
1942 assert!(input.build_messages(&State::default(), &request).is_empty());
1943 }
1944
1945 #[test]
1948 fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
1949 let messages = vec![
1950 Message::text(Role::System, "client instructions"),
1951 Message::text(Role::User, "initial task"),
1952 Message::text(Role::Assistant, "old response"),
1953 tool_call("call-1"),
1954 tool_result("call-1"),
1955 Message::text(Role::Assistant, "recent 1"),
1956 Message::text(Role::User, "recent 2"),
1957 Message::text(Role::Assistant, "recent 3"),
1958 Message::text(Role::User, "recent 4"),
1959 ];
1960
1961 let kept = trim_messages(&messages, 5);
1963
1964 assert_eq!(
1965 kept,
1966 vec![
1967 Message::text(Role::System, "client instructions"),
1968 Message::text(Role::User, "initial task"),
1969 tool_call("call-1"),
1970 tool_result("call-1"),
1971 Message::text(Role::Assistant, "recent 1"),
1972 Message::text(Role::User, "recent 2"),
1973 Message::text(Role::Assistant, "recent 3"),
1974 Message::text(Role::User, "recent 4"),
1975 ]
1976 );
1977 }
1978
1979 #[test]
1982 fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
1983 let messages = vec![
1984 Message::text(Role::System, "client instructions"),
1985 Message::text(Role::User, "initial task"),
1986 tool_call("x"),
1987 tool_result("x"),
1988 Message::text(Role::Assistant, "later"),
1989 tool_call("x"),
1990 tool_result("x"),
1991 ];
1992
1993 let kept = trim_messages(&messages, 4);
1995
1996 assert_eq!(
1997 kept,
1998 vec![
1999 Message::text(Role::System, "client instructions"),
2000 Message::text(Role::User, "initial task"),
2001 tool_call("x"),
2002 tool_result("x"),
2003 Message::text(Role::Assistant, "later"),
2004 tool_call("x"),
2005 tool_result("x"),
2006 ]
2007 );
2008 }
2009
2010 #[test]
2013 fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
2014 let messages = vec![
2015 Message::text(Role::System, "client instructions"),
2016 tool_call("orphan"),
2017 Message::text(Role::User, "initial task"),
2018 Message::text(Role::Assistant, "old response"),
2019 tool_result("orphan"),
2020 Message::text(Role::Assistant, "recent 1"),
2021 Message::text(Role::User, "recent 2"),
2022 ];
2023
2024 let kept = trim_messages(&messages, 3);
2025
2026 assert_eq!(
2027 kept,
2028 vec![
2029 Message::text(Role::System, "client instructions"),
2030 Message::text(Role::User, "initial task"),
2031 tool_result("orphan"),
2032 Message::text(Role::Assistant, "recent 1"),
2033 Message::text(Role::User, "recent 2"),
2034 ]
2035 );
2036 }
2037
2038 #[test]
2039 fn a_window_restates_the_routing_instruction_last() -> Result<()> {
2040 let contents = judged_contents(2)?;
2041
2042 assert_eq!(
2045 contents.last().map(String::as_str),
2046 Some(TRAILING_ROUTING_INSTRUCTION)
2047 );
2048 Ok(())
2049 }
2050
2051 #[test]
2056 fn a_window_drops_reasoning_but_keeps_visible_text_and_tool_pairs() {
2057 let messages = vec![
2058 Message::text(Role::System, "client instructions"),
2059 Message::text(Role::User, "initial task"),
2060 Message {
2061 role: Role::Assistant,
2062 content: vec![
2063 ContentBlock::Reasoning {
2064 text: "private chain of thought".to_string(),
2065 signature: None,
2066 details: Vec::new(),
2067 },
2068 ContentBlock::Text {
2069 text: "visible answer".to_string(),
2070 },
2071 ],
2072 },
2073 tool_call("call-1"),
2074 tool_result("call-1"),
2075 Message {
2076 role: Role::Assistant,
2077 content: vec![ContentBlock::Reasoning {
2078 text: "reasoning-only turn".to_string(),
2079 signature: None,
2080 details: Vec::new(),
2081 }],
2082 },
2083 Message::text(Role::User, "follow-up"),
2084 ];
2085 let request = Request {
2086 llm_request: LlmRequest {
2087 messages,
2088 ..LlmRequest::default()
2089 },
2090 raw_request: None,
2091 metadata: None,
2092 };
2093
2094 let built = TaskInput {
2095 recent_turn_window: Some(10),
2096 }
2097 .build_messages(&State::default(), &request);
2098
2099 assert!(
2100 !built
2101 .iter()
2102 .flat_map(|message| &message.content)
2103 .any(|block| matches!(block, ContentBlock::Reasoning { .. })),
2104 "{built:?}"
2105 );
2106 assert!(
2107 built
2108 .iter()
2109 .any(|message| message.text_content("\n").as_deref() == Some("visible answer"))
2110 );
2111 assert!(built.contains(&tool_call("call-1")));
2112 assert!(built.contains(&tool_result("call-1")));
2113 assert_eq!(built.len(), 7);
2116 assert!(built.iter().all(|message| !message.content.is_empty()));
2117 }
2118
2119 #[test]
2120 fn the_default_path_is_left_unchanged() -> Result<()> {
2121 let judge = capability_judge(None)?;
2124 let request = Request {
2125 llm_request: LlmRequest {
2126 messages: vec![Message::text(Role::User, "the task")],
2127 ..LlmRequest::default()
2128 },
2129 raw_request: None,
2130 metadata: None,
2131 };
2132
2133 let built = judge.build_request(&State::default(), &request);
2134
2135 assert_eq!(built.llm_request.messages.len(), 1);
2138 assert!(!built.llm_request.instructions.is_empty());
2139 assert!(
2140 !built
2141 .llm_request
2142 .messages
2143 .iter()
2144 .filter_map(|message| message.text_content("\n"))
2145 .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
2146 );
2147 Ok(())
2148 }
2149
2150 #[test]
2151 fn capability_judge_builds_a_structured_request() -> Result<()> {
2152 let judge = capability_judge(None)?;
2153 let request = Request {
2154 llm_request: LlmRequest {
2155 model: Some("inbound".to_string()),
2156 messages: vec![
2157 Message::text(Role::System, "client instructions"),
2158 Message::text(Role::Developer, "client developer instructions"),
2159 Message::text(Role::User, "initial task"),
2160 Message::text(Role::Assistant, "old response"),
2161 Message::text(Role::User, "old follow-up"),
2162 Message::text(Role::Assistant, "recent 1"),
2163 Message::text(Role::User, "recent 2"),
2164 Message::text(Role::Assistant, "recent 3"),
2165 Message::text(Role::User, "recent 4"),
2166 Message::text(Role::Assistant, "recent 5"),
2167 ],
2168 ..LlmRequest::default()
2169 },
2170 raw_request: None,
2171 metadata: None,
2172 };
2173 let judge_request = judge.build_request(&State::default(), &request);
2174
2175 assert_eq!(judge_request.llm_request.model, request.llm_request.model);
2176 assert_eq!(judge_request.llm_request.instructions.len(), 1);
2177 assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
2178 assert_eq!(
2179 judge_request.llm_request.instructions[0].content,
2180 InstructionBlock {
2181 role: Role::System,
2182 content: Message::text(Role::System, judge.contract().system_prompt()).content,
2183 }
2184 .content,
2185 );
2186 assert_eq!(judge_request.llm_request.messages.len(), 2);
2187 let contents = judge_request
2188 .llm_request
2189 .messages
2190 .iter()
2191 .filter_map(|message| message.text_content("\n"))
2192 .collect::<Vec<_>>();
2193 assert!(contents.contains(&"recent 4".to_string()));
2194 assert!(contents.contains(&"initial task".to_string()));
2195 assert!(!contents.contains(&"recent 5".to_string()));
2196 assert!(!contents.contains(&"client instructions".to_string()));
2197 assert_eq!(
2198 judge_request.llm_request.output.response_format,
2199 Some(judge.contract().response_format().clone())
2200 );
2201 assert_eq!(
2202 judge_request.llm_request.output.max_output_tokens,
2203 Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
2204 );
2205 Ok(())
2206 }
2207
2208 fn sample_value(spec: &Value) -> Value {
2209 if let Some(first) = spec
2210 .get("enum")
2211 .and_then(Value::as_array)
2212 .and_then(|values| values.first())
2213 {
2214 return first.clone();
2215 }
2216 match spec.get("type").and_then(Value::as_str) {
2217 Some("number") => serde_json::json!(0.5),
2218 Some("boolean") => serde_json::json!(false),
2219 _ => serde_json::json!("sample"),
2220 }
2221 }
2222
2223 fn schema_shaped_verdict(schema: &Value) -> Result<String> {
2224 let properties = schema
2225 .pointer("/json_schema/schema/properties")
2226 .and_then(Value::as_object)
2227 .ok_or_else(|| LibsyError::AlgorithmError {
2228 message: "packaged schema declares no properties".to_string(),
2229 })?;
2230 Ok(Value::Object(
2231 properties
2232 .iter()
2233 .map(|(name, spec)| (name.clone(), sample_value(spec)))
2234 .collect(),
2235 )
2236 .to_string())
2237 }
2238
2239 #[test]
2242 fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
2243 let contract =
2244 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
2245 let schema = contract.response_format();
2246 let reply = schema_shaped_verdict(schema)?;
2247 let judge: CapabilityJudge = StructuredJudge::new(
2248 TaskInput {
2249 recent_turn_window: None,
2250 },
2251 contract,
2252 SerdeDecoder::new(),
2253 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
2254 );
2255
2256 let verdict = judge.parse(&text_response(None, reply))?;
2257
2258 assert!(verdict.is_valid());
2259 assert!((0.0..=1.0).contains(&verdict.p_solve));
2260 Ok(())
2261 }
2262
2263 #[test]
2264 fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
2265 let contract =
2266 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
2267 let prompt = contract.system_prompt();
2268 let schema_name = contract
2269 .response_format()
2270 .pointer("/json_schema/name")
2271 .and_then(Value::as_str)
2272 .ok_or_else(|| LibsyError::AlgorithmError {
2273 message: "packaged response schema has no name".to_string(),
2274 })?;
2275 assert_eq!(schema_name, "CapabilityClassifierDecision");
2276 assert!(prompt.contains("SUP-1 [supported]"));
2277 assert!(prompt.contains("SUP-5 [supported]"));
2278 assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
2279 assert!(!prompt.contains("\"type\": \"object\""));
2280 assert!(!prompt.contains("\"json_schema\""));
2281 assert!(!prompt.contains(schema_name));
2282 let rule_values = contract
2283 .response_format()
2284 .pointer("/json_schema/schema/properties/primary_rule/enum")
2285 .and_then(Value::as_array)
2286 .ok_or_else(|| LibsyError::AlgorithmError {
2287 message: "rendered response schema has no primary rule enum".to_string(),
2288 })?;
2289 assert!(
2290 rule_values
2291 .iter()
2292 .any(|value| value.as_str() == Some("SUP-1"))
2293 );
2294 assert!(
2295 rule_values
2296 .iter()
2297 .any(|value| value.as_str() == Some("none"))
2298 );
2299 Ok(())
2300 }
2301}