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