1use std::collections::{BTreeMap, BTreeSet, HashSet};
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::{Deserialize, Deserializer};
11use serde_json::Value;
12use switchyard_protocol::{ContentBlock, Message, ModelId, Role};
13
14use super::escalation;
15use super::fall_through::{DefaultTarget, FallThrough};
16use super::util::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS;
17use super::util::affinity::{AffinityRouter, ClassifyTrigger};
18use super::util::classifier_contract::{
19 ClassifierContract, ClassifierContractConfig, ClassifierResponseFormat,
20};
21use super::util::escalation::EscalationJudgeConfig;
22use super::util::llm_judge::{
23 ClassifierInput, JsonSchemaDecoder, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig,
24 SerdeDecoder, StructuredJudge,
25};
26use super::util::target_selector::TargetSelectorPolicy;
27use crate::core::algorithm::{self, Algorithm, Driver};
28use crate::core::classifier::{Classification, Classifier, Score};
29use crate::core::state::State;
30use crate::{LibsyError, Result};
31use switchyard_protocol::{Request, Response};
32
33const PROMPT_TEMPLATE: &str = include_str!("../prompts/capability-classifier/prompt.md");
34const SCHEMA_TEMPLATE: &str = include_str!("../prompts/capability-classifier/schema.json");
35const ALGORITHM_NAME: &str = "llm_task_classifier";
37
38const TRAILING_ROUTING_INSTRUCTION: &str =
47 "Route the conversation above. Output ONLY the routing JSON object, nothing else.";
48
49#[derive(Deserialize)]
50#[serde(deny_unknown_fields)]
51struct TaskClassifierVerdict {
52 crux: String,
53 primary_rule: String,
54 capability_boundary: String,
55 p_solve: f64,
56}
57
58impl TaskClassifierVerdict {
59 fn is_valid(&self) -> bool {
61 (0.0..=1.0).contains(&self.p_solve)
62 && !self.crux.trim().is_empty()
63 && matches!(
64 (
65 self.primary_rule.as_str(),
66 self.capability_boundary.as_str()
67 ),
68 ("SUP-1" | "SUP-2" | "SUP-3" | "SUP-4" | "SUP-5", "supported")
69 | ("UNC-1" | "UNC-2", "uncertain")
70 | ("LIM-1" | "LIM-2", "unsupported")
71 | ("none", "unmatched")
72 )
73 }
74
75 fn boundary_steps(&self) -> Option<u8> {
77 match self.capability_boundary.as_str() {
78 "supported" => Some(0),
79 "uncertain" | "unmatched" => Some(1),
80 "unsupported" => Some(2),
81 _ => None,
82 }
83 }
84}
85
86fn trim_messages(messages: &[Message], recent_turn_window: usize) -> Vec<Message> {
96 let is_instruction = |message: &Message| matches!(message.role, Role::System | Role::Developer);
97 let mut kept: Vec<&Message> = messages.iter().filter(|m| is_instruction(m)).collect();
98 let Some(task) = messages.iter().position(|m| m.role == Role::User) else {
99 return kept.into_iter().cloned().collect();
100 };
101 kept.push(&messages[task]);
102
103 let tail: Vec<&Message> = messages[task + 1..]
104 .iter()
105 .filter(|m| !is_instruction(m))
106 .collect();
107 kept.extend(&tail[window_start(&tail, recent_turn_window)..]);
108 kept.into_iter().cloned().collect()
109}
110
111fn window_start(tail: &[&Message], recent_turn_window: usize) -> usize {
124 let counted = tail.len().saturating_sub(recent_turn_window);
125 if counted == tail.len() {
127 return counted;
128 }
129 let mut unpaired: HashSet<&str> = HashSet::new();
130 for (start, message) in tail.iter().enumerate().rev() {
131 for block in message.content.iter().rev() {
134 match block {
135 ContentBlock::ToolResult(result) => {
136 unpaired.insert(result.tool_call_id.as_str());
137 }
138 ContentBlock::ToolCall(call) => {
139 unpaired.remove(call.id.as_str());
140 }
141 _ => {}
142 }
143 }
144 if start <= counted && unpaired.is_empty() {
145 return start;
146 }
147 }
148 counted
149}
150
151fn task_messages(messages: &[Message]) -> Vec<Message> {
153 let mut user_messages = messages.iter().filter(|message| message.role == Role::User);
154 let Some(opening_task) = user_messages.next() else {
155 return Vec::new();
156 };
157 match user_messages.next_back() {
158 Some(latest_follow_up) => vec![opening_task.clone(), latest_follow_up.clone()],
159 None => vec![opening_task.clone()],
160 }
161}
162
163struct TaskInput {
165 recent_turn_window: Option<usize>,
166}
167
168impl ClassifierInput for TaskInput {
169 fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
170 let mut messages = match self.recent_turn_window {
173 Some(window) => trim_messages(&request.llm_request.messages, window),
174 None => task_messages(&request.llm_request.messages),
175 };
176 for message in &mut messages {
181 message
182 .content
183 .retain(|block| !matches!(block, ContentBlock::Reasoning { .. }));
184 }
185 messages.retain(|message| !message.content.is_empty());
186 if self.recent_turn_window.is_some() {
190 messages.push(Message::text(
191 Role::User,
192 TRAILING_ROUTING_INSTRUCTION.to_string(),
193 ));
194 }
195 messages
196 }
197}
198
199type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;
200
201struct TaskClassifierPolicy {
202 efficient_target: ModelId,
203 capable_target: ModelId,
204 base_threshold: f64,
205 threshold_step: f64,
206}
207
208impl TaskClassifierPolicy {
209 fn new(
210 efficient_target: impl Into<ModelId>,
211 capable_target: impl Into<ModelId>,
212 config: &TaskClassifierConfig,
213 ) -> Self {
214 Self {
215 efficient_target: efficient_target.into(),
216 capable_target: capable_target.into(),
217 base_threshold: config.base_threshold,
218 threshold_step: config.threshold_step,
219 }
220 }
221
222 fn threshold(&self, verdict: &TaskClassifierVerdict) -> Option<f64> {
224 Some(self.base_threshold + f64::from(verdict.boundary_steps()?) * self.threshold_step)
225 }
226}
227
228impl JudgePolicy for TaskClassifierPolicy {
229 type Verdict = TaskClassifierVerdict;
230
231 fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification {
232 let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
235 return Classification::Ambiguous(vec![]);
236 };
237 let Some(threshold) = self.threshold(verdict) else {
240 return Classification::Ambiguous(vec![]);
241 };
242 let target = if verdict.p_solve >= threshold
243 || (threshold - verdict.p_solve).abs() <= f64::EPSILON
244 {
245 &self.efficient_target
246 } else {
247 &self.capable_target
248 };
249 Classification::Scores(vec![Score {
250 target: target.clone(),
251 confidence: 1.0,
252 }])
253 }
254}
255
256fn capability_evidence(
258 policy: &TaskClassifierPolicy,
259 verdict: Option<&TaskClassifierVerdict>,
260) -> Option<Value> {
261 let verdict = verdict?;
262 let Some(threshold) = verdict
263 .is_valid()
264 .then(|| policy.threshold(verdict))
265 .flatten()
266 else {
267 return Some(serde_json::json!({
268 "source": "fail_open",
269 "reason_code": "invalid_verdict",
270 }));
271 };
272 Some(serde_json::json!({
273 "source": "llm-classifier",
274 "score": verdict.p_solve,
275 "threshold": threshold,
276 }))
277}
278
279#[derive(Clone, Debug)]
280pub struct TaskClassifierConfig {
282 pub base_threshold: f64,
284 pub threshold_step: f64,
289 pub classify_trigger: ClassifyTrigger,
291 pub message_hash_fallback: bool,
293 pub recent_turn_window: Option<usize>,
300 pub contract: ClassifierContractConfig,
302 pub max_output_tokens: u64,
304}
305
306#[derive(Deserialize)]
308#[serde(deny_unknown_fields)]
309struct TaskClassifierConfigWire {
310 base_threshold: f64,
311 #[serde(default)]
312 threshold_step: f64,
313 #[serde(default)]
314 classify_trigger: ClassifyTrigger,
315 #[serde(default)]
316 message_hash_fallback: bool,
317 #[serde(default)]
318 recent_turn_window: Option<usize>,
319 #[serde(default)]
320 prompt: Option<String>,
321 #[serde(default)]
322 response_format_type: ClassifierResponseFormat,
323 #[serde(default = "default_judge_max_output_tokens")]
324 max_output_tokens: u64,
325}
326
327impl<'de> Deserialize<'de> for TaskClassifierConfig {
328 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
329 where
330 D: Deserializer<'de>,
331 {
332 let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
333 let mut contract = ClassifierContractConfig::default();
334 if let Some(prompt) = wire.prompt {
335 contract = contract.with_prompt(prompt);
336 }
337 contract = contract.with_response_format_type(wire.response_format_type);
338 Ok(Self {
339 base_threshold: wire.base_threshold,
340 threshold_step: wire.threshold_step,
341 classify_trigger: wire.classify_trigger,
342 message_hash_fallback: wire.message_hash_fallback,
343 recent_turn_window: wire.recent_turn_window,
344 contract,
345 max_output_tokens: wire.max_output_tokens,
346 })
347 }
348}
349
350const fn default_judge_max_output_tokens() -> u64 {
351 DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
352}
353
354impl Default for TaskClassifierConfig {
355 fn default() -> Self {
356 Self {
357 base_threshold: 0.0,
358 threshold_step: 0.0,
359 classify_trigger: ClassifyTrigger::default(),
360 message_hash_fallback: false,
361 recent_turn_window: None,
362 contract: ClassifierContractConfig::default(),
363 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
364 }
365 }
366}
367
368impl TaskClassifierConfig {
369 fn validate(&self) -> Result<()> {
371 if !(0.0..=1.0).contains(&self.base_threshold) {
372 return Err(LibsyError::AlgorithmError {
373 message: format!(
374 "base_threshold must be between 0 and 1, got {}",
375 self.base_threshold
376 ),
377 });
378 }
379 if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
380 return Err(LibsyError::AlgorithmError {
381 message: format!(
382 "threshold_step must be finite and greater than or equal to 0, got {}",
383 self.threshold_step
384 ),
385 });
386 }
387 let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
388 if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
389 return Err(LibsyError::AlgorithmError {
390 message: format!(
391 "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
392 ),
393 });
394 }
395 if self.max_output_tokens == 0 {
396 return Err(LibsyError::AlgorithmError {
397 message: "max_output_tokens must be at least 1".to_string(),
398 });
399 }
400 if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
401 return Err(LibsyError::AlgorithmError {
402 message: "message_hash_fallback requires classify_trigger = new_session"
403 .to_string(),
404 });
405 }
406 Ok(())
407 }
408}
409
410#[derive(Clone, Debug)]
412pub enum CustomClassifierPolicy {
413 TargetSelector {
415 selector: String,
417 },
418}
419
420impl CustomClassifierPolicy {
421 pub fn target_selector(selector: impl Into<String>) -> Self {
423 Self::TargetSelector {
424 selector: selector.into(),
425 }
426 }
427}
428
429#[derive(Clone, Debug)]
431pub struct CustomClassifierConfig {
432 pub prompt: String,
434 pub response_schema: Value,
436 pub policy: CustomClassifierPolicy,
438 pub classify_trigger: ClassifyTrigger,
440 pub message_hash_fallback: bool,
442 pub recent_turn_window: Option<usize>,
444 pub max_output_tokens: u64,
446}
447
448impl CustomClassifierConfig {
449 pub fn new(
451 prompt: impl Into<String>,
452 response_schema: Value,
453 policy: CustomClassifierPolicy,
454 ) -> Self {
455 Self {
456 prompt: prompt.into(),
457 response_schema,
458 policy,
459 classify_trigger: ClassifyTrigger::default(),
460 message_hash_fallback: false,
461 recent_turn_window: None,
462 max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
463 }
464 }
465
466 fn validate(&self) -> Result<()> {
467 if self.max_output_tokens == 0 {
468 return Err(LibsyError::AlgorithmError {
469 message: "max_output_tokens must be at least 1".to_string(),
470 });
471 }
472 if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
473 return Err(LibsyError::AlgorithmError {
474 message: "message_hash_fallback requires classify_trigger = new_session"
475 .to_string(),
476 });
477 }
478 Ok(())
479 }
480}
481
482enum CustomPolicyRuntime {
483 TargetSelector(TargetSelectorPolicy),
484}
485
486impl JudgePolicy for CustomPolicyRuntime {
487 type Verdict = Value;
488
489 fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification {
490 match self {
491 Self::TargetSelector(policy) => policy.to_classification(verdict),
492 }
493 }
494}
495
496struct TaskClassifier {
497 classifier: JudgeClassifier<CapabilityJudge, TaskClassifierPolicy>,
498 capable_target: ModelId,
499}
500
501fn affinity_router(
503 trigger: ClassifyTrigger,
504 message_hash_fallback: bool,
505) -> Option<Arc<AffinityRouter>> {
506 let router = match trigger {
507 ClassifyTrigger::EveryRequest => return None,
508 ClassifyTrigger::NewSession => AffinityRouter::new(),
509 ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
510 };
511 let router = if message_hash_fallback {
512 router.with_message_hash_fallback()
513 } else {
514 router
515 };
516 Some(Arc::new(router))
517}
518
519pub struct LlmTaskClassifier {
521 route: FallThrough<State>,
522 inner: Arc<dyn Classifier<State>>,
524}
525
526struct ClassifierRouteConfig {
527 default_target: ModelId,
528 classify_trigger: ClassifyTrigger,
529 message_hash_fallback: bool,
530}
531
532#[derive(Clone)]
534#[non_exhaustive]
535pub enum LlmClassifierConfig {
536 Capability {
538 judge_target: ModelId,
540 efficient_target: ModelId,
542 capable_target: ModelId,
544 config: TaskClassifierConfig,
546 },
547 Escalation {
549 judge_target: ModelId,
551 efficient_target: ModelId,
553 capable_target: ModelId,
555 contract: ClassifierContractConfig,
557 config: EscalationJudgeConfig,
559 max_output_tokens: u64,
561 },
562 Custom {
564 judge_target: ModelId,
566 targets: Vec<(String, ModelId)>,
568 default_target: String,
570 config: CustomClassifierConfig,
572 },
573}
574
575impl LlmTaskClassifier {
576 pub fn new(config: LlmClassifierConfig) -> Result<Self> {
583 match config {
584 LlmClassifierConfig::Capability {
585 judge_target,
586 efficient_target,
587 capable_target,
588 config,
589 } => Self::build_capability(judge_target, efficient_target, capable_target, config),
590 LlmClassifierConfig::Escalation {
591 judge_target,
592 efficient_target,
593 capable_target,
594 contract,
595 config,
596 max_output_tokens,
597 } => Self::build_escalation(
598 judge_target,
599 efficient_target,
600 capable_target,
601 contract,
602 config,
603 max_output_tokens,
604 ),
605 LlmClassifierConfig::Custom {
606 judge_target,
607 targets,
608 default_target,
609 config,
610 } => Self::build_custom(judge_target, targets, default_target, config),
611 }
612 }
613
614 fn build_capability(
615 judge_target: ModelId,
616 efficient_target: ModelId,
617 capable_target: ModelId,
618 config: TaskClassifierConfig,
619 ) -> Result<Self> {
620 config.validate()?;
621 let contract = Self::load_capability_contract(&config.contract)?;
622 let targets = vec![efficient_target.clone(), capable_target.clone()];
623 let classify_trigger = config.classify_trigger;
624 let message_hash_fallback = config.message_hash_fallback;
625 let classifier = Arc::new(TaskClassifier {
626 classifier: JudgeClassifier::new(
627 StructuredJudge::new(
628 TaskInput {
629 recent_turn_window: config.recent_turn_window,
630 },
631 contract,
632 SerdeDecoder::new(),
633 JudgeRuntimeConfig::new(config.max_output_tokens)?,
634 ),
635 judge_target.clone(),
636 TaskClassifierPolicy::new(
637 efficient_target.clone(),
638 capable_target.clone(),
639 &config,
640 ),
641 )
642 .with_evidence(capability_evidence),
643 capable_target: capable_target.clone(),
644 });
645 let inner: Arc<dyn Classifier<State>> = classifier.clone();
646 Self::from_classifier(
647 targets,
648 inner,
649 ClassifierRouteConfig {
650 default_target: classifier.capable_target.clone(),
651 classify_trigger,
652 message_hash_fallback,
653 },
654 )
655 }
656
657 fn build_custom(
658 judge_target: ModelId,
659 targets: Vec<(String, ModelId)>,
660 default_target: String,
661 config: CustomClassifierConfig,
662 ) -> Result<Self> {
663 config.validate()?;
664 if targets.len() < 2 {
665 return Err(LibsyError::AlgorithmError {
666 message: "custom classifier requires at least two targets".to_string(),
667 });
668 }
669
670 let mut labels = BTreeSet::new();
671 let mut resolved_names = BTreeSet::new();
672 let mut target_map = BTreeMap::new();
673 let mut resolved_targets = Vec::with_capacity(targets.len());
674 for (label, target) in targets {
675 if label.trim().is_empty() || label.trim() != label {
676 return Err(LibsyError::AlgorithmError {
677 message: "custom classifier target labels must be non-empty and have no surrounding whitespace"
678 .to_string(),
679 });
680 }
681 if !labels.insert(label.clone()) {
682 return Err(LibsyError::AlgorithmError {
683 message: format!("custom classifier target label {label:?} is duplicated"),
684 });
685 }
686 if !resolved_names.insert(target.clone()) {
687 return Err(LibsyError::AlgorithmError {
688 message: format!("custom classifier resolved target {target:?} is duplicated"),
689 });
690 }
691 target_map.insert(label, target.clone());
692 resolved_targets.push(target);
693 }
694 let default_name =
695 target_map
696 .get(&default_target)
697 .cloned()
698 .ok_or_else(|| LibsyError::AlgorithmError {
699 message: format!(
700 "default_target {default_target:?} must be one of the configured targets"
701 ),
702 })?;
703
704 let CustomClassifierConfig {
705 prompt,
706 response_schema,
707 policy,
708 classify_trigger,
709 message_hash_fallback,
710 recent_turn_window,
711 max_output_tokens,
712 } = config;
713 let contract = ClassifierContract::from_inner_schema(&prompt, response_schema)?;
714 let policy = match policy {
715 CustomClassifierPolicy::TargetSelector { selector } => {
716 CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(
717 selector, target_map,
718 )?)
719 }
720 };
721 let classifier: Arc<dyn Classifier<State>> = Arc::new(JudgeClassifier::new(
722 StructuredJudge::new(
723 TaskInput { recent_turn_window },
724 contract,
725 JsonSchemaDecoder::new(),
726 JudgeRuntimeConfig::new(max_output_tokens)?,
727 ),
728 judge_target,
729 policy,
730 ));
731
732 Self::from_classifier(
733 resolved_targets,
734 classifier,
735 ClassifierRouteConfig {
736 default_target: default_name,
737 classify_trigger,
738 message_hash_fallback,
739 },
740 )
741 }
742
743 fn build_escalation(
744 judge_target: ModelId,
745 efficient_target: ModelId,
746 capable_target: ModelId,
747 contract_config: ClassifierContractConfig,
748 config: EscalationJudgeConfig,
749 max_output_tokens: u64,
750 ) -> Result<Self> {
751 let inner = escalation::build_classifier(
752 judge_target,
753 &efficient_target,
754 &capable_target,
755 contract_config,
756 config,
757 max_output_tokens,
758 )?;
759 let targets = vec![capable_target, efficient_target];
760 Ok(Self {
761 route: FallThrough::<State>::new_with_state(targets)
762 .with_name(ALGORITHM_NAME)
763 .with_classifier(Arc::clone(&inner)),
764 inner,
765 })
766 }
767
768 fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
770 ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
771 }
772
773 fn from_classifier(
775 targets: Vec<ModelId>,
776 inner: Arc<dyn Classifier<State>>,
777 config: ClassifierRouteConfig,
778 ) -> Result<Self> {
779 algorithm::ensure_model_is_target(&targets, &config.default_target)?;
780 if config.message_hash_fallback && config.classify_trigger != ClassifyTrigger::NewSession {
781 return Err(LibsyError::AlgorithmError {
782 message: "message_hash_fallback requires classify_trigger = new_session"
783 .to_string(),
784 });
785 }
786 let mut route = FallThrough::<State>::new_with_state(targets).with_name(ALGORITHM_NAME);
788 if let Some(affinity) =
789 affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
790 {
791 route = route
793 .with_processor(affinity.clone())
794 .with_classifier(affinity.clone());
795 }
796 let fallback = DefaultTarget::new(config.default_target);
797 Ok(Self {
798 route: route
799 .with_classifier(inner.clone())
800 .with_classifier(Arc::new(fallback)),
801 inner,
802 })
803 }
804}
805
806#[async_trait]
807impl Classifier<State> for TaskClassifier {
808 async fn score(
809 &self,
810 state: &mut State,
811 request: &mut Request,
812 driver: Option<&Driver>,
813 ) -> Result<(Classification, Option<Response>)> {
814 self.classifier.score(state, request, driver).await
815 }
816}
817
818#[async_trait]
819impl Classifier<State> for LlmTaskClassifier {
820 async fn score(
821 &self,
822 state: &mut State,
823 request: &mut Request,
824 driver: Option<&Driver>,
825 ) -> Result<(Classification, Option<Response>)> {
826 self.inner.score(state, request, driver).await
827 }
828}
829
830#[async_trait]
831impl Algorithm for LlmTaskClassifier {
832 fn name(&self) -> &str {
833 "llm_task_classifier"
834 }
835
836 async fn route(
837 self: Arc<Self>,
838 driver: Driver,
839 request: Request,
840 ) -> Result<crate::RoutingOutcome> {
841 self.route.execute(driver, request).await
842 }
843}
844
845#[cfg(test)]
846mod tests {
847 use std::sync::Arc;
848
849 use parking_lot::Mutex;
850 use serde_json::Value;
851
852 use super::*;
853 use switchyard_protocol::{
854 ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ToolCall, ToolResult,
855 completion_text, text_request, text_response,
856 };
857
858 use crate::algorithms::util::llm_judge::Judge;
859 use crate::core::testing::{Serve, test_drive};
860 use switchyard_protocol::{LlmResponse, Response};
861
862 const TEST_THRESHOLD: f64 = 0.5;
863
864 fn test_config(base_threshold: f64) -> TaskClassifierConfig {
865 TaskClassifierConfig {
866 base_threshold,
867 ..TaskClassifierConfig::default()
868 }
869 }
870
871 fn policy() -> TaskClassifierPolicy {
872 TaskClassifierPolicy::new("efficient", "capable", &test_config(TEST_THRESHOLD))
873 }
874
875 fn verdict(
876 p_solve: f64,
877 capability_boundary: &str,
878 primary_rule: &str,
879 ) -> TaskClassifierVerdict {
880 TaskClassifierVerdict {
881 crux: "test crux".to_string(),
882 primary_rule: primary_rule.to_string(),
883 capability_boundary: capability_boundary.to_string(),
884 p_solve,
885 }
886 }
887
888 fn selected(
889 policy: &TaskClassifierPolicy,
890 verdict: Option<&TaskClassifierVerdict>,
891 ) -> Result<ModelId> {
892 policy
893 .to_classification(verdict)
894 .argmax(false)?
895 .map(|score| score.target)
896 .ok_or_else(|| LibsyError::AlgorithmError {
897 message: "policy abstained".to_string(),
898 })
899 }
900
901 #[derive(Default)]
904 struct Recorder {
905 calls: Mutex<Vec<String>>,
906 call_roles: Mutex<Vec<(String, bool)>>,
907 judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
908 judge_system_prompts: Mutex<Vec<String>>,
909 }
910
911 impl Recorder {
912 fn calls(&self) -> Vec<String> {
913 self.calls.lock().clone()
914 }
915
916 fn call_roles(&self) -> Vec<(String, bool)> {
917 self.call_roles.lock().clone()
918 }
919
920 fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
921 self.judge_max_output_tokens.lock().clone()
922 }
923
924 fn judge_system_prompts(&self) -> Vec<String> {
925 self.judge_system_prompts.lock().clone()
926 }
927
928 fn serve(self: &Arc<Self>) -> impl Serve {
929 let recorder = Arc::clone(self);
930 move |model: ModelId, request: Request| {
931 let recorder = Arc::clone(&recorder);
932 async move {
933 let model = model.to_string();
934 recorder.calls.lock().push(model.clone());
935 recorder
936 .call_roles
937 .lock()
938 .push((model.clone(), model != "judge"));
939 let completion = if model == "judge" {
940 recorder
941 .judge_max_output_tokens
942 .lock()
943 .push(request.llm_request.output.max_output_tokens);
944 recorder.judge_system_prompts.lock().extend(
945 request
946 .llm_request
947 .instructions
948 .first()
949 .and_then(|instruction| {
950 instruction.content.iter().find_map(|b| {
951 if let ContentBlock::Text { text } = b {
952 Some(text.clone())
953 } else {
954 None
955 }
956 })
957 }),
958 );
959 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
960 } else {
961 format!("answer from {model}")
962 };
963 Ok(Response {
964 llm_response: LlmResponse::Agg(text_response(None, completion)),
965 metadata: request.metadata,
966 })
967 }
968 }
969 }
970 }
971
972 fn unreachable_judge() -> impl Serve {
974 |model: ModelId, request: Request| async move {
975 let model = model.to_string();
976 if model == "judge" {
977 return Err(LlmClientError::Timeout {
978 source: Box::new(std::io::Error::other("judge unreachable")),
979 });
980 }
981 Ok(Response {
982 llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
983 metadata: request.metadata,
984 })
985 }
986 }
987
988 fn router() -> Result<Arc<LlmTaskClassifier>> {
989 Ok(Arc::new(LlmTaskClassifier::new(
990 LlmClassifierConfig::Capability {
991 judge_target: ModelId::from("judge"),
992 efficient_target: ModelId::from("efficient"),
993 capable_target: ModelId::from("capable"),
994 config: test_config(TEST_THRESHOLD),
995 },
996 )?))
997 }
998
999 fn classify_request() -> Request {
1000 Request {
1001 llm_request: text_request(Some("auto".to_string()), "classify this task"),
1002 raw_request: None,
1003 metadata: None,
1004 }
1005 }
1006
1007 fn classify_session_request() -> Request {
1008 Request {
1009 metadata: Some(Metadata {
1010 session_id: Some("session-1".to_string()),
1011 ..Metadata::default()
1012 }),
1013 ..classify_request()
1014 }
1015 }
1016
1017 fn classify_follow_up_request() -> Request {
1018 let mut request = classify_request();
1019 request
1020 .llm_request
1021 .messages
1022 .push(Message::text(Role::Assistant, "I will add the test."));
1023 request.llm_request.messages.push(Message::text(
1024 Role::User,
1025 "Now run the test suite and report the result.",
1026 ));
1027 request
1028 }
1029
1030 #[tokio::test]
1031 async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
1032 let router = router()?;
1033
1034 let (selected_model, response) =
1035 test_drive(router, classify_request(), unreachable_judge()).await?;
1036
1037 assert_eq!(selected_model, "capable");
1038 assert_eq!(
1039 response.llm_response.as_agg().map(completion_text),
1040 Some("answer from capable".to_string())
1041 );
1042 Ok(())
1043 }
1044
1045 #[tokio::test]
1046 async fn classifier_judges_each_request_without_affinity() -> Result<()> {
1047 let recorder = Arc::new(Recorder::default());
1048 let router = router()?;
1049 let request = classify_request;
1050
1051 test_drive(router.clone(), request(), recorder.serve()).await?;
1052 test_drive(router.clone(), request(), recorder.serve()).await?;
1053
1054 assert_eq!(
1055 recorder.calls(),
1056 vec!["judge", "efficient", "judge", "efficient"]
1057 );
1058 assert_eq!(
1059 recorder.call_roles(),
1060 vec![
1061 ("judge".to_string(), false),
1062 ("efficient".to_string(), true),
1063 ("judge".to_string(), false),
1064 ("efficient".to_string(), true),
1065 ]
1066 );
1067 Ok(())
1068 }
1069
1070 #[tokio::test]
1071 async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1072 let recorder = Arc::new(Recorder::default());
1073 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1074 judge_target: ModelId::from("judge"),
1075 efficient_target: ModelId::from("efficient"),
1076 capable_target: ModelId::from("capable"),
1077 config: TaskClassifierConfig {
1078 max_output_tokens: 512,
1079 ..test_config(TEST_THRESHOLD)
1080 },
1081 })?);
1082
1083 test_drive(router, classify_request(), recorder.serve()).await?;
1084
1085 assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1086 Ok(())
1087 }
1088
1089 #[tokio::test]
1090 async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1091 let recorder = Arc::new(Recorder::default());
1092 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1093 judge_target: ModelId::from("judge"),
1094 efficient_target: ModelId::from("efficient"),
1095 capable_target: ModelId::from("capable"),
1096 config: TaskClassifierConfig {
1097 contract: ClassifierContractConfig::default()
1098 .with_prompt("Custom capability rubric."),
1099 ..test_config(TEST_THRESHOLD)
1100 },
1101 })?);
1102
1103 test_drive(router, classify_request(), recorder.serve()).await?;
1104
1105 let prompts = recorder.judge_system_prompts();
1106 assert_eq!(prompts.len(), 1);
1107 assert_eq!(prompts[0], "Custom capability rubric.");
1108 Ok(())
1109 }
1110
1111 #[tokio::test]
1112 async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1113 let recorder = Arc::new(Recorder::default());
1114 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1115 judge_target: ModelId::from("judge"),
1116 efficient_target: ModelId::from("efficient"),
1117 capable_target: ModelId::from("capable"),
1118 config: TaskClassifierConfig {
1119 classify_trigger: ClassifyTrigger::NewSession,
1120 ..test_config(TEST_THRESHOLD)
1121 },
1122 })?);
1123
1124 let session_request = classify_session_request;
1125 test_drive(router.clone(), session_request(), recorder.serve()).await?;
1126 test_drive(router.clone(), session_request(), recorder.serve()).await?;
1127
1128 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1129 Ok(())
1130 }
1131
1132 #[tokio::test]
1133 async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1134 let recorder = Arc::new(Recorder::default());
1135 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1136 judge_target: ModelId::from("judge"),
1137 efficient_target: ModelId::from("efficient"),
1138 capable_target: ModelId::from("capable"),
1139 config: TaskClassifierConfig {
1140 classify_trigger: ClassifyTrigger::NewSession,
1141 message_hash_fallback: true,
1142 recent_turn_window: None,
1143 ..test_config(TEST_THRESHOLD)
1144 },
1145 })?);
1146
1147 test_drive(router.clone(), classify_request(), recorder.serve()).await?;
1148 test_drive(
1149 router.clone(),
1150 classify_follow_up_request(),
1151 recorder.serve(),
1152 )
1153 .await?;
1154
1155 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1156 Ok(())
1157 }
1158
1159 #[test]
1160 fn the_threshold_boundary_is_inclusive() -> Result<()> {
1161 let policy = policy();
1162 let at_threshold = verdict(0.5, "supported", "SUP-1");
1163 let below_threshold = verdict(0.49, "supported", "SUP-1");
1164 assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1165 assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1166 Ok(())
1167 }
1168
1169 #[test]
1170 fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1171 let borderline = verdict(0.5, "supported", "SUP-1");
1172 let strict = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.9));
1173 let lenient = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.1));
1174 assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1175 assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1176 Ok(())
1177 }
1178
1179 #[test]
1180 fn classifier_config_rejects_unknown_fields() {
1181 let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1182 "base_threshold": 0.5,
1183 "classifier_magic": true,
1184 }))
1185 .expect_err("unknown classifier fields must be rejected");
1186
1187 assert!(
1188 error
1189 .to_string()
1190 .contains("unknown field `classifier_magic`"),
1191 "{error}"
1192 );
1193 }
1194
1195 #[test]
1196 fn invalid_classifier_config_is_rejected() -> Result<()> {
1197 for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1198 assert!(
1199 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1200 judge_target: ModelId::from("judge"),
1201 efficient_target: ModelId::from("e"),
1202 capable_target: ModelId::from("c"),
1203 config: test_config(bad),
1204 })
1205 .is_err(),
1206 "base threshold {bad} should be rejected"
1207 );
1208 }
1209 for config in [
1210 TaskClassifierConfig {
1211 base_threshold: 0.5,
1212 threshold_step: -0.1,
1213 ..TaskClassifierConfig::default()
1214 },
1215 TaskClassifierConfig {
1216 base_threshold: 0.8,
1217 threshold_step: 0.11,
1218 ..TaskClassifierConfig::default()
1219 },
1220 TaskClassifierConfig {
1221 base_threshold: 0.5,
1222 message_hash_fallback: true,
1223 ..TaskClassifierConfig::default()
1224 },
1225 TaskClassifierConfig {
1226 base_threshold: 0.5,
1227 max_output_tokens: 0,
1228 ..TaskClassifierConfig::default()
1229 },
1230 ] {
1231 assert!(
1232 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1233 judge_target: ModelId::from("judge"),
1234 efficient_target: ModelId::from("e"),
1235 capable_target: ModelId::from("c"),
1236 config,
1237 })
1238 .is_err()
1239 );
1240 }
1241 for base_threshold in [0.0, 1.0] {
1242 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1243 judge_target: ModelId::from("judge"),
1244 efficient_target: ModelId::from("e"),
1245 capable_target: ModelId::from("c"),
1246 config: test_config(base_threshold),
1247 })?;
1248 }
1249 Ok(())
1250 }
1251
1252 #[test]
1253 fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1254 let policy = policy();
1255 let inconsistent_rule = TaskClassifierVerdict {
1256 capability_boundary: "uncertain".to_string(),
1257 ..verdict(1.0, "supported", "SUP-1")
1258 };
1259 let empty_crux = TaskClassifierVerdict {
1260 crux: " ".to_string(),
1261 ..verdict(1.0, "supported", "SUP-1")
1262 };
1263 let unusable = [
1264 Some(verdict(1.1, "supported", "SUP-1")),
1265 Some(inconsistent_rule),
1266 Some(empty_crux),
1267 None,
1268 ];
1269 for verdict in unusable {
1270 let classification = policy.to_classification(verdict.as_ref());
1271 assert!(matches!(classification, Classification::Ambiguous(_)));
1272 assert!(classification.argmax(false)?.is_none());
1273 assert!(classification.argmax(true)?.is_none());
1274 }
1275 Ok(())
1276 }
1277
1278 #[test]
1279 fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
1280 let policy = TaskClassifierPolicy::new(
1281 "efficient",
1282 "capable",
1283 &TaskClassifierConfig {
1284 threshold_step: 0.1,
1285 ..test_config(0.4)
1286 },
1287 );
1288
1289 assert_eq!(
1290 selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
1291 "efficient"
1292 );
1293 assert_eq!(
1294 selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
1295 "capable"
1296 );
1297 assert_eq!(
1298 selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
1299 "efficient"
1300 );
1301 assert_eq!(
1302 selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
1303 "efficient"
1304 );
1305 assert_eq!(
1306 selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
1307 "capable"
1308 );
1309 assert_eq!(
1310 selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
1311 "efficient"
1312 );
1313 Ok(())
1314 }
1315
1316 fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
1319 Ok(StructuredJudge::new(
1320 TaskInput { recent_turn_window },
1321 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
1322 SerdeDecoder::new(),
1323 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1324 ))
1325 }
1326
1327 fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
1328 let judge = capability_judge(Some(recent_turn_window))?;
1329 let request = Request {
1330 llm_request: LlmRequest {
1331 messages: vec![
1332 Message::text(Role::System, "client instructions"),
1333 Message::text(Role::User, "initial task"),
1334 Message::text(Role::Assistant, "old response"),
1335 Message::text(Role::User, "old follow-up"),
1336 Message::text(Role::Assistant, "recent 1"),
1337 Message::text(Role::User, "recent 2"),
1338 ],
1339 ..LlmRequest::default()
1340 },
1341 raw_request: None,
1342 metadata: None,
1343 };
1344 Ok(judge
1345 .build_request(&State::default(), &request)
1346 .llm_request
1347 .messages
1348 .iter()
1349 .filter_map(|message| message.text_content("\n"))
1350 .collect())
1351 }
1352
1353 #[test]
1354 fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
1355 let contents = judged_contents(2)?;
1357 assert!(contents.contains(&"client instructions".to_string()));
1358 assert!(contents.contains(&"initial task".to_string()));
1359 assert!(contents.contains(&"recent 1".to_string()));
1360 assert!(contents.contains(&"recent 2".to_string()));
1361 assert!(!contents.contains(&"old response".to_string()));
1362 Ok(())
1363 }
1364
1365 #[test]
1366 fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
1367 let contents = judged_contents(0)?;
1368 assert!(contents.contains(&"client instructions".to_string()));
1369 assert!(contents.contains(&"initial task".to_string()));
1370 assert!(!contents.contains(&"recent 2".to_string()));
1371 Ok(())
1372 }
1373
1374 fn tool_call(id: &str) -> Message {
1375 Message {
1376 role: Role::Assistant,
1377 content: vec![ContentBlock::ToolCall(ToolCall {
1378 id: id.to_string(),
1379 name: "search".to_string(),
1380 arguments: Value::Null,
1381 })],
1382 }
1383 }
1384
1385 fn tool_result(id: &str) -> Message {
1386 Message {
1387 role: Role::Tool,
1388 content: vec![ContentBlock::ToolResult(ToolResult {
1389 tool_call_id: id.to_string(),
1390 content: vec![ContentBlock::Text {
1391 text: "tool output".to_string(),
1392 }],
1393 is_error: None,
1394 })],
1395 }
1396 }
1397
1398 #[test]
1401 fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
1402 let messages = vec![
1403 Message::text(Role::System, "client instructions"),
1404 Message::text(Role::User, "initial task"),
1405 Message::text(Role::Assistant, "old response"),
1406 tool_call("call-1"),
1407 tool_result("call-1"),
1408 Message::text(Role::Assistant, "recent 1"),
1409 Message::text(Role::User, "recent 2"),
1410 Message::text(Role::Assistant, "recent 3"),
1411 Message::text(Role::User, "recent 4"),
1412 ];
1413
1414 let kept = trim_messages(&messages, 5);
1416
1417 assert_eq!(
1418 kept,
1419 vec![
1420 Message::text(Role::System, "client instructions"),
1421 Message::text(Role::User, "initial task"),
1422 tool_call("call-1"),
1423 tool_result("call-1"),
1424 Message::text(Role::Assistant, "recent 1"),
1425 Message::text(Role::User, "recent 2"),
1426 Message::text(Role::Assistant, "recent 3"),
1427 Message::text(Role::User, "recent 4"),
1428 ]
1429 );
1430 }
1431
1432 #[test]
1435 fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
1436 let messages = vec![
1437 Message::text(Role::System, "client instructions"),
1438 Message::text(Role::User, "initial task"),
1439 tool_call("x"),
1440 tool_result("x"),
1441 Message::text(Role::Assistant, "later"),
1442 tool_call("x"),
1443 tool_result("x"),
1444 ];
1445
1446 let kept = trim_messages(&messages, 4);
1448
1449 assert_eq!(
1450 kept,
1451 vec![
1452 Message::text(Role::System, "client instructions"),
1453 Message::text(Role::User, "initial task"),
1454 tool_call("x"),
1455 tool_result("x"),
1456 Message::text(Role::Assistant, "later"),
1457 tool_call("x"),
1458 tool_result("x"),
1459 ]
1460 );
1461 }
1462
1463 #[test]
1466 fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
1467 let messages = vec![
1468 Message::text(Role::System, "client instructions"),
1469 tool_call("orphan"),
1470 Message::text(Role::User, "initial task"),
1471 Message::text(Role::Assistant, "old response"),
1472 tool_result("orphan"),
1473 Message::text(Role::Assistant, "recent 1"),
1474 Message::text(Role::User, "recent 2"),
1475 ];
1476
1477 let kept = trim_messages(&messages, 3);
1478
1479 assert_eq!(
1480 kept,
1481 vec![
1482 Message::text(Role::System, "client instructions"),
1483 Message::text(Role::User, "initial task"),
1484 tool_result("orphan"),
1485 Message::text(Role::Assistant, "recent 1"),
1486 Message::text(Role::User, "recent 2"),
1487 ]
1488 );
1489 }
1490
1491 #[test]
1492 fn a_window_restates_the_routing_instruction_last() -> Result<()> {
1493 let contents = judged_contents(2)?;
1494
1495 assert_eq!(
1498 contents.last().map(String::as_str),
1499 Some(TRAILING_ROUTING_INSTRUCTION)
1500 );
1501 Ok(())
1502 }
1503
1504 #[test]
1509 fn a_window_drops_reasoning_but_keeps_visible_text_and_tool_pairs() {
1510 let messages = vec![
1511 Message::text(Role::System, "client instructions"),
1512 Message::text(Role::User, "initial task"),
1513 Message {
1514 role: Role::Assistant,
1515 content: vec![
1516 ContentBlock::Reasoning {
1517 text: "private chain of thought".to_string(),
1518 signature: None,
1519 details: Vec::new(),
1520 },
1521 ContentBlock::Text {
1522 text: "visible answer".to_string(),
1523 },
1524 ],
1525 },
1526 tool_call("call-1"),
1527 tool_result("call-1"),
1528 Message {
1529 role: Role::Assistant,
1530 content: vec![ContentBlock::Reasoning {
1531 text: "reasoning-only turn".to_string(),
1532 signature: None,
1533 details: Vec::new(),
1534 }],
1535 },
1536 Message::text(Role::User, "follow-up"),
1537 ];
1538 let request = Request {
1539 llm_request: LlmRequest {
1540 messages,
1541 ..LlmRequest::default()
1542 },
1543 raw_request: None,
1544 metadata: None,
1545 };
1546
1547 let built = TaskInput {
1548 recent_turn_window: Some(10),
1549 }
1550 .build_messages(&State::default(), &request);
1551
1552 assert!(
1553 !built
1554 .iter()
1555 .flat_map(|message| &message.content)
1556 .any(|block| matches!(block, ContentBlock::Reasoning { .. })),
1557 "{built:?}"
1558 );
1559 assert!(
1560 built
1561 .iter()
1562 .any(|message| message.text_content("\n").as_deref() == Some("visible answer"))
1563 );
1564 assert!(built.contains(&tool_call("call-1")));
1565 assert!(built.contains(&tool_result("call-1")));
1566 assert_eq!(built.len(), 7);
1569 assert!(built.iter().all(|message| !message.content.is_empty()));
1570 }
1571
1572 #[test]
1573 fn the_default_path_is_left_unchanged() -> Result<()> {
1574 let judge = capability_judge(None)?;
1577 let request = Request {
1578 llm_request: LlmRequest {
1579 messages: vec![Message::text(Role::User, "the task")],
1580 ..LlmRequest::default()
1581 },
1582 raw_request: None,
1583 metadata: None,
1584 };
1585
1586 let built = judge.build_request(&State::default(), &request);
1587
1588 assert_eq!(built.llm_request.messages.len(), 1);
1591 assert!(!built.llm_request.instructions.is_empty());
1592 assert!(
1593 !built
1594 .llm_request
1595 .messages
1596 .iter()
1597 .filter_map(|message| message.text_content("\n"))
1598 .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
1599 );
1600 Ok(())
1601 }
1602
1603 #[test]
1604 fn capability_judge_builds_a_structured_request() -> Result<()> {
1605 let judge = capability_judge(None)?;
1606 let request = Request {
1607 llm_request: LlmRequest {
1608 model: Some("inbound".to_string()),
1609 messages: vec![
1610 Message::text(Role::System, "client instructions"),
1611 Message::text(Role::Developer, "client developer instructions"),
1612 Message::text(Role::User, "initial task"),
1613 Message::text(Role::Assistant, "old response"),
1614 Message::text(Role::User, "old follow-up"),
1615 Message::text(Role::Assistant, "recent 1"),
1616 Message::text(Role::User, "recent 2"),
1617 Message::text(Role::Assistant, "recent 3"),
1618 Message::text(Role::User, "recent 4"),
1619 Message::text(Role::Assistant, "recent 5"),
1620 ],
1621 ..LlmRequest::default()
1622 },
1623 raw_request: None,
1624 metadata: None,
1625 };
1626 let judge_request = judge.build_request(&State::default(), &request);
1627
1628 assert_eq!(judge_request.llm_request.model, request.llm_request.model);
1629 assert_eq!(judge_request.llm_request.instructions.len(), 1);
1630 assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
1631 assert_eq!(
1632 judge_request.llm_request.instructions[0].content,
1633 InstructionBlock {
1634 role: Role::System,
1635 content: Message::text(Role::System, judge.contract().system_prompt()).content,
1636 }
1637 .content,
1638 );
1639 assert_eq!(judge_request.llm_request.messages.len(), 2);
1640 let contents = judge_request
1641 .llm_request
1642 .messages
1643 .iter()
1644 .filter_map(|message| message.text_content("\n"))
1645 .collect::<Vec<_>>();
1646 assert!(contents.contains(&"recent 4".to_string()));
1647 assert!(contents.contains(&"initial task".to_string()));
1648 assert!(!contents.contains(&"recent 5".to_string()));
1649 assert!(!contents.contains(&"client instructions".to_string()));
1650 assert_eq!(
1651 judge_request.llm_request.output.response_format,
1652 Some(judge.contract().response_format().clone())
1653 );
1654 assert_eq!(
1655 judge_request.llm_request.output.max_output_tokens,
1656 Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
1657 );
1658 Ok(())
1659 }
1660
1661 fn sample_value(spec: &Value) -> Value {
1662 if let Some(first) = spec
1663 .get("enum")
1664 .and_then(Value::as_array)
1665 .and_then(|values| values.first())
1666 {
1667 return first.clone();
1668 }
1669 match spec.get("type").and_then(Value::as_str) {
1670 Some("number") => serde_json::json!(0.5),
1671 Some("boolean") => serde_json::json!(false),
1672 _ => serde_json::json!("sample"),
1673 }
1674 }
1675
1676 fn schema_shaped_verdict(schema: &Value) -> Result<String> {
1677 let properties = schema
1678 .pointer("/json_schema/schema/properties")
1679 .and_then(Value::as_object)
1680 .ok_or_else(|| LibsyError::AlgorithmError {
1681 message: "packaged schema declares no properties".to_string(),
1682 })?;
1683 Ok(Value::Object(
1684 properties
1685 .iter()
1686 .map(|(name, spec)| (name.clone(), sample_value(spec)))
1687 .collect(),
1688 )
1689 .to_string())
1690 }
1691
1692 #[test]
1695 fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
1696 let contract =
1697 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1698 let schema = contract.response_format();
1699 let reply = schema_shaped_verdict(schema)?;
1700 let judge: CapabilityJudge = StructuredJudge::new(
1701 TaskInput {
1702 recent_turn_window: None,
1703 },
1704 contract,
1705 SerdeDecoder::new(),
1706 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1707 );
1708
1709 let verdict = judge.parse(&text_response(None, reply))?;
1710
1711 assert!(verdict.is_valid());
1712 assert!((0.0..=1.0).contains(&verdict.p_solve));
1713 Ok(())
1714 }
1715
1716 #[test]
1717 fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
1718 let contract =
1719 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1720 let prompt = contract.system_prompt();
1721 let schema_name = contract
1722 .response_format()
1723 .pointer("/json_schema/name")
1724 .and_then(Value::as_str)
1725 .ok_or_else(|| LibsyError::AlgorithmError {
1726 message: "packaged response schema has no name".to_string(),
1727 })?;
1728 assert_eq!(schema_name, "CapabilityClassifierDecision");
1729 assert!(prompt.contains("SUP-1 [supported]"));
1730 assert!(prompt.contains("SUP-5 [supported]"));
1731 assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
1732 assert!(!prompt.contains("\"type\": \"object\""));
1733 assert!(!prompt.contains("\"json_schema\""));
1734 assert!(!prompt.contains(schema_name));
1735 let rule_values = contract
1736 .response_format()
1737 .pointer("/json_schema/schema/properties/primary_rule/enum")
1738 .and_then(Value::as_array)
1739 .ok_or_else(|| LibsyError::AlgorithmError {
1740 message: "rendered response schema has no primary rule enum".to_string(),
1741 })?;
1742 assert!(
1743 rule_values
1744 .iter()
1745 .any(|value| value.as_str() == Some("SUP-1"))
1746 );
1747 assert!(
1748 rule_values
1749 .iter()
1750 .any(|value| value.as_str() == Some("none"))
1751 );
1752 Ok(())
1753 }
1754}