1use std::collections::HashSet;
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::{Deserialize, Deserializer};
11use serde_json::Value;
12use switchyard_protocol::{Category, ContentBlock, Message, Role};
13
14use super::escalation;
15use super::fall_through::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::{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
199struct TaskClassifierPolicy {
200 base_threshold: f64,
201 threshold_step: f64,
202}
203
204impl TaskClassifierPolicy {
205 fn new(config: &TaskClassifierConfig) -> Self {
206 Self {
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(
222 &self,
223 verdict: Option<&Self::Verdict>,
224 driver: &Driver,
225 ) -> Result<Classification> {
226 let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
229 return Ok(Classification::Ambiguous(vec![]));
230 };
231 let Some(threshold) = self.threshold(verdict) else {
234 return Ok(Classification::Ambiguous(vec![]));
235 };
236 let category = if verdict.p_solve >= threshold
237 || (threshold - verdict.p_solve).abs() <= f64::EPSILON
238 {
239 Category::Efficient
240 } else {
241 Category::Capable
242 };
243 let Some(target) = driver.models_for(&category).first().cloned() else {
246 return Ok(Classification::Ambiguous(vec![]));
247 };
248 Ok(Classification::Scores(vec![Score {
249 target,
250 confidence: 1.0,
251 category: Some(category),
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(
490 &self,
491 verdict: Option<&Self::Verdict>,
492 driver: &Driver,
493 ) -> Result<Classification> {
494 match self {
495 Self::TargetSelector(policy) => policy.to_classification(verdict, driver),
496 }
497 }
498}
499
500fn affinity_router(
502 trigger: ClassifyTrigger,
503 message_hash_fallback: bool,
504) -> Option<Arc<AffinityRouter>> {
505 let router = match trigger {
506 ClassifyTrigger::EveryRequest => return None,
507 ClassifyTrigger::NewSession => AffinityRouter::new(),
508 ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
509 };
510 let router = if message_hash_fallback {
511 router.with_message_hash_fallback()
512 } else {
513 router
514 };
515 Some(Arc::new(router))
516}
517
518pub struct LlmTaskClassifier {
520 route: FallThrough<State>,
521 inner: Arc<dyn Classifier<State>>,
523}
524
525struct ClassifierRouteConfig {
526 default_target: Category,
527 classify_trigger: ClassifyTrigger,
528 message_hash_fallback: bool,
529}
530
531pub struct DefaultCategoryClassifier(pub Category);
534
535#[async_trait]
536impl<S: Send> Classifier<S> for DefaultCategoryClassifier {
537 async fn score(
538 &self,
539 _state: &mut S,
540 _request: &mut Request,
541 driver: &Driver,
542 ) -> Result<(Classification, Option<Response>)> {
543 let target = driver.first_model_for(&self.0)?;
544 driver.set_evidence_if_empty(serde_json::json!({"source": "fall_open"}));
545 Ok((
546 Classification::Scores(vec![Score {
547 target: target.clone(),
548 confidence: 0.0,
549 category: Some(self.0.clone()),
550 }]),
551 None,
552 ))
553 }
554}
555
556#[derive(Clone)]
558#[non_exhaustive]
559pub enum LlmClassifierConfig {
560 Capability {
562 config: TaskClassifierConfig,
564 },
565 Escalation {
567 contract: ClassifierContractConfig,
569 config: EscalationJudgeConfig,
571 max_output_tokens: u64,
573 },
574 Custom {
576 default_target: Category,
578 config: CustomClassifierConfig,
580 },
581}
582
583impl LlmTaskClassifier {
584 pub fn new(config: LlmClassifierConfig) -> Result<Self> {
591 match config {
592 LlmClassifierConfig::Capability { config } => Self::build_capability(config),
593 LlmClassifierConfig::Escalation {
594 contract,
595 config,
596 max_output_tokens,
597 } => Self::build_escalation(contract, config, max_output_tokens),
598 LlmClassifierConfig::Custom {
599 default_target,
600 config,
601 } => Self::build_custom(default_target, config),
602 }
603 }
604
605 fn build_capability(config: TaskClassifierConfig) -> Result<Self> {
606 config.validate()?;
607 let contract = Self::load_capability_contract(&config.contract)?;
608 let classify_trigger = config.classify_trigger;
609 let message_hash_fallback = config.message_hash_fallback;
610 let classifier: Arc<dyn Classifier<State>> = Arc::new(
611 JudgeClassifier::new(
612 StructuredJudge::new(
613 TaskInput {
614 recent_turn_window: config.recent_turn_window,
615 },
616 contract,
617 SerdeDecoder::new(),
618 JudgeRuntimeConfig::new(config.max_output_tokens)?,
619 ),
620 TaskClassifierPolicy::new(&config),
621 )
622 .with_evidence(capability_evidence),
623 );
624 Self::from_classifier(
625 classifier,
626 ClassifierRouteConfig {
627 default_target: Category::Capable,
628 classify_trigger,
629 message_hash_fallback,
630 },
631 )
632 }
633
634 fn build_custom(default_target: Category, config: CustomClassifierConfig) -> Result<Self> {
635 config.validate()?;
636 let CustomClassifierConfig {
637 prompt,
638 response_schema,
639 policy,
640 classify_trigger,
641 message_hash_fallback,
642 recent_turn_window,
643 max_output_tokens,
644 } = config;
645 let contract = ClassifierContract::from_inner_schema(&prompt, response_schema)?;
646 let policy = match policy {
647 CustomClassifierPolicy::TargetSelector { selector } => {
648 CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(selector)?)
649 }
650 };
651 let classifier: Arc<dyn Classifier<State>> = Arc::new(JudgeClassifier::new(
652 StructuredJudge::new(
653 TaskInput { recent_turn_window },
654 contract,
655 JsonSchemaDecoder::new(),
656 JudgeRuntimeConfig::new(max_output_tokens)?,
657 ),
658 policy,
659 ));
660
661 Self::from_classifier(
662 classifier,
663 ClassifierRouteConfig {
664 default_target,
665 classify_trigger,
666 message_hash_fallback,
667 },
668 )
669 }
670
671 fn build_escalation(
672 contract_config: ClassifierContractConfig,
673 config: EscalationJudgeConfig,
674 max_output_tokens: u64,
675 ) -> Result<Self> {
676 let inner = escalation::build_classifier(contract_config, config, max_output_tokens)?;
677 Ok(Self {
678 route: FallThrough::<State>::new_with_state()
679 .with_name(ALGORITHM_NAME)
680 .with_classifier(Arc::clone(&inner)),
681 inner,
682 })
683 }
684
685 fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
687 ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
688 }
689
690 fn from_classifier(
692 inner: Arc<dyn Classifier<State>>,
693 config: ClassifierRouteConfig,
694 ) -> Result<Self> {
695 if config.message_hash_fallback && config.classify_trigger != ClassifyTrigger::NewSession {
696 return Err(LibsyError::AlgorithmError {
697 message: "message_hash_fallback requires classify_trigger = new_session"
698 .to_string(),
699 });
700 }
701 let mut route = FallThrough::<State>::new_with_state().with_name(ALGORITHM_NAME);
703 if let Some(affinity) =
704 affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
705 {
706 route = route
708 .with_processor(affinity.clone())
709 .with_classifier(affinity.clone());
710 }
711 let fallback = DefaultCategoryClassifier(config.default_target);
712 Ok(Self {
713 route: route
714 .with_classifier(inner.clone())
715 .with_classifier(Arc::new(fallback)),
716 inner,
717 })
718 }
719}
720
721#[async_trait]
722impl Classifier<State> for LlmTaskClassifier {
723 async fn score(
724 &self,
725 state: &mut State,
726 request: &mut Request,
727 driver: &Driver,
728 ) -> Result<(Classification, Option<Response>)> {
729 self.inner.score(state, request, driver).await
730 }
731}
732
733#[async_trait]
734impl Algorithm for LlmTaskClassifier {
735 fn name(&self) -> &str {
736 "llm_task_classifier"
737 }
738
739 async fn route(
740 self: Arc<Self>,
741 driver: Driver,
742 request: Request,
743 ) -> Result<crate::RoutingOutcome> {
744 self.route.execute(driver, request).await
745 }
746}
747
748#[cfg(test)]
749mod tests {
750 use std::collections::HashMap;
751 use std::sync::Arc;
752
753 use parking_lot::Mutex;
754 use serde_json::Value;
755
756 use super::*;
757 use switchyard_protocol::{
758 ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ModelId, ToolCall,
759 ToolResult, completion_text, text_request, text_response,
760 };
761
762 use crate::algorithms::util::llm_judge::Judge;
763 use crate::core::testing::{Serve, test_drive_with_models};
764 use switchyard_protocol::{LlmResponse, Response};
765
766 const TEST_THRESHOLD: f64 = 0.5;
767
768 type CapabilityJudge = StructuredJudge<TaskInput, SerdeDecoder<TaskClassifierVerdict>>;
769
770 fn test_config(base_threshold: f64) -> TaskClassifierConfig {
771 TaskClassifierConfig {
772 base_threshold,
773 ..TaskClassifierConfig::default()
774 }
775 }
776
777 fn policy() -> TaskClassifierPolicy {
778 TaskClassifierPolicy::new(&test_config(TEST_THRESHOLD))
779 }
780
781 fn runtime_models() -> HashMap<Category, Vec<ModelId>> {
782 [
783 (Category::Judge, vec![ModelId::from("judge")]),
784 (Category::Efficient, vec![ModelId::from("efficient")]),
785 (Category::Capable, vec![ModelId::from("capable")]),
786 (
787 Category::Any,
788 vec![ModelId::from("efficient"), ModelId::from("capable")],
789 ),
790 ]
791 .into()
792 }
793
794 fn policy_driver() -> Driver {
795 Driver::new("test", Arc::new(runtime_models().into())).0
796 }
797
798 fn verdict(
799 p_solve: f64,
800 capability_boundary: &str,
801 primary_rule: &str,
802 ) -> TaskClassifierVerdict {
803 TaskClassifierVerdict {
804 crux: "test crux".to_string(),
805 primary_rule: primary_rule.to_string(),
806 capability_boundary: capability_boundary.to_string(),
807 p_solve,
808 }
809 }
810
811 fn selected(
812 policy: &TaskClassifierPolicy,
813 verdict: Option<&TaskClassifierVerdict>,
814 ) -> Result<ModelId> {
815 policy
816 .to_classification(verdict, &policy_driver())?
817 .argmax(false)?
818 .map(|score| score.target)
819 .ok_or_else(|| LibsyError::AlgorithmError {
820 message: "policy abstained".to_string(),
821 })
822 }
823
824 #[derive(Default)]
827 struct Recorder {
828 calls: Mutex<Vec<String>>,
829 call_roles: Mutex<Vec<(String, bool)>>,
830 judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
831 judge_system_prompts: Mutex<Vec<String>>,
832 }
833
834 impl Recorder {
835 fn calls(&self) -> Vec<String> {
836 self.calls.lock().clone()
837 }
838
839 fn call_roles(&self) -> Vec<(String, bool)> {
840 self.call_roles.lock().clone()
841 }
842
843 fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
844 self.judge_max_output_tokens.lock().clone()
845 }
846
847 fn judge_system_prompts(&self) -> Vec<String> {
848 self.judge_system_prompts.lock().clone()
849 }
850
851 fn serve(self: &Arc<Self>) -> impl Serve {
852 let recorder = Arc::clone(self);
853 move |model: ModelId, request: Request| {
854 let recorder = Arc::clone(&recorder);
855 async move {
856 let model = model.to_string();
857 recorder.calls.lock().push(model.clone());
858 recorder
859 .call_roles
860 .lock()
861 .push((model.clone(), model != "judge"));
862 let completion = if model == "judge" {
863 recorder
864 .judge_max_output_tokens
865 .lock()
866 .push(request.llm_request.output.max_output_tokens);
867 recorder.judge_system_prompts.lock().extend(
868 request
869 .llm_request
870 .instructions
871 .first()
872 .and_then(|instruction| {
873 instruction.content.iter().find_map(|b| {
874 if let ContentBlock::Text { text } = b {
875 Some(text.clone())
876 } else {
877 None
878 }
879 })
880 }),
881 );
882 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
883 } else {
884 format!("answer from {model}")
885 };
886 Ok(Response {
887 llm_response: LlmResponse::Agg(text_response(None, completion)),
888 metadata: request.metadata,
889 upstream_headers: http::HeaderMap::new(),
890 })
891 }
892 }
893 }
894 }
895
896 fn unreachable_judge() -> impl Serve {
898 |model: ModelId, request: Request| async move {
899 let model = model.to_string();
900 if model == "judge" {
901 return Err(LlmClientError::Timeout {
902 source: Box::new(std::io::Error::other("judge unreachable")),
903 });
904 }
905 Ok(Response {
906 llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
907 metadata: request.metadata,
908 upstream_headers: http::HeaderMap::new(),
909 })
910 }
911 }
912
913 fn router() -> Result<Arc<LlmTaskClassifier>> {
914 Ok(Arc::new(LlmTaskClassifier::new(
915 LlmClassifierConfig::Capability {
916 config: test_config(TEST_THRESHOLD),
917 },
918 )?))
919 }
920
921 fn classify_request() -> Request {
922 Request {
923 llm_request: text_request(Some("auto".to_string()), "classify this task"),
924 raw_request: None,
925 metadata: None,
926 }
927 }
928
929 fn classify_session_request() -> Request {
930 Request {
931 metadata: Some(Metadata {
932 session_id: Some("session-1".to_string()),
933 ..Metadata::default()
934 }),
935 ..classify_request()
936 }
937 }
938
939 fn classify_follow_up_request() -> Request {
940 let mut request = classify_request();
941 request
942 .llm_request
943 .messages
944 .push(Message::text(Role::Assistant, "I will add the test."));
945 request.llm_request.messages.push(Message::text(
946 Role::User,
947 "Now run the test suite and report the result.",
948 ));
949 request
950 }
951
952 #[tokio::test]
953 async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
954 let router = router()?;
955
956 let (selected_model, response) = test_drive_with_models(
957 router,
958 classify_request(),
959 runtime_models(),
960 unreachable_judge(),
961 )
962 .await?;
963
964 assert_eq!(selected_model, "capable");
965 assert_eq!(
966 response.llm_response.as_agg().map(completion_text),
967 Some("answer from capable".to_string())
968 );
969 Ok(())
970 }
971
972 #[tokio::test]
973 async fn classifier_judges_each_request_without_affinity() -> Result<()> {
974 let recorder = Arc::new(Recorder::default());
975 let router = router()?;
976 let request = classify_request();
977 let models = runtime_models();
978
979 test_drive_with_models(
980 router.clone(),
981 request.clone(),
982 models.clone(),
983 recorder.serve(),
984 )
985 .await?;
986 test_drive_with_models(router, request, models, recorder.serve()).await?;
987
988 assert_eq!(
989 recorder.calls(),
990 vec!["judge", "efficient", "judge", "efficient"]
991 );
992 assert_eq!(
993 recorder.call_roles(),
994 vec![
995 ("judge".to_string(), false),
996 ("efficient".to_string(), true),
997 ("judge".to_string(), false),
998 ("efficient".to_string(), true),
999 ]
1000 );
1001 Ok(())
1002 }
1003
1004 #[tokio::test]
1005 async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1006 let recorder = Arc::new(Recorder::default());
1007 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1008 config: TaskClassifierConfig {
1009 max_output_tokens: 512,
1010 ..test_config(TEST_THRESHOLD)
1011 },
1012 })?);
1013
1014 test_drive_with_models(
1015 router,
1016 classify_request(),
1017 runtime_models(),
1018 recorder.serve(),
1019 )
1020 .await?;
1021
1022 assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1023 Ok(())
1024 }
1025
1026 #[tokio::test]
1027 async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1028 let recorder = Arc::new(Recorder::default());
1029 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1030 config: TaskClassifierConfig {
1031 contract: ClassifierContractConfig::default()
1032 .with_prompt("Custom capability rubric."),
1033 ..test_config(TEST_THRESHOLD)
1034 },
1035 })?);
1036
1037 test_drive_with_models(
1038 router,
1039 classify_request(),
1040 runtime_models(),
1041 recorder.serve(),
1042 )
1043 .await?;
1044
1045 let prompts = recorder.judge_system_prompts();
1046 assert_eq!(prompts.len(), 1);
1047 assert_eq!(prompts[0], "Custom capability rubric.");
1048 Ok(())
1049 }
1050
1051 #[tokio::test]
1052 async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1053 let recorder = Arc::new(Recorder::default());
1054 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1055 config: TaskClassifierConfig {
1056 classify_trigger: ClassifyTrigger::NewSession,
1057 ..test_config(TEST_THRESHOLD)
1058 },
1059 })?);
1060
1061 let request = classify_session_request();
1062 let models = runtime_models();
1063 test_drive_with_models(
1064 router.clone(),
1065 request.clone(),
1066 models.clone(),
1067 recorder.serve(),
1068 )
1069 .await?;
1070 test_drive_with_models(router, request, models, recorder.serve()).await?;
1071
1072 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1073 Ok(())
1074 }
1075
1076 #[tokio::test]
1077 async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1078 let recorder = Arc::new(Recorder::default());
1079 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1080 config: TaskClassifierConfig {
1081 classify_trigger: ClassifyTrigger::NewSession,
1082 message_hash_fallback: true,
1083 recent_turn_window: None,
1084 ..test_config(TEST_THRESHOLD)
1085 },
1086 })?);
1087
1088 let models = runtime_models();
1089 test_drive_with_models(
1090 router.clone(),
1091 classify_request(),
1092 models.clone(),
1093 recorder.serve(),
1094 )
1095 .await?;
1096 test_drive_with_models(
1097 router,
1098 classify_follow_up_request(),
1099 models,
1100 recorder.serve(),
1101 )
1102 .await?;
1103
1104 assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1105 Ok(())
1106 }
1107
1108 #[tokio::test]
1109 async fn one_classifier_uses_each_requests_runtime_models() -> Result<()> {
1110 let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1111 config: TaskClassifierConfig {
1112 classify_trigger: ClassifyTrigger::NewSession,
1113 ..test_config(TEST_THRESHOLD)
1114 },
1115 })?);
1116 let calls = Arc::new(Mutex::new(Vec::new()));
1117 let serve = |calls: Arc<Mutex<Vec<String>>>| {
1118 move |model: ModelId, _request: Request| {
1119 let calls = Arc::clone(&calls);
1120 async move {
1121 calls.lock().push(model.to_string());
1122 let text = if model.as_str().starts_with("judge-") {
1123 r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
1124 } else {
1125 model.to_string()
1126 };
1127 Ok(Response {
1128 llm_response: LlmResponse::Agg(text_response(None, text)),
1129 metadata: None,
1130 upstream_headers: Default::default(),
1131 })
1132 }
1133 }
1134 };
1135 let models = |suffix: &str| -> HashMap<Category, Vec<ModelId>> {
1136 [
1137 (
1138 Category::Judge,
1139 vec![ModelId::from(format!("judge-{suffix}"))],
1140 ),
1141 (
1142 Category::Efficient,
1143 vec![ModelId::from(format!("efficient-{suffix}"))],
1144 ),
1145 (
1146 Category::Capable,
1147 vec![ModelId::from(format!("capable-{suffix}"))],
1148 ),
1149 (
1150 Category::Any,
1151 vec![
1152 ModelId::from(format!("efficient-{suffix}")),
1153 ModelId::from(format!("capable-{suffix}")),
1154 ],
1155 ),
1156 ]
1157 .into()
1158 };
1159 let request = classify_session_request();
1160
1161 let (first, _) = test_drive_with_models(
1162 router.clone(),
1163 request.clone(),
1164 models("a"),
1165 serve(Arc::clone(&calls)),
1166 )
1167 .await?;
1168 let (second, _) =
1169 test_drive_with_models(router, request, models("b"), serve(Arc::clone(&calls))).await?;
1170
1171 assert_eq!(first, "efficient-a");
1172 assert_eq!(second, "efficient-b");
1173 assert_eq!(
1174 &*calls.lock(),
1175 &["judge-a", "efficient-a", "judge-b", "efficient-b"]
1176 );
1177 Ok(())
1178 }
1179
1180 #[test]
1181 fn the_threshold_boundary_is_inclusive() -> Result<()> {
1182 let policy = policy();
1183 let at_threshold = verdict(0.5, "supported", "SUP-1");
1184 let below_threshold = verdict(0.49, "supported", "SUP-1");
1185 assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1186 assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1187 Ok(())
1188 }
1189
1190 #[test]
1191 fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1192 let borderline = verdict(0.5, "supported", "SUP-1");
1193 let strict = TaskClassifierPolicy::new(&test_config(0.9));
1194 let lenient = TaskClassifierPolicy::new(&test_config(0.1));
1195 assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1196 assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1197 Ok(())
1198 }
1199
1200 #[test]
1201 fn classifier_config_rejects_unknown_fields() {
1202 let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1203 "base_threshold": 0.5,
1204 "classifier_magic": true,
1205 }))
1206 .expect_err("unknown classifier fields must be rejected");
1207
1208 assert!(
1209 error
1210 .to_string()
1211 .contains("unknown field `classifier_magic`"),
1212 "{error}"
1213 );
1214 }
1215
1216 #[test]
1217 fn invalid_classifier_config_is_rejected() -> Result<()> {
1218 for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1219 assert!(
1220 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1221 config: test_config(bad),
1222 })
1223 .is_err(),
1224 "base threshold {bad} should be rejected"
1225 );
1226 }
1227 for config in [
1228 TaskClassifierConfig {
1229 base_threshold: 0.5,
1230 threshold_step: -0.1,
1231 ..TaskClassifierConfig::default()
1232 },
1233 TaskClassifierConfig {
1234 base_threshold: 0.8,
1235 threshold_step: 0.11,
1236 ..TaskClassifierConfig::default()
1237 },
1238 TaskClassifierConfig {
1239 base_threshold: 0.5,
1240 message_hash_fallback: true,
1241 ..TaskClassifierConfig::default()
1242 },
1243 TaskClassifierConfig {
1244 base_threshold: 0.5,
1245 max_output_tokens: 0,
1246 ..TaskClassifierConfig::default()
1247 },
1248 ] {
1249 assert!(LlmTaskClassifier::new(LlmClassifierConfig::Capability { config }).is_err());
1250 }
1251 for base_threshold in [0.0, 1.0] {
1252 LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1253 config: test_config(base_threshold),
1254 })?;
1255 }
1256 Ok(())
1257 }
1258
1259 #[test]
1260 fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1261 let policy = policy();
1262 let inconsistent_rule = TaskClassifierVerdict {
1263 capability_boundary: "uncertain".to_string(),
1264 ..verdict(1.0, "supported", "SUP-1")
1265 };
1266 let empty_crux = TaskClassifierVerdict {
1267 crux: " ".to_string(),
1268 ..verdict(1.0, "supported", "SUP-1")
1269 };
1270 let unusable = [
1271 Some(verdict(1.1, "supported", "SUP-1")),
1272 Some(inconsistent_rule),
1273 Some(empty_crux),
1274 None,
1275 ];
1276 for verdict in unusable {
1277 let classification = policy.to_classification(verdict.as_ref(), &policy_driver())?;
1278 assert!(matches!(classification, Classification::Ambiguous(_)));
1279 assert!(classification.argmax(false)?.is_none());
1280 assert!(classification.argmax(true)?.is_none());
1281 }
1282 Ok(())
1283 }
1284
1285 #[test]
1286 fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
1287 let policy = TaskClassifierPolicy::new(&TaskClassifierConfig {
1288 threshold_step: 0.1,
1289 ..test_config(0.4)
1290 });
1291
1292 assert_eq!(
1293 selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
1294 "efficient"
1295 );
1296 assert_eq!(
1297 selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
1298 "capable"
1299 );
1300 assert_eq!(
1301 selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
1302 "efficient"
1303 );
1304 assert_eq!(
1305 selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
1306 "efficient"
1307 );
1308 assert_eq!(
1309 selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
1310 "capable"
1311 );
1312 assert_eq!(
1313 selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
1314 "efficient"
1315 );
1316 Ok(())
1317 }
1318
1319 fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
1322 Ok(StructuredJudge::new(
1323 TaskInput { recent_turn_window },
1324 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
1325 SerdeDecoder::new(),
1326 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1327 ))
1328 }
1329
1330 fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
1331 let judge = capability_judge(Some(recent_turn_window))?;
1332 let request = Request {
1333 llm_request: LlmRequest {
1334 messages: vec![
1335 Message::text(Role::System, "client instructions"),
1336 Message::text(Role::User, "initial task"),
1337 Message::text(Role::Assistant, "old response"),
1338 Message::text(Role::User, "old follow-up"),
1339 Message::text(Role::Assistant, "recent 1"),
1340 Message::text(Role::User, "recent 2"),
1341 ],
1342 ..LlmRequest::default()
1343 },
1344 raw_request: None,
1345 metadata: None,
1346 };
1347 Ok(judge
1348 .build_request(&State::default(), &request)
1349 .llm_request
1350 .messages
1351 .iter()
1352 .filter_map(|message| message.text_content("\n"))
1353 .collect())
1354 }
1355
1356 #[test]
1357 fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
1358 let contents = judged_contents(2)?;
1360 assert!(contents.contains(&"client instructions".to_string()));
1361 assert!(contents.contains(&"initial task".to_string()));
1362 assert!(contents.contains(&"recent 1".to_string()));
1363 assert!(contents.contains(&"recent 2".to_string()));
1364 assert!(!contents.contains(&"old response".to_string()));
1365 Ok(())
1366 }
1367
1368 #[test]
1369 fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
1370 let contents = judged_contents(0)?;
1371 assert!(contents.contains(&"client instructions".to_string()));
1372 assert!(contents.contains(&"initial task".to_string()));
1373 assert!(!contents.contains(&"recent 2".to_string()));
1374 Ok(())
1375 }
1376
1377 fn tool_call(id: &str) -> Message {
1378 Message {
1379 role: Role::Assistant,
1380 content: vec![ContentBlock::ToolCall(ToolCall {
1381 id: id.to_string(),
1382 name: "search".to_string(),
1383 arguments: Value::Null,
1384 })],
1385 }
1386 }
1387
1388 fn tool_result(id: &str) -> Message {
1389 Message {
1390 role: Role::Tool,
1391 content: vec![ContentBlock::ToolResult(ToolResult {
1392 tool_call_id: id.to_string(),
1393 content: vec![ContentBlock::Text {
1394 text: "tool output".to_string(),
1395 }],
1396 is_error: None,
1397 })],
1398 }
1399 }
1400
1401 #[test]
1404 fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
1405 let messages = vec![
1406 Message::text(Role::System, "client instructions"),
1407 Message::text(Role::User, "initial task"),
1408 Message::text(Role::Assistant, "old response"),
1409 tool_call("call-1"),
1410 tool_result("call-1"),
1411 Message::text(Role::Assistant, "recent 1"),
1412 Message::text(Role::User, "recent 2"),
1413 Message::text(Role::Assistant, "recent 3"),
1414 Message::text(Role::User, "recent 4"),
1415 ];
1416
1417 let kept = trim_messages(&messages, 5);
1419
1420 assert_eq!(
1421 kept,
1422 vec![
1423 Message::text(Role::System, "client instructions"),
1424 Message::text(Role::User, "initial task"),
1425 tool_call("call-1"),
1426 tool_result("call-1"),
1427 Message::text(Role::Assistant, "recent 1"),
1428 Message::text(Role::User, "recent 2"),
1429 Message::text(Role::Assistant, "recent 3"),
1430 Message::text(Role::User, "recent 4"),
1431 ]
1432 );
1433 }
1434
1435 #[test]
1438 fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
1439 let messages = vec![
1440 Message::text(Role::System, "client instructions"),
1441 Message::text(Role::User, "initial task"),
1442 tool_call("x"),
1443 tool_result("x"),
1444 Message::text(Role::Assistant, "later"),
1445 tool_call("x"),
1446 tool_result("x"),
1447 ];
1448
1449 let kept = trim_messages(&messages, 4);
1451
1452 assert_eq!(
1453 kept,
1454 vec![
1455 Message::text(Role::System, "client instructions"),
1456 Message::text(Role::User, "initial task"),
1457 tool_call("x"),
1458 tool_result("x"),
1459 Message::text(Role::Assistant, "later"),
1460 tool_call("x"),
1461 tool_result("x"),
1462 ]
1463 );
1464 }
1465
1466 #[test]
1469 fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
1470 let messages = vec![
1471 Message::text(Role::System, "client instructions"),
1472 tool_call("orphan"),
1473 Message::text(Role::User, "initial task"),
1474 Message::text(Role::Assistant, "old response"),
1475 tool_result("orphan"),
1476 Message::text(Role::Assistant, "recent 1"),
1477 Message::text(Role::User, "recent 2"),
1478 ];
1479
1480 let kept = trim_messages(&messages, 3);
1481
1482 assert_eq!(
1483 kept,
1484 vec![
1485 Message::text(Role::System, "client instructions"),
1486 Message::text(Role::User, "initial task"),
1487 tool_result("orphan"),
1488 Message::text(Role::Assistant, "recent 1"),
1489 Message::text(Role::User, "recent 2"),
1490 ]
1491 );
1492 }
1493
1494 #[test]
1495 fn a_window_restates_the_routing_instruction_last() -> Result<()> {
1496 let contents = judged_contents(2)?;
1497
1498 assert_eq!(
1501 contents.last().map(String::as_str),
1502 Some(TRAILING_ROUTING_INSTRUCTION)
1503 );
1504 Ok(())
1505 }
1506
1507 #[test]
1512 fn a_window_drops_reasoning_but_keeps_visible_text_and_tool_pairs() {
1513 let messages = vec![
1514 Message::text(Role::System, "client instructions"),
1515 Message::text(Role::User, "initial task"),
1516 Message {
1517 role: Role::Assistant,
1518 content: vec![
1519 ContentBlock::Reasoning {
1520 text: "private chain of thought".to_string(),
1521 signature: None,
1522 details: Vec::new(),
1523 },
1524 ContentBlock::Text {
1525 text: "visible answer".to_string(),
1526 },
1527 ],
1528 },
1529 tool_call("call-1"),
1530 tool_result("call-1"),
1531 Message {
1532 role: Role::Assistant,
1533 content: vec![ContentBlock::Reasoning {
1534 text: "reasoning-only turn".to_string(),
1535 signature: None,
1536 details: Vec::new(),
1537 }],
1538 },
1539 Message::text(Role::User, "follow-up"),
1540 ];
1541 let request = Request {
1542 llm_request: LlmRequest {
1543 messages,
1544 ..LlmRequest::default()
1545 },
1546 raw_request: None,
1547 metadata: None,
1548 };
1549
1550 let built = TaskInput {
1551 recent_turn_window: Some(10),
1552 }
1553 .build_messages(&State::default(), &request);
1554
1555 assert!(
1556 !built
1557 .iter()
1558 .flat_map(|message| &message.content)
1559 .any(|block| matches!(block, ContentBlock::Reasoning { .. })),
1560 "{built:?}"
1561 );
1562 assert!(
1563 built
1564 .iter()
1565 .any(|message| message.text_content("\n").as_deref() == Some("visible answer"))
1566 );
1567 assert!(built.contains(&tool_call("call-1")));
1568 assert!(built.contains(&tool_result("call-1")));
1569 assert_eq!(built.len(), 7);
1572 assert!(built.iter().all(|message| !message.content.is_empty()));
1573 }
1574
1575 #[test]
1576 fn the_default_path_is_left_unchanged() -> Result<()> {
1577 let judge = capability_judge(None)?;
1580 let request = Request {
1581 llm_request: LlmRequest {
1582 messages: vec![Message::text(Role::User, "the task")],
1583 ..LlmRequest::default()
1584 },
1585 raw_request: None,
1586 metadata: None,
1587 };
1588
1589 let built = judge.build_request(&State::default(), &request);
1590
1591 assert_eq!(built.llm_request.messages.len(), 1);
1594 assert!(!built.llm_request.instructions.is_empty());
1595 assert!(
1596 !built
1597 .llm_request
1598 .messages
1599 .iter()
1600 .filter_map(|message| message.text_content("\n"))
1601 .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
1602 );
1603 Ok(())
1604 }
1605
1606 #[test]
1607 fn capability_judge_builds_a_structured_request() -> Result<()> {
1608 let judge = capability_judge(None)?;
1609 let request = Request {
1610 llm_request: LlmRequest {
1611 model: Some("inbound".to_string()),
1612 messages: vec![
1613 Message::text(Role::System, "client instructions"),
1614 Message::text(Role::Developer, "client developer instructions"),
1615 Message::text(Role::User, "initial task"),
1616 Message::text(Role::Assistant, "old response"),
1617 Message::text(Role::User, "old follow-up"),
1618 Message::text(Role::Assistant, "recent 1"),
1619 Message::text(Role::User, "recent 2"),
1620 Message::text(Role::Assistant, "recent 3"),
1621 Message::text(Role::User, "recent 4"),
1622 Message::text(Role::Assistant, "recent 5"),
1623 ],
1624 ..LlmRequest::default()
1625 },
1626 raw_request: None,
1627 metadata: None,
1628 };
1629 let judge_request = judge.build_request(&State::default(), &request);
1630
1631 assert_eq!(judge_request.llm_request.model, request.llm_request.model);
1632 assert_eq!(judge_request.llm_request.instructions.len(), 1);
1633 assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
1634 assert_eq!(
1635 judge_request.llm_request.instructions[0].content,
1636 InstructionBlock {
1637 role: Role::System,
1638 content: Message::text(Role::System, judge.contract().system_prompt()).content,
1639 }
1640 .content,
1641 );
1642 assert_eq!(judge_request.llm_request.messages.len(), 2);
1643 let contents = judge_request
1644 .llm_request
1645 .messages
1646 .iter()
1647 .filter_map(|message| message.text_content("\n"))
1648 .collect::<Vec<_>>();
1649 assert!(contents.contains(&"recent 4".to_string()));
1650 assert!(contents.contains(&"initial task".to_string()));
1651 assert!(!contents.contains(&"recent 5".to_string()));
1652 assert!(!contents.contains(&"client instructions".to_string()));
1653 assert_eq!(
1654 judge_request.llm_request.output.response_format,
1655 Some(judge.contract().response_format().clone())
1656 );
1657 assert_eq!(
1658 judge_request.llm_request.output.max_output_tokens,
1659 Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
1660 );
1661 Ok(())
1662 }
1663
1664 fn sample_value(spec: &Value) -> Value {
1665 if let Some(first) = spec
1666 .get("enum")
1667 .and_then(Value::as_array)
1668 .and_then(|values| values.first())
1669 {
1670 return first.clone();
1671 }
1672 match spec.get("type").and_then(Value::as_str) {
1673 Some("number") => serde_json::json!(0.5),
1674 Some("boolean") => serde_json::json!(false),
1675 _ => serde_json::json!("sample"),
1676 }
1677 }
1678
1679 fn schema_shaped_verdict(schema: &Value) -> Result<String> {
1680 let properties = schema
1681 .pointer("/json_schema/schema/properties")
1682 .and_then(Value::as_object)
1683 .ok_or_else(|| LibsyError::AlgorithmError {
1684 message: "packaged schema declares no properties".to_string(),
1685 })?;
1686 Ok(Value::Object(
1687 properties
1688 .iter()
1689 .map(|(name, spec)| (name.clone(), sample_value(spec)))
1690 .collect(),
1691 )
1692 .to_string())
1693 }
1694
1695 #[test]
1698 fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
1699 let contract =
1700 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1701 let schema = contract.response_format();
1702 let reply = schema_shaped_verdict(schema)?;
1703 let judge: CapabilityJudge = StructuredJudge::new(
1704 TaskInput {
1705 recent_turn_window: None,
1706 },
1707 contract,
1708 SerdeDecoder::new(),
1709 JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1710 );
1711
1712 let verdict = judge.parse(&text_response(None, reply))?;
1713
1714 assert!(verdict.is_valid());
1715 assert!((0.0..=1.0).contains(&verdict.p_solve));
1716 Ok(())
1717 }
1718
1719 #[test]
1720 fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
1721 let contract =
1722 LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1723 let prompt = contract.system_prompt();
1724 let schema_name = contract
1725 .response_format()
1726 .pointer("/json_schema/name")
1727 .and_then(Value::as_str)
1728 .ok_or_else(|| LibsyError::AlgorithmError {
1729 message: "packaged response schema has no name".to_string(),
1730 })?;
1731 assert_eq!(schema_name, "CapabilityClassifierDecision");
1732 assert!(prompt.contains("SUP-1 [supported]"));
1733 assert!(prompt.contains("SUP-5 [supported]"));
1734 assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
1735 assert!(!prompt.contains("\"type\": \"object\""));
1736 assert!(!prompt.contains("\"json_schema\""));
1737 assert!(!prompt.contains(schema_name));
1738 let rule_values = contract
1739 .response_format()
1740 .pointer("/json_schema/schema/properties/primary_rule/enum")
1741 .and_then(Value::as_array)
1742 .ok_or_else(|| LibsyError::AlgorithmError {
1743 message: "rendered response schema has no primary rule enum".to_string(),
1744 })?;
1745 assert!(
1746 rule_values
1747 .iter()
1748 .any(|value| value.as_str() == Some("SUP-1"))
1749 );
1750 assert!(
1751 rule_values
1752 .iter()
1753 .any(|value| value.as_str() == Some("none"))
1754 );
1755 Ok(())
1756 }
1757}