Skip to main content

switchyard_libsy/algorithms/
llm_class.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Judge-backed capability, escalation, and custom-policy routing.
5
6use 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");
35/// Telemetry label for this algorithm's spans, metrics, and logs.
36const ALGORITHM_NAME: &str = "llm_task_classifier";
37
38/// Restates the task after the conversation the judge is asked to route.
39///
40/// The rubric leads the request, which works while the payload is a single short
41/// message. Once `recent_turn_window` includes real conversation, those instructions
42/// sit far from the generation point and the judge sometimes *answers* the
43/// conversation instead of classifying it — the reply then fails to parse, the
44/// verdict is unavailable, and the turn silently falls back. Repeating the task at
45/// the end is what reliably prevents that.
46const 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    /// Rejects malformed or internally inconsistent verdicts before policy evaluation.
60    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    /// Returns the number of threshold steps assigned to this capability boundary.
76    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
86/// Keeps the opening task and the last `recent_turn_window` turns after it. A
87/// window of `0` keeps the task alone.
88///
89/// Inbound decoders normalize client system and developer content into
90/// `LlmRequest::instructions`, so it never reaches this list.
91///
92/// Selects by reference and clones only what survives — a coding-agent
93/// conversation carries every tool result, so cloning it whole to keep a window
94/// would copy the transcript on each judged turn.
95fn 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
111/// The first index of the trailing window.
112///
113/// Counting messages alone can start the window between an assistant tool call and the
114/// result answering it, leaving the judge a result whose call id was never introduced. The
115/// start therefore moves back to the nearest one that keeps every tool pair whole.
116///
117/// One newest-to-oldest pass carries the ids still waiting for a call. Direction is what
118/// makes it correct: ids repeat across a conversation, and in this order a call is only
119/// ever seen after the results it could answer, so a later call — already passed — clears
120/// nothing. A result whose call sits before the opening task, which trimming never reaches,
121/// keeps the set non-empty to the end and falls back to the counted start, so an unpairable
122/// result costs one pass and cannot widen the window to the whole conversation.
123fn window_start(tail: &[&Message], recent_turn_window: usize) -> usize {
124    let counted = tail.len().saturating_sub(recent_turn_window);
125    // An empty window holds no result to pair, and the loop below never visits its start.
126    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        // Blocks reverse too, so a call answers a result only when it precedes it inside
132        // one message as well as across messages.
133        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
151/// Keeps the opening task and the latest user follow-up when they differ.
152fn 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
163/// Selects the task messages shown to capability and custom-schema classifiers.
164struct 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        // The default preserves the whole-task anchor and latest user update. A
171        // configured window widens that to the surrounding conversation.
172        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        // Only the windowed path carries assistant turns and tool traffic for the judge
177        // to be distracted by. The default path is user task messages only — the anchor
178        // and the latest follow-up — so there is nothing there to outrank.
179        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    /// Returns the required solve probability for one validated verdict.
213    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        // Judge output is untrusted. An absent, invalid, or inconsistent verdict is
223        // ambiguous so the surrounding router applies its configured fallback.
224        let Some(verdict) = verdict.filter(|verdict| verdict.is_valid()) else {
225            return Classification::Ambiguous(vec![]);
226        };
227        // A usable verdict below the capability threshold is still a decision: the judge
228        // does not trust the efficient tier with this task.
229        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/// Maps valid verdicts to scores, invalid verdicts to a reason, and leaves absent verdicts alone.
247fn capability_evidence(
248    policy: &TaskClassifierPolicy,
249    verdict: Option<&TaskClassifierVerdict>,
250) -> Option<Value> {
251    let verdict = verdict?;
252    let Some(threshold) = verdict
253        .is_valid()
254        .then(|| policy.threshold(verdict))
255        .flatten()
256    else {
257        return Some(serde_json::json!({
258            "source": "fail_open",
259            "reason_code": "invalid_verdict",
260        }));
261    };
262    Some(serde_json::json!({
263        "source": "llm-classifier",
264        "score": verdict.p_solve,
265        "threshold": threshold,
266    }))
267}
268
269#[derive(Clone, Debug)]
270/// Settings that control capability classifier prompting and routing.
271pub struct TaskClassifierConfig {
272    /// Lowest solve probability that routes a supported task to the efficient target.
273    pub base_threshold: f64,
274    /// Amount added per capability-boundary step.
275    ///
276    /// Supported verdicts use `base_threshold`, uncertain and unmatched verdicts use one
277    /// step, and unsupported verdicts use two steps.
278    pub threshold_step: f64,
279    /// How often the classifier re-decides this session's target.
280    pub classify_trigger: ClassifyTrigger,
281    /// Uses the first user message as the SessionKey for sticky routing when session metadata is unavailable.
282    pub message_hash_fallback: bool,
283    /// Trailing conversation turns the judge sees on top of the client
284    /// instructions and the opening task.
285    ///
286    /// `None` (the default) judges the opening task and latest user follow-up.
287    /// `Some(n)` widens that to the client instructions, the opening task, and
288    /// the last `n` turns after it.
289    pub recent_turn_window: Option<usize>,
290    /// Prompt and verdict contract settings for the classifier judge.
291    pub contract: ClassifierContractConfig,
292    /// Maximum completion tokens available to the classifier verdict.
293    pub max_output_tokens: u64,
294}
295
296/// Flat serialized shape that maps prompt settings into the runtime contract.
297#[derive(Deserialize)]
298#[serde(deny_unknown_fields)]
299struct TaskClassifierConfigWire {
300    base_threshold: f64,
301    #[serde(default)]
302    threshold_step: f64,
303    #[serde(default)]
304    classify_trigger: ClassifyTrigger,
305    #[serde(default)]
306    message_hash_fallback: bool,
307    #[serde(default)]
308    recent_turn_window: Option<usize>,
309    #[serde(default)]
310    prompt: Option<String>,
311    #[serde(default)]
312    response_format_type: ClassifierResponseFormat,
313    #[serde(default = "default_judge_max_output_tokens")]
314    max_output_tokens: u64,
315}
316
317impl<'de> Deserialize<'de> for TaskClassifierConfig {
318    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
319    where
320        D: Deserializer<'de>,
321    {
322        let wire = TaskClassifierConfigWire::deserialize(deserializer)?;
323        let mut contract = ClassifierContractConfig::default();
324        if let Some(prompt) = wire.prompt {
325            contract = contract.with_prompt(prompt);
326        }
327        contract = contract.with_response_format_type(wire.response_format_type);
328        Ok(Self {
329            base_threshold: wire.base_threshold,
330            threshold_step: wire.threshold_step,
331            classify_trigger: wire.classify_trigger,
332            message_hash_fallback: wire.message_hash_fallback,
333            recent_turn_window: wire.recent_turn_window,
334            contract,
335            max_output_tokens: wire.max_output_tokens,
336        })
337    }
338}
339
340const fn default_judge_max_output_tokens() -> u64 {
341    DEFAULT_JUDGE_MAX_OUTPUT_TOKENS
342}
343
344impl Default for TaskClassifierConfig {
345    fn default() -> Self {
346        Self {
347            base_threshold: 0.0,
348            threshold_step: 0.0,
349            classify_trigger: ClassifyTrigger::default(),
350            message_hash_fallback: false,
351            recent_turn_window: None,
352            contract: ClassifierContractConfig::default(),
353            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
354        }
355    }
356}
357
358impl TaskClassifierConfig {
359    /// Validates routing thresholds before the classifier is constructed.
360    fn validate(&self) -> Result<()> {
361        if !(0.0..=1.0).contains(&self.base_threshold) {
362            return Err(LibsyError::AlgorithmError {
363                message: format!(
364                    "base_threshold must be between 0 and 1, got {}",
365                    self.base_threshold
366                ),
367            });
368        }
369        if !self.threshold_step.is_finite() || self.threshold_step < 0.0 {
370            return Err(LibsyError::AlgorithmError {
371                message: format!(
372                    "threshold_step must be finite and greater than or equal to 0, got {}",
373                    self.threshold_step
374                ),
375            });
376        }
377        let unsupported_threshold = self.base_threshold + 2.0 * self.threshold_step;
378        if unsupported_threshold > 1.0 && unsupported_threshold - 1.0 > f64::EPSILON {
379            return Err(LibsyError::AlgorithmError {
380                message: format!(
381                    "base_threshold + 2 * threshold_step must be at most 1, got {unsupported_threshold}"
382                ),
383            });
384        }
385        if self.max_output_tokens == 0 {
386            return Err(LibsyError::AlgorithmError {
387                message: "max_output_tokens must be at least 1".to_string(),
388            });
389        }
390        if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
391            return Err(LibsyError::AlgorithmError {
392                message: "message_hash_fallback requires classify_trigger = new_session"
393                    .to_string(),
394            });
395        }
396        Ok(())
397    }
398}
399
400/// Policy that maps a custom classifier verdict to a routing target.
401#[derive(Clone, Debug)]
402pub enum CustomClassifierPolicy {
403    /// Resolves a JSON Pointer and treats its string value as a configured target label.
404    TargetSelector {
405        /// JSON Pointer evaluated against each schema-validated verdict.
406        selector: String,
407    },
408}
409
410impl CustomClassifierPolicy {
411    /// Creates a policy that selects a target label through a JSON Pointer.
412    pub fn target_selector(selector: impl Into<String>) -> Self {
413        Self::TargetSelector {
414            selector: selector.into(),
415        }
416    }
417}
418
419/// Settings for a classifier whose JSON Schema and target-selection policy are user supplied.
420#[derive(Clone, Debug)]
421pub struct CustomClassifierConfig {
422    /// System prompt sent to the classifier judge.
423    pub prompt: String,
424    /// Inner JSON Schema placed inside the provider's structured-output wrapper.
425    pub response_schema: Value,
426    /// Deterministic policy applied after the verdict passes schema validation.
427    pub policy: CustomClassifierPolicy,
428    /// How often the classifier re-decides this session's target.
429    pub classify_trigger: ClassifyTrigger,
430    /// Uses the first user message when session metadata is unavailable.
431    pub message_hash_fallback: bool,
432    /// Trailing conversation turns shown to the classifier judge.
433    pub recent_turn_window: Option<usize>,
434    /// Maximum completion tokens available to the classifier verdict.
435    pub max_output_tokens: u64,
436}
437
438impl CustomClassifierConfig {
439    /// Creates a custom-schema classifier contract with conservative runtime defaults.
440    pub fn new(
441        prompt: impl Into<String>,
442        response_schema: Value,
443        policy: CustomClassifierPolicy,
444    ) -> Self {
445        Self {
446            prompt: prompt.into(),
447            response_schema,
448            policy,
449            classify_trigger: ClassifyTrigger::default(),
450            message_hash_fallback: false,
451            recent_turn_window: None,
452            max_output_tokens: DEFAULT_JUDGE_MAX_OUTPUT_TOKENS,
453        }
454    }
455
456    fn validate(&self) -> Result<()> {
457        if self.max_output_tokens == 0 {
458            return Err(LibsyError::AlgorithmError {
459                message: "max_output_tokens must be at least 1".to_string(),
460            });
461        }
462        if self.message_hash_fallback && self.classify_trigger != ClassifyTrigger::NewSession {
463            return Err(LibsyError::AlgorithmError {
464                message: "message_hash_fallback requires classify_trigger = new_session"
465                    .to_string(),
466            });
467        }
468        Ok(())
469    }
470}
471
472enum CustomPolicyRuntime {
473    TargetSelector(TargetSelectorPolicy),
474}
475
476impl JudgePolicy for CustomPolicyRuntime {
477    type Verdict = Value;
478
479    fn to_classification(&self, verdict: Option<&Self::Verdict>) -> Classification {
480        match self {
481            Self::TargetSelector(policy) => policy.to_classification(verdict),
482        }
483    }
484}
485
486struct TaskClassifier {
487    classifier: JudgeClassifier<CapabilityJudge, TaskClassifierPolicy>,
488    capable_target: ModelId,
489}
490
491/// Builds the affinity router a trigger calls for, if any.
492fn affinity_router(
493    trigger: ClassifyTrigger,
494    message_hash_fallback: bool,
495) -> Option<Arc<AffinityRouter>> {
496    let router = match trigger {
497        ClassifyTrigger::EveryRequest => return None,
498        ClassifyTrigger::NewSession => AffinityRouter::new(),
499        ClassifyTrigger::UserTurn => AffinityRouter::new().with_release_on_user_turn(),
500    };
501    let router = if message_hash_fallback {
502        router.with_message_hash_fallback()
503    } else {
504        router
505    };
506    Some(Arc::new(router))
507}
508
509/// Routes requests through a capability, escalation, or custom classifier mode.
510pub struct LlmTaskClassifier {
511    route: FallThrough<State>,
512    /// Classifier used when this router is embedded in another cascade.
513    inner: Arc<dyn Classifier<State>>,
514}
515
516struct ClassifierRouteConfig {
517    default_target: ModelId,
518    classify_trigger: ClassifyTrigger,
519    message_hash_fallback: bool,
520}
521
522/// Complete construction settings for one LLM classifier mode.
523#[derive(Clone)]
524#[non_exhaustive]
525pub enum LlmClassifierConfig {
526    /// Routes between efficient and capable targets from a task-level verdict.
527    Capability {
528        /// Target that produces classifier verdicts.
529        judge_target: ModelId,
530        /// Target used when the efficient tier can handle the task.
531        efficient_target: ModelId,
532        /// Target used when the task needs the capable tier.
533        capable_target: ModelId,
534        /// Capability classifier settings.
535        config: TaskClassifierConfig,
536    },
537    /// Judges efficient responses and escalates after a confirmed streak.
538    Escalation {
539        /// Target that produces escalation verdicts.
540        judge_target: ModelId,
541        /// Target called before each escalation decision.
542        efficient_target: ModelId,
543        /// Target used after escalation is confirmed.
544        capable_target: ModelId,
545        /// Prompt and verdict contract settings for the escalation judge.
546        contract: ClassifierContractConfig,
547        /// Escalation policy settings.
548        config: EscalationJudgeConfig,
549        /// Maximum completion tokens available to the escalation verdict.
550        max_output_tokens: u64,
551    },
552    /// Routes among named targets using a user-supplied schema and policy.
553    Custom {
554        /// Target that produces classifier verdicts.
555        judge_target: ModelId,
556        /// User-facing labels paired with their resolved routing targets.
557        targets: Vec<(String, ModelId)>,
558        /// Label selected when the judge does not produce a usable verdict.
559        default_target: String,
560        /// Custom classifier settings.
561        config: CustomClassifierConfig,
562    },
563}
564
565impl LlmTaskClassifier {
566    /// Builds the classifier mode described by `config`.
567    ///
568    /// # Errors
569    ///
570    /// Returns an error when the selected mode's targets, contract, policy, or runtime
571    /// settings are invalid.
572    pub fn new(config: LlmClassifierConfig) -> Result<Self> {
573        match config {
574            LlmClassifierConfig::Capability {
575                judge_target,
576                efficient_target,
577                capable_target,
578                config,
579            } => Self::build_capability(judge_target, efficient_target, capable_target, config),
580            LlmClassifierConfig::Escalation {
581                judge_target,
582                efficient_target,
583                capable_target,
584                contract,
585                config,
586                max_output_tokens,
587            } => Self::build_escalation(
588                judge_target,
589                efficient_target,
590                capable_target,
591                contract,
592                config,
593                max_output_tokens,
594            ),
595            LlmClassifierConfig::Custom {
596                judge_target,
597                targets,
598                default_target,
599                config,
600            } => Self::build_custom(judge_target, targets, default_target, config),
601        }
602    }
603
604    fn build_capability(
605        judge_target: ModelId,
606        efficient_target: ModelId,
607        capable_target: ModelId,
608        config: TaskClassifierConfig,
609    ) -> Result<Self> {
610        config.validate()?;
611        let contract = Self::load_capability_contract(&config.contract)?;
612        let targets = vec![efficient_target.clone(), capable_target.clone()];
613        let classify_trigger = config.classify_trigger;
614        let message_hash_fallback = config.message_hash_fallback;
615        let classifier = Arc::new(TaskClassifier {
616            classifier: JudgeClassifier::new(
617                StructuredJudge::new(
618                    TaskInput {
619                        recent_turn_window: config.recent_turn_window,
620                    },
621                    contract,
622                    SerdeDecoder::new(),
623                    JudgeRuntimeConfig::new(config.max_output_tokens)?,
624                ),
625                judge_target.clone(),
626                TaskClassifierPolicy::new(
627                    efficient_target.clone(),
628                    capable_target.clone(),
629                    &config,
630                ),
631            )
632            .with_evidence(capability_evidence),
633            capable_target: capable_target.clone(),
634        });
635        let inner: Arc<dyn Classifier<State>> = classifier.clone();
636        Self::from_classifier(
637            targets,
638            inner,
639            ClassifierRouteConfig {
640                default_target: classifier.capable_target.clone(),
641                classify_trigger,
642                message_hash_fallback,
643            },
644        )
645    }
646
647    fn build_custom(
648        judge_target: ModelId,
649        targets: Vec<(String, ModelId)>,
650        default_target: String,
651        config: CustomClassifierConfig,
652    ) -> Result<Self> {
653        config.validate()?;
654        if targets.len() < 2 {
655            return Err(LibsyError::AlgorithmError {
656                message: "custom classifier requires at least two targets".to_string(),
657            });
658        }
659
660        let mut labels = BTreeSet::new();
661        let mut resolved_names = BTreeSet::new();
662        let mut target_map = BTreeMap::new();
663        let mut resolved_targets = Vec::with_capacity(targets.len());
664        for (label, target) in targets {
665            if label.trim().is_empty() || label.trim() != label {
666                return Err(LibsyError::AlgorithmError {
667                    message: "custom classifier target labels must be non-empty and have no surrounding whitespace"
668                        .to_string(),
669                });
670            }
671            if !labels.insert(label.clone()) {
672                return Err(LibsyError::AlgorithmError {
673                    message: format!("custom classifier target label {label:?} is duplicated"),
674                });
675            }
676            if !resolved_names.insert(target.clone()) {
677                return Err(LibsyError::AlgorithmError {
678                    message: format!("custom classifier resolved target {target:?} is duplicated"),
679                });
680            }
681            target_map.insert(label, target.clone());
682            resolved_targets.push(target);
683        }
684        let default_name =
685            target_map
686                .get(&default_target)
687                .cloned()
688                .ok_or_else(|| LibsyError::AlgorithmError {
689                    message: format!(
690                        "default_target {default_target:?} must be one of the configured targets"
691                    ),
692                })?;
693
694        let CustomClassifierConfig {
695            prompt,
696            response_schema,
697            policy,
698            classify_trigger,
699            message_hash_fallback,
700            recent_turn_window,
701            max_output_tokens,
702        } = config;
703        let contract = ClassifierContract::from_inner_schema(&prompt, response_schema)?;
704        let policy = match policy {
705            CustomClassifierPolicy::TargetSelector { selector } => {
706                CustomPolicyRuntime::TargetSelector(TargetSelectorPolicy::new(
707                    selector, target_map,
708                )?)
709            }
710        };
711        let classifier: Arc<dyn Classifier<State>> = Arc::new(JudgeClassifier::new(
712            StructuredJudge::new(
713                TaskInput { recent_turn_window },
714                contract,
715                JsonSchemaDecoder::new(),
716                JudgeRuntimeConfig::new(max_output_tokens)?,
717            ),
718            judge_target,
719            policy,
720        ));
721
722        Self::from_classifier(
723            resolved_targets,
724            classifier,
725            ClassifierRouteConfig {
726                default_target: default_name,
727                classify_trigger,
728                message_hash_fallback,
729            },
730        )
731    }
732
733    fn build_escalation(
734        judge_target: ModelId,
735        efficient_target: ModelId,
736        capable_target: ModelId,
737        contract_config: ClassifierContractConfig,
738        config: EscalationJudgeConfig,
739        max_output_tokens: u64,
740    ) -> Result<Self> {
741        let inner = escalation::build_classifier(
742            judge_target,
743            &efficient_target,
744            &capable_target,
745            contract_config,
746            config,
747            max_output_tokens,
748        )?;
749        let targets = vec![capable_target, efficient_target];
750        Ok(Self {
751            route: FallThrough::<State>::new_with_state(targets)
752                .with_name(ALGORITHM_NAME)
753                .with_classifier(Arc::clone(&inner)),
754            inner,
755        })
756    }
757
758    /// Loads the packaged capability-classifier contract.
759    fn load_capability_contract(config: &ClassifierContractConfig) -> Result<ClassifierContract> {
760        ClassifierContract::from_config(config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)
761    }
762
763    /// Keeps affinity and fallback ordering identical across judge-backed modes.
764    fn from_classifier(
765        targets: Vec<ModelId>,
766        inner: Arc<dyn Classifier<State>>,
767        config: ClassifierRouteConfig,
768    ) -> Result<Self> {
769        algorithm::ensure_model_is_target(&targets, &config.default_target)?;
770        if config.message_hash_fallback && config.classify_trigger != ClassifyTrigger::NewSession {
771            return Err(LibsyError::AlgorithmError {
772                message: "message_hash_fallback requires classify_trigger = new_session"
773                    .to_string(),
774            });
775        }
776        // Affinity comes first so a retained assignment short-circuits the judge call.
777        let mut route = FallThrough::<State>::new_with_state(targets).with_name(ALGORITHM_NAME);
778        if let Some(affinity) =
779            affinity_router(config.classify_trigger, config.message_hash_fallback).as_ref()
780        {
781            // Both roles must share one `Arc` so the classifier reads what the processor wrote.
782            route = route
783                .with_processor(affinity.clone())
784                .with_classifier(affinity.clone());
785        }
786        let fallback = DefaultTarget::new(config.default_target);
787        Ok(Self {
788            route: route
789                .with_classifier(inner.clone())
790                .with_classifier(Arc::new(fallback)),
791            inner,
792        })
793    }
794}
795
796#[async_trait]
797impl Classifier<State> for TaskClassifier {
798    async fn score(
799        &self,
800        state: &mut State,
801        request: &mut Request,
802        driver: Option<&Driver>,
803    ) -> Result<(Classification, Option<Response>)> {
804        self.classifier.score(state, request, driver).await
805    }
806}
807
808#[async_trait]
809impl Classifier<State> for LlmTaskClassifier {
810    async fn score(
811        &self,
812        state: &mut State,
813        request: &mut Request,
814        driver: Option<&Driver>,
815    ) -> Result<(Classification, Option<Response>)> {
816        self.inner.score(state, request, driver).await
817    }
818}
819
820#[async_trait]
821impl Algorithm for LlmTaskClassifier {
822    fn name(&self) -> &str {
823        "llm_task_classifier"
824    }
825
826    async fn route(
827        self: Arc<Self>,
828        driver: Driver,
829        request: Request,
830    ) -> Result<crate::RoutingOutcome> {
831        self.route.execute(driver, request).await
832    }
833}
834
835#[cfg(test)]
836mod tests {
837    use std::sync::Arc;
838
839    use parking_lot::Mutex;
840    use serde_json::Value;
841
842    use super::*;
843    use switchyard_protocol::{
844        ContentBlock, InstructionBlock, LlmClientError, LlmRequest, Metadata, ToolCall, ToolResult,
845        completion_text, text_request, text_response,
846    };
847
848    use crate::algorithms::util::llm_judge::Judge;
849    use crate::core::testing::{Serve, test_drive};
850    use switchyard_protocol::{LlmResponse, Response};
851
852    const TEST_THRESHOLD: f64 = 0.5;
853
854    fn test_config(base_threshold: f64) -> TaskClassifierConfig {
855        TaskClassifierConfig {
856            base_threshold,
857            ..TaskClassifierConfig::default()
858        }
859    }
860
861    fn policy() -> TaskClassifierPolicy {
862        TaskClassifierPolicy::new("efficient", "capable", &test_config(TEST_THRESHOLD))
863    }
864
865    fn verdict(
866        p_solve: f64,
867        capability_boundary: &str,
868        primary_rule: &str,
869    ) -> TaskClassifierVerdict {
870        TaskClassifierVerdict {
871            crux: "test crux".to_string(),
872            primary_rule: primary_rule.to_string(),
873            capability_boundary: capability_boundary.to_string(),
874            p_solve,
875        }
876    }
877
878    fn selected(
879        policy: &TaskClassifierPolicy,
880        verdict: Option<&TaskClassifierVerdict>,
881    ) -> Result<ModelId> {
882        policy
883            .to_classification(verdict)
884            .argmax(false)?
885            .map(|score| score.target)
886            .ok_or_else(|| LibsyError::AlgorithmError {
887                message: "policy abstained".to_string(),
888            })
889    }
890
891    /// Records what each target received; answers the judge with a supported verdict and
892    /// every other target with a plain completion.
893    #[derive(Default)]
894    struct Recorder {
895        calls: Mutex<Vec<String>>,
896        call_roles: Mutex<Vec<(String, bool)>>,
897        judge_max_output_tokens: Mutex<Vec<Option<u64>>>,
898        judge_system_prompts: Mutex<Vec<String>>,
899    }
900
901    impl Recorder {
902        fn calls(&self) -> Vec<String> {
903            self.calls.lock().clone()
904        }
905
906        fn call_roles(&self) -> Vec<(String, bool)> {
907            self.call_roles.lock().clone()
908        }
909
910        fn judge_max_output_tokens(&self) -> Vec<Option<u64>> {
911            self.judge_max_output_tokens.lock().clone()
912        }
913
914        fn judge_system_prompts(&self) -> Vec<String> {
915            self.judge_system_prompts.lock().clone()
916        }
917
918        fn serve(self: &Arc<Self>) -> impl Serve {
919            let recorder = Arc::clone(self);
920            move |model: ModelId, request: Request| {
921                let recorder = Arc::clone(&recorder);
922                async move {
923                    let model = model.to_string();
924                    recorder.calls.lock().push(model.clone());
925                    recorder
926                        .call_roles
927                        .lock()
928                        .push((model.clone(), model != "judge"));
929                    let completion = if model == "judge" {
930                        recorder
931                            .judge_max_output_tokens
932                            .lock()
933                            .push(request.llm_request.output.max_output_tokens);
934                        recorder.judge_system_prompts.lock().extend(
935                            request
936                                .llm_request
937                                .instructions
938                                .first()
939                                .and_then(|instruction| {
940                                    instruction.content.iter().find_map(|b| {
941                                        if let ContentBlock::Text { text } = b {
942                                            Some(text.clone())
943                                        } else {
944                                            None
945                                        }
946                                    })
947                                }),
948                        );
949                        r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#.to_string()
950                    } else {
951                        format!("answer from {model}")
952                    };
953                    Ok(Response {
954                        llm_response: LlmResponse::Agg(text_response(None, completion)),
955                        metadata: request.metadata,
956                    })
957                }
958            }
959        }
960    }
961
962    /// The judge times out; every other target answers normally.
963    fn unreachable_judge() -> impl Serve {
964        |model: ModelId, request: Request| async move {
965            let model = model.to_string();
966            if model == "judge" {
967                return Err(LlmClientError::Timeout {
968                    source: Box::new(std::io::Error::other("judge unreachable")),
969                });
970            }
971            Ok(Response {
972                llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))),
973                metadata: request.metadata,
974            })
975        }
976    }
977
978    fn router() -> Result<Arc<LlmTaskClassifier>> {
979        Ok(Arc::new(LlmTaskClassifier::new(
980            LlmClassifierConfig::Capability {
981                judge_target: ModelId::from("judge"),
982                efficient_target: ModelId::from("efficient"),
983                capable_target: ModelId::from("capable"),
984                config: test_config(TEST_THRESHOLD),
985            },
986        )?))
987    }
988
989    fn classify_request() -> Request {
990        Request {
991            llm_request: text_request(Some("auto".to_string()), "classify this task"),
992            raw_request: None,
993            metadata: None,
994        }
995    }
996
997    fn classify_session_request() -> Request {
998        Request {
999            metadata: Some(Metadata {
1000                session_id: Some("session-1".to_string()),
1001                ..Metadata::default()
1002            }),
1003            ..classify_request()
1004        }
1005    }
1006
1007    fn classify_follow_up_request() -> Request {
1008        let mut request = classify_request();
1009        request
1010            .llm_request
1011            .messages
1012            .push(Message::text(Role::Assistant, "I will add the test."));
1013        request.llm_request.messages.push(Message::text(
1014            Role::User,
1015            "Now run the test suite and report the result.",
1016        ));
1017        request
1018    }
1019
1020    #[tokio::test]
1021    async fn an_unreachable_judge_routes_capable_instead_of_failing_the_request() -> Result<()> {
1022        let router = router()?;
1023
1024        let (selected_model, response) =
1025            test_drive(router, classify_request(), unreachable_judge()).await?;
1026
1027        assert_eq!(selected_model, "capable");
1028        assert_eq!(
1029            response.llm_response.as_agg().map(completion_text),
1030            Some("answer from capable".to_string())
1031        );
1032        Ok(())
1033    }
1034
1035    #[tokio::test]
1036    async fn classifier_judges_each_request_without_affinity() -> Result<()> {
1037        let recorder = Arc::new(Recorder::default());
1038        let router = router()?;
1039        let request = classify_request;
1040
1041        test_drive(router.clone(), request(), recorder.serve()).await?;
1042        test_drive(router.clone(), request(), recorder.serve()).await?;
1043
1044        assert_eq!(
1045            recorder.calls(),
1046            vec!["judge", "efficient", "judge", "efficient"]
1047        );
1048        assert_eq!(
1049            recorder.call_roles(),
1050            vec![
1051                ("judge".to_string(), false),
1052                ("efficient".to_string(), true),
1053                ("judge".to_string(), false),
1054                ("efficient".to_string(), true),
1055            ]
1056        );
1057        Ok(())
1058    }
1059
1060    #[tokio::test]
1061    async fn classifier_config_sets_the_judge_completion_cap() -> Result<()> {
1062        let recorder = Arc::new(Recorder::default());
1063        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1064            judge_target: ModelId::from("judge"),
1065            efficient_target: ModelId::from("efficient"),
1066            capable_target: ModelId::from("capable"),
1067            config: TaskClassifierConfig {
1068                max_output_tokens: 512,
1069                ..test_config(TEST_THRESHOLD)
1070            },
1071        })?);
1072
1073        test_drive(router, classify_request(), recorder.serve()).await?;
1074
1075        assert_eq!(recorder.judge_max_output_tokens(), vec![Some(512)]);
1076        Ok(())
1077    }
1078
1079    #[tokio::test]
1080    async fn classifier_config_overrides_the_packaged_prompt() -> Result<()> {
1081        let recorder = Arc::new(Recorder::default());
1082        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1083            judge_target: ModelId::from("judge"),
1084            efficient_target: ModelId::from("efficient"),
1085            capable_target: ModelId::from("capable"),
1086            config: TaskClassifierConfig {
1087                contract: ClassifierContractConfig::default()
1088                    .with_prompt("Custom capability rubric."),
1089                ..test_config(TEST_THRESHOLD)
1090            },
1091        })?);
1092
1093        test_drive(router, classify_request(), recorder.serve()).await?;
1094
1095        let prompts = recorder.judge_system_prompts();
1096        assert_eq!(prompts.len(), 1);
1097        assert_eq!(prompts[0], "Custom capability rubric.");
1098        Ok(())
1099    }
1100
1101    #[tokio::test]
1102    async fn classifier_config_enables_new_session_trigger() -> Result<()> {
1103        let recorder = Arc::new(Recorder::default());
1104        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1105            judge_target: ModelId::from("judge"),
1106            efficient_target: ModelId::from("efficient"),
1107            capable_target: ModelId::from("capable"),
1108            config: TaskClassifierConfig {
1109                classify_trigger: ClassifyTrigger::NewSession,
1110                ..test_config(TEST_THRESHOLD)
1111            },
1112        })?);
1113
1114        let session_request = classify_session_request;
1115        test_drive(router.clone(), session_request(), recorder.serve()).await?;
1116        test_drive(router.clone(), session_request(), recorder.serve()).await?;
1117
1118        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1119        Ok(())
1120    }
1121
1122    #[tokio::test]
1123    async fn classifier_config_reuses_message_hash_affinity_for_a_follow_up() -> Result<()> {
1124        let recorder = Arc::new(Recorder::default());
1125        let router = Arc::new(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1126            judge_target: ModelId::from("judge"),
1127            efficient_target: ModelId::from("efficient"),
1128            capable_target: ModelId::from("capable"),
1129            config: TaskClassifierConfig {
1130                classify_trigger: ClassifyTrigger::NewSession,
1131                message_hash_fallback: true,
1132                recent_turn_window: None,
1133                ..test_config(TEST_THRESHOLD)
1134            },
1135        })?);
1136
1137        test_drive(router.clone(), classify_request(), recorder.serve()).await?;
1138        test_drive(
1139            router.clone(),
1140            classify_follow_up_request(),
1141            recorder.serve(),
1142        )
1143        .await?;
1144
1145        assert_eq!(recorder.calls(), vec!["judge", "efficient", "efficient"]);
1146        Ok(())
1147    }
1148
1149    #[test]
1150    fn the_threshold_boundary_is_inclusive() -> Result<()> {
1151        let policy = policy();
1152        let at_threshold = verdict(0.5, "supported", "SUP-1");
1153        let below_threshold = verdict(0.49, "supported", "SUP-1");
1154        assert_eq!(selected(&policy, Some(&at_threshold))?, "efficient");
1155        assert_eq!(selected(&policy, Some(&below_threshold))?, "capable");
1156        Ok(())
1157    }
1158
1159    #[test]
1160    fn the_threshold_moves_the_routing_boundary() -> Result<()> {
1161        let borderline = verdict(0.5, "supported", "SUP-1");
1162        let strict = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.9));
1163        let lenient = TaskClassifierPolicy::new("efficient", "capable", &test_config(0.1));
1164        assert_eq!(selected(&strict, Some(&borderline))?, "capable");
1165        assert_eq!(selected(&lenient, Some(&borderline))?, "efficient");
1166        Ok(())
1167    }
1168
1169    #[test]
1170    fn classifier_config_rejects_unknown_fields() {
1171        let error = serde_json::from_value::<TaskClassifierConfig>(serde_json::json!({
1172            "base_threshold": 0.5,
1173            "classifier_magic": true,
1174        }))
1175        .expect_err("unknown classifier fields must be rejected");
1176
1177        assert!(
1178            error
1179                .to_string()
1180                .contains("unknown field `classifier_magic`"),
1181            "{error}"
1182        );
1183    }
1184
1185    #[test]
1186    fn invalid_classifier_config_is_rejected() -> Result<()> {
1187        for bad in [1.5, -0.1, f64::NAN, f64::INFINITY] {
1188            assert!(
1189                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1190                    judge_target: ModelId::from("judge"),
1191                    efficient_target: ModelId::from("e"),
1192                    capable_target: ModelId::from("c"),
1193                    config: test_config(bad),
1194                })
1195                .is_err(),
1196                "base threshold {bad} should be rejected"
1197            );
1198        }
1199        for config in [
1200            TaskClassifierConfig {
1201                base_threshold: 0.5,
1202                threshold_step: -0.1,
1203                ..TaskClassifierConfig::default()
1204            },
1205            TaskClassifierConfig {
1206                base_threshold: 0.8,
1207                threshold_step: 0.11,
1208                ..TaskClassifierConfig::default()
1209            },
1210            TaskClassifierConfig {
1211                base_threshold: 0.5,
1212                message_hash_fallback: true,
1213                ..TaskClassifierConfig::default()
1214            },
1215            TaskClassifierConfig {
1216                base_threshold: 0.5,
1217                max_output_tokens: 0,
1218                ..TaskClassifierConfig::default()
1219            },
1220        ] {
1221            assert!(
1222                LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1223                    judge_target: ModelId::from("judge"),
1224                    efficient_target: ModelId::from("e"),
1225                    capable_target: ModelId::from("c"),
1226                    config,
1227                })
1228                .is_err()
1229            );
1230        }
1231        for base_threshold in [0.0, 1.0] {
1232            LlmTaskClassifier::new(LlmClassifierConfig::Capability {
1233                judge_target: ModelId::from("judge"),
1234                efficient_target: ModelId::from("e"),
1235                capable_target: ModelId::from("c"),
1236                config: test_config(base_threshold),
1237            })?;
1238        }
1239        Ok(())
1240    }
1241
1242    #[test]
1243    fn an_unusable_verdict_is_ambiguous() -> Result<()> {
1244        let policy = policy();
1245        let inconsistent_rule = TaskClassifierVerdict {
1246            capability_boundary: "uncertain".to_string(),
1247            ..verdict(1.0, "supported", "SUP-1")
1248        };
1249        let empty_crux = TaskClassifierVerdict {
1250            crux: "  ".to_string(),
1251            ..verdict(1.0, "supported", "SUP-1")
1252        };
1253        let unusable = [
1254            Some(verdict(1.1, "supported", "SUP-1")),
1255            Some(inconsistent_rule),
1256            Some(empty_crux),
1257            None,
1258        ];
1259        for verdict in unusable {
1260            let classification = policy.to_classification(verdict.as_ref());
1261            assert!(matches!(classification, Classification::Ambiguous(_)));
1262            assert!(classification.argmax(false)?.is_none());
1263            assert!(classification.argmax(true)?.is_none());
1264        }
1265        Ok(())
1266    }
1267
1268    #[test]
1269    fn capability_boundaries_apply_monotonic_threshold_steps() -> Result<()> {
1270        let policy = TaskClassifierPolicy::new(
1271            "efficient",
1272            "capable",
1273            &TaskClassifierConfig {
1274                threshold_step: 0.1,
1275                ..test_config(0.4)
1276            },
1277        );
1278
1279        assert_eq!(
1280            selected(&policy, Some(&verdict(0.4, "supported", "SUP-2")))?,
1281            "efficient"
1282        );
1283        assert_eq!(
1284            selected(&policy, Some(&verdict(0.49, "uncertain", "UNC-1")))?,
1285            "capable"
1286        );
1287        assert_eq!(
1288            selected(&policy, Some(&verdict(0.5, "uncertain", "UNC-1")))?,
1289            "efficient"
1290        );
1291        assert_eq!(
1292            selected(&policy, Some(&verdict(0.5, "unmatched", "none")))?,
1293            "efficient"
1294        );
1295        assert_eq!(
1296            selected(&policy, Some(&verdict(0.59, "unsupported", "LIM-1")))?,
1297            "capable"
1298        );
1299        assert_eq!(
1300            selected(&policy, Some(&verdict(0.6, "unsupported", "LIM-1")))?,
1301            "efficient"
1302        );
1303        Ok(())
1304    }
1305
1306    /// The text of each message a judge with `recent_turn_window` would be sent.
1307    /// The no-window case is covered by `capability_judge_builds_a_structured_request`.
1308    fn capability_judge(recent_turn_window: Option<usize>) -> Result<CapabilityJudge> {
1309        Ok(StructuredJudge::new(
1310            TaskInput { recent_turn_window },
1311            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?,
1312            SerdeDecoder::new(),
1313            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1314        ))
1315    }
1316
1317    fn judged_contents(recent_turn_window: usize) -> Result<Vec<String>> {
1318        let judge = capability_judge(Some(recent_turn_window))?;
1319        let request = Request {
1320            llm_request: LlmRequest {
1321                messages: vec![
1322                    Message::text(Role::System, "client instructions"),
1323                    Message::text(Role::User, "initial task"),
1324                    Message::text(Role::Assistant, "old response"),
1325                    Message::text(Role::User, "old follow-up"),
1326                    Message::text(Role::Assistant, "recent 1"),
1327                    Message::text(Role::User, "recent 2"),
1328                ],
1329                ..LlmRequest::default()
1330            },
1331            raw_request: None,
1332            metadata: None,
1333        };
1334        Ok(judge
1335            .build_request(&State::default(), &request)
1336            .llm_request
1337            .messages
1338            .iter()
1339            .filter_map(|message| message.text_content("\n"))
1340            .collect())
1341    }
1342
1343    #[test]
1344    fn a_window_widens_the_judge_to_the_surrounding_conversation() -> Result<()> {
1345        // Client instructions and the opening task, plus the last two turns.
1346        let contents = judged_contents(2)?;
1347        assert!(contents.contains(&"client instructions".to_string()));
1348        assert!(contents.contains(&"initial task".to_string()));
1349        assert!(contents.contains(&"recent 1".to_string()));
1350        assert!(contents.contains(&"recent 2".to_string()));
1351        assert!(!contents.contains(&"old response".to_string()));
1352        Ok(())
1353    }
1354
1355    #[test]
1356    fn a_zero_window_keeps_only_the_instructions_and_the_task() -> Result<()> {
1357        let contents = judged_contents(0)?;
1358        assert!(contents.contains(&"client instructions".to_string()));
1359        assert!(contents.contains(&"initial task".to_string()));
1360        assert!(!contents.contains(&"recent 2".to_string()));
1361        Ok(())
1362    }
1363
1364    fn tool_call(id: &str) -> Message {
1365        Message {
1366            role: Role::Assistant,
1367            content: vec![ContentBlock::ToolCall(ToolCall {
1368                id: id.to_string(),
1369                name: "search".to_string(),
1370                arguments: Value::Null,
1371            })],
1372        }
1373    }
1374
1375    fn tool_result(id: &str) -> Message {
1376        Message {
1377            role: Role::Tool,
1378            content: vec![ContentBlock::ToolResult(ToolResult {
1379                tool_call_id: id.to_string(),
1380                content: vec![ContentBlock::Text {
1381                    text: "tool output".to_string(),
1382                }],
1383                is_error: None,
1384            })],
1385        }
1386    }
1387
1388    /// A count-based window can begin on a tool result, which leaves the call that
1389    /// introduced its id outside the window and the classifier history invalid.
1390    #[test]
1391    fn trimming_keeps_the_call_that_introduced_a_kept_tool_result() {
1392        let messages = vec![
1393            Message::text(Role::System, "client instructions"),
1394            Message::text(Role::User, "initial task"),
1395            Message::text(Role::Assistant, "old response"),
1396            tool_call("call-1"),
1397            tool_result("call-1"),
1398            Message::text(Role::Assistant, "recent 1"),
1399            Message::text(Role::User, "recent 2"),
1400            Message::text(Role::Assistant, "recent 3"),
1401            Message::text(Role::User, "recent 4"),
1402        ];
1403
1404        // The five-message tail begins exactly on the tool result.
1405        let kept = trim_messages(&messages, 5);
1406
1407        assert_eq!(
1408            kept,
1409            vec![
1410                Message::text(Role::System, "client instructions"),
1411                Message::text(Role::User, "initial task"),
1412                tool_call("call-1"),
1413                tool_result("call-1"),
1414                Message::text(Role::Assistant, "recent 1"),
1415                Message::text(Role::User, "recent 2"),
1416                Message::text(Role::Assistant, "recent 3"),
1417                Message::text(Role::User, "recent 4"),
1418            ]
1419        );
1420    }
1421
1422    /// Ids repeat across a conversation, so a later call must not stand in for the one that
1423    /// answers an earlier result.
1424    #[test]
1425    fn trimming_pairs_a_repeated_id_with_the_call_that_precedes_it() {
1426        let messages = vec![
1427            Message::text(Role::System, "client instructions"),
1428            Message::text(Role::User, "initial task"),
1429            tool_call("x"),
1430            tool_result("x"),
1431            Message::text(Role::Assistant, "later"),
1432            tool_call("x"),
1433            tool_result("x"),
1434        ];
1435
1436        // The four-message tail begins on the first result, whose own call sits one earlier.
1437        let kept = trim_messages(&messages, 4);
1438
1439        assert_eq!(
1440            kept,
1441            vec![
1442                Message::text(Role::System, "client instructions"),
1443                Message::text(Role::User, "initial task"),
1444                tool_call("x"),
1445                tool_result("x"),
1446                Message::text(Role::Assistant, "later"),
1447                tool_call("x"),
1448                tool_result("x"),
1449            ]
1450        );
1451    }
1452
1453    /// A result whose call precedes the opening task can never be paired, because trimming
1454    /// never reaches behind the task. The window must not widen hunting for it.
1455    #[test]
1456    fn trimming_keeps_the_counted_window_when_a_result_cannot_be_paired() {
1457        let messages = vec![
1458            Message::text(Role::System, "client instructions"),
1459            tool_call("orphan"),
1460            Message::text(Role::User, "initial task"),
1461            Message::text(Role::Assistant, "old response"),
1462            tool_result("orphan"),
1463            Message::text(Role::Assistant, "recent 1"),
1464            Message::text(Role::User, "recent 2"),
1465        ];
1466
1467        let kept = trim_messages(&messages, 3);
1468
1469        assert_eq!(
1470            kept,
1471            vec![
1472                Message::text(Role::System, "client instructions"),
1473                Message::text(Role::User, "initial task"),
1474                tool_result("orphan"),
1475                Message::text(Role::Assistant, "recent 1"),
1476                Message::text(Role::User, "recent 2"),
1477            ]
1478        );
1479    }
1480
1481    #[test]
1482    fn a_window_restates_the_routing_instruction_last() -> Result<()> {
1483        let contents = judged_contents(2)?;
1484
1485        // Last, not merely present: the instruction only works in final position,
1486        // after the conversation content it is meant to outrank.
1487        assert_eq!(
1488            contents.last().map(String::as_str),
1489            Some(TRAILING_ROUTING_INSTRUCTION)
1490        );
1491        Ok(())
1492    }
1493
1494    #[test]
1495    fn the_default_path_is_left_unchanged() -> Result<()> {
1496        // No window means no conversation to be distracted by, so the default
1497        // request shape stays exactly as it was.
1498        let judge = capability_judge(None)?;
1499        let request = Request {
1500            llm_request: LlmRequest {
1501                messages: vec![Message::text(Role::User, "the task")],
1502                ..LlmRequest::default()
1503            },
1504            raw_request: None,
1505            metadata: None,
1506        };
1507
1508        let built = judge.build_request(&State::default(), &request);
1509
1510        // The task message alone; the rubric reaches the judge as an instruction block
1511        // rather than a message, so nothing was dropped by it not being counted here.
1512        assert_eq!(built.llm_request.messages.len(), 1);
1513        assert!(!built.llm_request.instructions.is_empty());
1514        assert!(
1515            !built
1516                .llm_request
1517                .messages
1518                .iter()
1519                .filter_map(|message| message.text_content("\n"))
1520                .any(|text| text.contains(TRAILING_ROUTING_INSTRUCTION))
1521        );
1522        Ok(())
1523    }
1524
1525    #[test]
1526    fn capability_judge_builds_a_structured_request() -> Result<()> {
1527        let judge = capability_judge(None)?;
1528        let request = Request {
1529            llm_request: LlmRequest {
1530                model: Some("inbound".to_string()),
1531                messages: vec![
1532                    Message::text(Role::System, "client instructions"),
1533                    Message::text(Role::Developer, "client developer instructions"),
1534                    Message::text(Role::User, "initial task"),
1535                    Message::text(Role::Assistant, "old response"),
1536                    Message::text(Role::User, "old follow-up"),
1537                    Message::text(Role::Assistant, "recent 1"),
1538                    Message::text(Role::User, "recent 2"),
1539                    Message::text(Role::Assistant, "recent 3"),
1540                    Message::text(Role::User, "recent 4"),
1541                    Message::text(Role::Assistant, "recent 5"),
1542                ],
1543                ..LlmRequest::default()
1544            },
1545            raw_request: None,
1546            metadata: None,
1547        };
1548        let judge_request = judge.build_request(&State::default(), &request);
1549
1550        assert_eq!(judge_request.llm_request.model, request.llm_request.model);
1551        assert_eq!(judge_request.llm_request.instructions.len(), 1);
1552        assert_eq!(judge_request.llm_request.instructions[0].role, Role::System);
1553        assert_eq!(
1554            judge_request.llm_request.instructions[0].content,
1555            InstructionBlock {
1556                role: Role::System,
1557                content: Message::text(Role::System, judge.contract().system_prompt()).content,
1558            }
1559            .content,
1560        );
1561        assert_eq!(judge_request.llm_request.messages.len(), 2);
1562        let contents = judge_request
1563            .llm_request
1564            .messages
1565            .iter()
1566            .filter_map(|message| message.text_content("\n"))
1567            .collect::<Vec<_>>();
1568        assert!(contents.contains(&"recent 4".to_string()));
1569        assert!(contents.contains(&"initial task".to_string()));
1570        assert!(!contents.contains(&"recent 5".to_string()));
1571        assert!(!contents.contains(&"client instructions".to_string()));
1572        assert_eq!(
1573            judge_request.llm_request.output.response_format,
1574            Some(judge.contract().response_format().clone())
1575        );
1576        assert_eq!(
1577            judge_request.llm_request.output.max_output_tokens,
1578            Some(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
1579        );
1580        Ok(())
1581    }
1582
1583    fn sample_value(spec: &Value) -> Value {
1584        if let Some(first) = spec
1585            .get("enum")
1586            .and_then(Value::as_array)
1587            .and_then(|values| values.first())
1588        {
1589            return first.clone();
1590        }
1591        match spec.get("type").and_then(Value::as_str) {
1592            Some("number") => serde_json::json!(0.5),
1593            Some("boolean") => serde_json::json!(false),
1594            _ => serde_json::json!("sample"),
1595        }
1596    }
1597
1598    fn schema_shaped_verdict(schema: &Value) -> Result<String> {
1599        let properties = schema
1600            .pointer("/json_schema/schema/properties")
1601            .and_then(Value::as_object)
1602            .ok_or_else(|| LibsyError::AlgorithmError {
1603                message: "packaged schema declares no properties".to_string(),
1604            })?;
1605        Ok(Value::Object(
1606            properties
1607                .iter()
1608                .map(|(name, spec)| (name.clone(), sample_value(spec)))
1609                .collect(),
1610        )
1611        .to_string())
1612    }
1613
1614    /// Built from the schema so a property added there fails here rather than silently
1615    /// rejecting every production verdict.
1616    #[test]
1617    fn every_schema_property_round_trips_through_the_judge_parser() -> Result<()> {
1618        let contract =
1619            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1620        let schema = contract.response_format();
1621        let reply = schema_shaped_verdict(schema)?;
1622        let judge: CapabilityJudge = StructuredJudge::new(
1623            TaskInput {
1624                recent_turn_window: None,
1625            },
1626            contract,
1627            SerdeDecoder::new(),
1628            JudgeRuntimeConfig::new(DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?,
1629        );
1630
1631        let verdict = judge.parse(&text_response(None, reply))?;
1632
1633        assert!(verdict.is_valid());
1634        assert!((0.0..=1.0).contains(&verdict.p_solve));
1635        Ok(())
1636    }
1637
1638    #[test]
1639    fn packaged_prompt_keeps_the_schema_in_the_structured_request() -> Result<()> {
1640        let contract =
1641            LlmTaskClassifier::load_capability_contract(&ClassifierContractConfig::default())?;
1642        let prompt = contract.system_prompt();
1643        let schema_name = contract
1644            .response_format()
1645            .pointer("/json_schema/name")
1646            .and_then(Value::as_str)
1647            .ok_or_else(|| LibsyError::AlgorithmError {
1648                message: "packaged response schema has no name".to_string(),
1649            })?;
1650        assert_eq!(schema_name, "CapabilityClassifierDecision");
1651        assert!(prompt.contains("SUP-1 [supported]"));
1652        assert!(prompt.contains("SUP-5 [supported]"));
1653        assert!(!prompt.contains("{{RESPONSE_SCHEMA}}"));
1654        assert!(!prompt.contains("\"type\": \"object\""));
1655        assert!(!prompt.contains("\"json_schema\""));
1656        assert!(!prompt.contains(schema_name));
1657        let rule_values = contract
1658            .response_format()
1659            .pointer("/json_schema/schema/properties/primary_rule/enum")
1660            .and_then(Value::as_array)
1661            .ok_or_else(|| LibsyError::AlgorithmError {
1662                message: "rendered response schema has no primary rule enum".to_string(),
1663            })?;
1664        assert!(
1665            rule_values
1666                .iter()
1667                .any(|value| value.as_str() == Some("SUP-1"))
1668        );
1669        assert!(
1670            rule_values
1671                .iter()
1672                .any(|value| value.as_str() == Some("none"))
1673        );
1674        Ok(())
1675    }
1676}