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