1use serde::Deserialize;
12use switchyard_protocol::{ContentBlock, Message, Role};
13
14use super::classifier_contract::{ClassifierContract, ClassifierContractConfig};
15use super::llm_judge::{
16 ClassifierInput, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig, SerdeDecoder,
17 StructuredJudge,
18};
19use crate::core::algorithm::LlmTarget;
20use crate::core::classifier::{Classification, Score};
21use crate::core::state::State;
22use crate::{LibsyError, Result};
23use switchyard_protocol::Request;
24
25const PROMPT_TEMPLATE: &str = include_str!("../../prompts/escalation/prompt.md");
26const RECOVERY_PROMPT_TEMPLATE: &str = include_str!("../../prompts/escalation/recovery_prompt.md");
27const SCHEMA_TEMPLATE: &str = include_str!("../../prompts/escalation/schema.json");
28
29const TRIM_MARKER: &str = " ...[trimmed] ";
31
32const TRUNCATION_SUFFIX: &str = "...<truncated>";
34
35const SYSTEM_CHARS: usize = 1_000;
38
39const FIRST_USER_CHARS: usize = 2_000;
41
42const MAX_REQUEST_CHARS: usize = 18_000;
44
45#[derive(Clone, Debug, Deserialize)]
50#[serde(default, deny_unknown_fields)]
51pub struct EscalationJudgeConfig {
52 pub confirmations: u32,
57 pub recovery_confirmations: u32,
62 pub recent_turn_window: usize,
64 pub window_message_chars: usize,
66}
67
68impl EscalationJudgeConfig {
69 fn validate(&self) -> Result<()> {
71 let reject = |message: String| Err(LibsyError::AlgorithmError { message });
72 if self.confirmations == 0 {
73 return reject("confirmations must be at least 1".to_string());
74 }
75 if self.recent_turn_window == 0 {
76 return reject("recent_turn_window must be at least 1".to_string());
77 }
78 if self.window_message_chars < 50 {
79 return reject(format!(
80 "window_message_chars must be at least 50, got {}",
81 self.window_message_chars
82 ));
83 }
84 Ok(())
85 }
86}
87
88impl Default for EscalationJudgeConfig {
89 fn default() -> Self {
90 Self {
91 confirmations: 2,
92 recovery_confirmations: 0,
93 recent_turn_window: 28,
94 window_message_chars: 500,
95 }
96 }
97}
98
99#[derive(Deserialize)]
103pub(crate) struct EscalationVerdict {
104 escalate: bool,
105}
106
107pub(crate) struct EscalationInput {
109 config: EscalationJudgeConfig,
110}
111
112impl ClassifierInput for EscalationInput {
113 fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
114 let messages = &request.llm_request.messages;
115 let summary = summarize_for_judge(messages, conversation_turn(request), &self.config);
116 vec![Message::text(Role::User, summary)]
117 }
118}
119
120pub(crate) type EscalationJudge = StructuredJudge<EscalationInput, SerdeDecoder<EscalationVerdict>>;
122
123pub(crate) struct EscalationPolicy {
129 capable: String,
130 efficient: String,
131}
132
133impl JudgePolicy for EscalationPolicy {
134 type Verdict = EscalationVerdict;
135
136 fn to_classification(&self, verdict: Option<&EscalationVerdict>) -> Classification {
137 match verdict {
138 Some(verdict) if verdict.escalate => Classification::Scores(vec![Score {
139 target: self.capable.clone(),
140 confidence: 1.0,
141 }]),
142 Some(_) => Classification::Scores(vec![Score {
143 target: self.efficient.clone(),
144 confidence: 1.0,
145 }]),
146 None => Classification::Ambiguous(Vec::new()),
147 }
148 }
149}
150
151pub(crate) fn build_judge(
156 judge_target: LlmTarget,
157 capable: String,
158 efficient: String,
159 contract_config: &ClassifierContractConfig,
160 config: EscalationJudgeConfig,
161 max_output_tokens: u64,
162) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
163 let contract =
164 ClassifierContract::from_config(contract_config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)?;
165 assemble_judge(
166 judge_target,
167 capable,
168 efficient,
169 contract,
170 config,
171 max_output_tokens,
172 )
173}
174
175pub(crate) fn build_recovery_judge(
184 judge_target: LlmTarget,
185 capable: String,
186 efficient: String,
187 config: EscalationJudgeConfig,
188 max_output_tokens: u64,
189) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
190 let contract = ClassifierContract::from_config(
191 &ClassifierContractConfig::default(),
192 RECOVERY_PROMPT_TEMPLATE,
193 SCHEMA_TEMPLATE,
194 )?;
195 assemble_judge(
196 judge_target,
197 capable,
198 efficient,
199 contract,
200 config,
201 max_output_tokens,
202 )
203}
204
205fn assemble_judge(
207 judge_target: LlmTarget,
208 capable: String,
209 efficient: String,
210 contract: ClassifierContract,
211 config: EscalationJudgeConfig,
212 max_output_tokens: u64,
213) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
214 config.validate()?;
215 Ok(JudgeClassifier::new(
216 StructuredJudge::new(
217 EscalationInput { config },
218 contract,
219 SerdeDecoder::new(),
220 JudgeRuntimeConfig::new(max_output_tokens)?,
221 ),
222 judge_target,
223 EscalationPolicy { capable, efficient },
224 ))
225}
226
227pub(crate) fn conversation_turn(request: &Request) -> usize {
236 request
237 .llm_request
238 .messages
239 .iter()
240 .filter(|message| message.role == Role::Assistant)
241 .count()
242}
243
244fn message_text(message: &Message) -> String {
250 let mut parts = Vec::new();
251 collect_text(&message.content, &mut parts);
252 parts.join(" ")
253}
254
255fn collect_text(content: &[ContentBlock], parts: &mut Vec<String>) {
257 for block in content {
258 match block {
259 ContentBlock::Text { text } | ContentBlock::Refusal { text } => {
260 parts.push(text.clone());
261 }
262 ContentBlock::ToolCall(call) => {
263 parts.push(format!("tool_call {}({})", call.name, call.arguments));
264 }
265 ContentBlock::ToolResult(result) => collect_text(&result.content, parts),
266 _ => {}
267 }
268 }
269}
270
271fn truncate_middle(text: &str, limit: usize) -> String {
276 let chars: Vec<char> = text.chars().collect();
277 if chars.len() <= limit {
278 return text.to_string();
279 }
280 let keep = limit
281 .saturating_sub(TRIM_MARKER.chars().count())
282 .max(20)
283 .min(chars.len());
284 let head = keep * 2 / 3;
285 let tail = keep - head;
286 let mut out: String = chars[..head].iter().collect();
287 out.push_str(TRIM_MARKER);
288 out.extend(chars[chars.len() - tail..].iter());
289 out
290}
291
292fn summarize_for_judge(
302 messages: &[Message],
303 turn: usize,
304 config: &EscalationJudgeConfig,
305) -> String {
306 let mut anchors: Vec<String> = Vec::new();
307 let mut window: Vec<String> = Vec::new();
308 let mut first_user_seen = false;
309
310 for message in messages {
311 let text = message_text(message);
312 match message.role {
313 Role::System | Role::Developer => anchors.push(format!(
314 "[{}] {}",
315 role_label(message.role),
316 truncate_middle(&text, SYSTEM_CHARS)
317 )),
318 Role::User if !first_user_seen => {
319 first_user_seen = true;
320 anchors.push(format!(
321 "[user (task)] {}",
322 truncate_middle(&text, FIRST_USER_CHARS)
323 ));
324 }
325 role => window.push(format!(
326 "[{}] {}",
327 role_label(role),
328 truncate_middle(&text, config.window_message_chars)
329 )),
330 }
331 }
332
333 if window.len() > config.recent_turn_window {
334 window.drain(..window.len() - config.recent_turn_window);
335 }
336
337 let assemble = |window: &[String]| {
338 let header = format!(
339 "Conversation turn {turn}; showing the last {} of {} messages after the task framing.",
340 window.len(),
341 messages.len(),
342 );
343 std::iter::once(header)
344 .chain(anchors.iter().cloned())
345 .chain(window.iter().cloned())
346 .collect::<Vec<_>>()
347 .join("\n")
348 };
349
350 let mut text = assemble(&window);
351 while text.chars().count() > MAX_REQUEST_CHARS && !window.is_empty() {
352 window.remove(0);
353 text = assemble(&window);
354 }
355 if text.chars().count() > MAX_REQUEST_CHARS {
356 let keep = MAX_REQUEST_CHARS.saturating_sub(TRUNCATION_SUFFIX.chars().count() + 1);
357 text = text.chars().take(keep).collect::<String>() + TRUNCATION_SUFFIX;
358 }
359 text
360}
361
362fn role_label(role: Role) -> &'static str {
364 match role {
365 Role::System => "system",
366 Role::Developer => "developer",
367 Role::User => "user",
368 Role::Assistant => "assistant",
369 Role::Tool => "tool",
370 }
371}
372
373#[cfg(test)]
378pub(crate) fn request_at_turn(session_id: Option<&str>, turn: usize) -> Request {
379 use switchyard_protocol::{LlmRequest, Metadata};
380
381 let mut messages = vec![Message::text(Role::User, "What is 2+2?")];
382 for attempt in 1..turn {
383 messages.push(Message::text(Role::Assistant, format!("attempt {attempt}")));
384 messages.push(Message::text(Role::User, format!("still wrong {attempt}")));
385 }
386 Request {
387 llm_request: LlmRequest {
388 model: Some("auto".to_string()),
389 messages,
390 ..LlmRequest::default()
391 },
392 raw_request: None,
393 metadata: session_id.map(|id| Metadata {
394 session_id: Some(id.to_string()),
395 ..Metadata::default()
396 }),
397 }
398}
399
400#[cfg(test)]
401mod tests {
402 use serde_json::json;
403 use switchyard_protocol::{ContentBlock, Message, Role, ToolCall, ToolResult};
404
405 use super::*;
406 use crate::algorithms::util::llm_judge::Judge;
407
408 fn escalation_judge(max_output_tokens: u64) -> Result<EscalationJudge> {
409 Ok(StructuredJudge::new(
410 EscalationInput {
411 config: EscalationJudgeConfig::default(),
412 },
413 ClassifierContract::from_config(
414 &ClassifierContractConfig::default(),
415 PROMPT_TEMPLATE,
416 SCHEMA_TEMPLATE,
417 )?,
418 SerdeDecoder::new(),
419 JudgeRuntimeConfig::new(max_output_tokens)?,
420 ))
421 }
422
423 #[test]
424 fn judge_request_is_rubric_plus_summary_under_a_completion_cap() -> Result<()> {
425 let judge = escalation_judge(super::super::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?;
426
427 let mut judged = request_at_turn(None, 4);
429 judged
430 .llm_request
431 .messages
432 .push(Message::text(Role::Assistant, "this turn's reply"));
433 let built = judge.build_request(&State::default(), &judged);
434
435 assert_eq!(built.llm_request.instructions.len(), 1);
437 assert_eq!(built.llm_request.instructions[0].role, Role::System);
438 assert_eq!(built.llm_request.messages.len(), 1);
439 assert_eq!(built.llm_request.messages[0].role, Role::User);
440 assert!(
441 built.llm_request.messages[0]
442 .text_content("")
443 .is_some_and(|text| text.contains("Conversation turn 4"))
444 );
445 assert_eq!(
447 built.llm_request.output.max_output_tokens,
448 Some(super::super::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
449 );
450 assert!(built.llm_request.output.response_format.is_some());
451 Ok(())
452 }
453
454 #[test]
455 fn judge_request_uses_the_configured_completion_cap() -> Result<()> {
456 let judge = escalation_judge(512)?;
457
458 let built = judge.build_request(&State::default(), &request_at_turn(None, 1));
459
460 assert_eq!(built.llm_request.output.max_output_tokens, Some(512));
461 Ok(())
462 }
463
464 #[test]
465 fn conversation_turn_counts_assistant_replies() {
466 for turn in [1, 5] {
469 let mut judged = request_at_turn(None, turn);
470 judged
471 .llm_request
472 .messages
473 .push(Message::text(Role::Assistant, "this turn's reply"));
474 assert_eq!(conversation_turn(&judged), turn);
475 }
476 }
477
478 #[test]
479 fn message_text_keeps_tool_calls_and_results() {
480 let call = Message {
481 role: Role::Assistant,
482 content: vec![
483 ContentBlock::Text {
484 text: "running it".to_string(),
485 },
486 ContentBlock::ToolCall(ToolCall {
487 id: "call-1".to_string(),
488 name: "bash".to_string(),
489 arguments: json!({"cmd": "ls"}),
490 }),
491 ],
492 };
493 let text = message_text(&call);
494 assert!(text.contains("running it"), "{text}");
495 assert!(text.contains(r#"tool_call bash({"cmd":"ls"})"#), "{text}");
496
497 let result = Message {
498 role: Role::Tool,
499 content: vec![ContentBlock::ToolResult(ToolResult {
500 tool_call_id: "call-1".to_string(),
501 content: vec![ContentBlock::Text {
502 text: "no such file".to_string(),
503 }],
504 is_error: Some(true),
505 })],
506 };
507 assert_eq!(message_text(&result), "no such file");
508 }
509
510 #[test]
511 fn truncate_middle_keeps_head_and_tail() {
512 let text = "a".repeat(40) + &"z".repeat(40);
513 let trimmed = truncate_middle(&text, 50);
514 assert!(trimmed.chars().count() <= 50, "{trimmed}");
515 assert!(trimmed.starts_with('a'));
516 assert!(trimmed.ends_with('z'));
517 assert!(trimmed.contains("[trimmed]"));
518
519 assert_eq!(truncate_middle("short", 50), "short");
521 }
522
523 #[test]
524 fn summary_keeps_anchors_and_the_recent_window() {
525 let mut messages = vec![
526 Message::text(Role::System, "you are a coding agent"),
527 Message::text(Role::User, "fix the failing test"),
528 ];
529 for i in 0..10 {
530 messages.push(Message::text(Role::Assistant, format!("step {i}")));
531 }
532 let config = EscalationJudgeConfig {
533 recent_turn_window: 3,
534 ..EscalationJudgeConfig::default()
535 };
536
537 let summary = summarize_for_judge(&messages, 11, &config);
538
539 assert!(
540 summary.contains("[system] you are a coding agent"),
541 "{summary}"
542 );
543 assert!(
544 summary.contains("[user (task)] fix the failing test"),
545 "{summary}"
546 );
547 assert!(summary.contains("Conversation turn 11; showing the last 3 of 12 messages"));
548 assert!(summary.contains("step 9"), "{summary}");
550 assert!(summary.contains("step 7"), "{summary}");
551 assert!(!summary.contains("step 6"), "{summary}");
552 }
553
554 #[test]
555 fn summary_drops_oldest_window_lines_under_the_char_cap() {
556 let mut messages = vec![
560 Message::text(Role::System, "framing"),
561 Message::text(Role::User, "task"),
562 ];
563 for i in 0..20 {
564 messages.push(Message::text(
565 Role::Assistant,
566 format!("{i} {}", "x".repeat(2_000)),
567 ));
568 }
569 let config = EscalationJudgeConfig {
570 window_message_chars: 2_000,
571 ..EscalationJudgeConfig::default()
572 };
573
574 let summary = summarize_for_judge(&messages, 21, &config);
575
576 assert!(
577 summary.chars().count() <= MAX_REQUEST_CHARS,
578 "{}",
579 summary.chars().count()
580 );
581 assert!(summary.contains("[system] framing"), "{summary}");
583 assert!(summary.contains("[user (task)] task"), "{summary}");
584 assert!(summary.contains("19 xxx"), "{summary}");
585 assert!(!summary.contains("0 xxx"), "{summary}");
586 }
587}