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