1use serde::Deserialize;
11use switchyard_protocol::{ContentBlock, Message, ModelId, Role};
12
13use super::classifier_contract::{ClassifierContract, ClassifierContractConfig};
14use super::llm_judge::{
15 ClassifierInput, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig, SerdeDecoder,
16 StructuredJudge,
17};
18use crate::core::classifier::{Classification, Score};
19use crate::core::state::State;
20use crate::{LibsyError, Result};
21use switchyard_protocol::Request;
22
23const PROMPT_TEMPLATE: &str = include_str!("../../prompts/escalation/prompt.md");
24const SCHEMA_TEMPLATE: &str = include_str!("../../prompts/escalation/schema.json");
25
26const TRIM_MARKER: &str = " ...[trimmed] ";
28
29const TRUNCATION_SUFFIX: &str = "...<truncated>";
31
32const SYSTEM_CHARS: usize = 1_000;
35
36const FIRST_USER_CHARS: usize = 2_000;
38
39const MAX_REQUEST_CHARS: usize = 18_000;
41
42#[derive(Clone, Debug, Deserialize)]
47#[serde(default, deny_unknown_fields)]
48pub struct EscalationJudgeConfig {
49 pub confirmations: u32,
54 pub recent_turn_window: usize,
56 pub window_message_chars: usize,
58}
59
60impl EscalationJudgeConfig {
61 fn validate(&self) -> Result<()> {
63 let reject = |message: String| Err(LibsyError::AlgorithmError { message });
64 if self.confirmations == 0 {
65 return reject("confirmations must be at least 1".to_string());
66 }
67 if self.recent_turn_window == 0 {
68 return reject("recent_turn_window must be at least 1".to_string());
69 }
70 if self.window_message_chars < 50 {
71 return reject(format!(
72 "window_message_chars must be at least 50, got {}",
73 self.window_message_chars
74 ));
75 }
76 Ok(())
77 }
78}
79
80impl Default for EscalationJudgeConfig {
81 fn default() -> Self {
82 Self {
83 confirmations: 2,
84 recent_turn_window: 28,
85 window_message_chars: 500,
86 }
87 }
88}
89
90#[derive(Deserialize)]
94pub(crate) struct EscalationVerdict {
95 escalate: bool,
96}
97
98pub(crate) struct EscalationInput {
100 config: EscalationJudgeConfig,
101}
102
103impl ClassifierInput for EscalationInput {
104 fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
105 let messages = &request.llm_request.messages;
106 let summary = summarize_for_judge(messages, conversation_turn(request), &self.config);
107 vec![Message::text(Role::User, summary)]
108 }
109}
110
111pub(crate) type EscalationJudge = StructuredJudge<EscalationInput, SerdeDecoder<EscalationVerdict>>;
113
114pub(crate) struct EscalationPolicy {
120 capable: ModelId,
121 efficient: ModelId,
122}
123
124impl JudgePolicy for EscalationPolicy {
125 type Verdict = EscalationVerdict;
126
127 fn to_classification(&self, verdict: Option<&EscalationVerdict>) -> Classification {
128 match verdict {
129 Some(verdict) if verdict.escalate => Classification::Scores(vec![Score {
130 target: self.capable.clone(),
131 confidence: 1.0,
132 }]),
133 Some(_) => Classification::Scores(vec![Score {
134 target: self.efficient.clone(),
135 confidence: 1.0,
136 }]),
137 None => Classification::Ambiguous(Vec::new()),
138 }
139 }
140}
141
142pub(crate) fn build_judge(
147 judge_target: ModelId,
148 capable: ModelId,
149 efficient: ModelId,
150 contract_config: &ClassifierContractConfig,
151 config: EscalationJudgeConfig,
152 max_output_tokens: u64,
153) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
154 config.validate()?;
155 let contract =
156 ClassifierContract::from_config(contract_config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)?;
157 Ok(JudgeClassifier::new(
158 StructuredJudge::new(
159 EscalationInput { config },
160 contract,
161 SerdeDecoder::new(),
162 JudgeRuntimeConfig::new(max_output_tokens)?,
163 ),
164 judge_target,
165 EscalationPolicy { capable, efficient },
166 ))
167}
168
169pub(crate) fn conversation_turn(request: &Request) -> usize {
178 request
179 .llm_request
180 .messages
181 .iter()
182 .filter(|message| message.role == Role::Assistant)
183 .count()
184}
185
186fn message_text(message: &Message) -> String {
192 let mut parts = Vec::new();
193 collect_text(&message.content, &mut parts);
194 parts.join(" ")
195}
196
197fn collect_text(content: &[ContentBlock], parts: &mut Vec<String>) {
199 for block in content {
200 match block {
201 ContentBlock::Text { text } | ContentBlock::Refusal { text } => {
202 parts.push(text.clone());
203 }
204 ContentBlock::ToolCall(call) => {
205 parts.push(format!("tool_call {}({})", call.name, call.arguments));
206 }
207 ContentBlock::ToolResult(result) => collect_text(&result.content, parts),
208 _ => {}
209 }
210 }
211}
212
213fn truncate_middle(text: &str, limit: usize) -> String {
218 let chars: Vec<char> = text.chars().collect();
219 if chars.len() <= limit {
220 return text.to_string();
221 }
222 let keep = limit
223 .saturating_sub(TRIM_MARKER.chars().count())
224 .max(20)
225 .min(chars.len());
226 let head = keep * 2 / 3;
227 let tail = keep - head;
228 let mut out: String = chars[..head].iter().collect();
229 out.push_str(TRIM_MARKER);
230 out.extend(chars[chars.len() - tail..].iter());
231 out
232}
233
234fn summarize_for_judge(
244 messages: &[Message],
245 turn: usize,
246 config: &EscalationJudgeConfig,
247) -> String {
248 let mut anchors: Vec<String> = Vec::new();
249 let mut window: Vec<String> = Vec::new();
250 let mut first_user_seen = false;
251
252 for message in messages {
253 let text = message_text(message);
254 match message.role {
255 Role::System | Role::Developer => anchors.push(format!(
256 "[{}] {}",
257 role_label(message.role),
258 truncate_middle(&text, SYSTEM_CHARS)
259 )),
260 Role::User if !first_user_seen => {
261 first_user_seen = true;
262 anchors.push(format!(
263 "[user (task)] {}",
264 truncate_middle(&text, FIRST_USER_CHARS)
265 ));
266 }
267 role => window.push(format!(
268 "[{}] {}",
269 role_label(role),
270 truncate_middle(&text, config.window_message_chars)
271 )),
272 }
273 }
274
275 if window.len() > config.recent_turn_window {
276 window.drain(..window.len() - config.recent_turn_window);
277 }
278
279 let assemble = |window: &[String]| {
280 let header = format!(
281 "Conversation turn {turn}; showing the last {} of {} messages after the task framing.",
282 window.len(),
283 messages.len(),
284 );
285 std::iter::once(header)
286 .chain(anchors.iter().cloned())
287 .chain(window.iter().cloned())
288 .collect::<Vec<_>>()
289 .join("\n")
290 };
291
292 let mut text = assemble(&window);
293 while text.chars().count() > MAX_REQUEST_CHARS && !window.is_empty() {
294 window.remove(0);
295 text = assemble(&window);
296 }
297 if text.chars().count() > MAX_REQUEST_CHARS {
298 let keep = MAX_REQUEST_CHARS.saturating_sub(TRUNCATION_SUFFIX.chars().count() + 1);
299 text = text.chars().take(keep).collect::<String>() + TRUNCATION_SUFFIX;
300 }
301 text
302}
303
304fn role_label(role: Role) -> &'static str {
306 match role {
307 Role::System => "system",
308 Role::Developer => "developer",
309 Role::User => "user",
310 Role::Assistant => "assistant",
311 Role::Tool => "tool",
312 }
313}
314
315#[cfg(test)]
320pub(crate) fn request_at_turn(session_id: Option<&str>, turn: usize) -> Request {
321 use switchyard_protocol::{LlmRequest, Metadata};
322
323 let mut messages = vec![Message::text(Role::User, "What is 2+2?")];
324 for attempt in 1..turn {
325 messages.push(Message::text(Role::Assistant, format!("attempt {attempt}")));
326 messages.push(Message::text(Role::User, format!("still wrong {attempt}")));
327 }
328 Request {
329 llm_request: LlmRequest {
330 model: Some("auto".to_string()),
331 messages,
332 ..LlmRequest::default()
333 },
334 raw_request: None,
335 metadata: session_id.map(|id| Metadata {
336 session_id: Some(id.to_string()),
337 ..Metadata::default()
338 }),
339 }
340}
341
342#[cfg(test)]
343mod tests {
344 use serde_json::json;
345 use switchyard_protocol::{ContentBlock, Message, Role, ToolCall, ToolResult};
346
347 use super::*;
348 use crate::algorithms::util::llm_judge::Judge;
349
350 fn escalation_judge(max_output_tokens: u64) -> Result<EscalationJudge> {
351 Ok(StructuredJudge::new(
352 EscalationInput {
353 config: EscalationJudgeConfig::default(),
354 },
355 ClassifierContract::from_config(
356 &ClassifierContractConfig::default(),
357 PROMPT_TEMPLATE,
358 SCHEMA_TEMPLATE,
359 )?,
360 SerdeDecoder::new(),
361 JudgeRuntimeConfig::new(max_output_tokens)?,
362 ))
363 }
364
365 #[test]
366 fn judge_request_is_rubric_plus_summary_under_a_completion_cap() -> Result<()> {
367 let judge = escalation_judge(super::super::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?;
368
369 let mut judged = request_at_turn(None, 4);
371 judged
372 .llm_request
373 .messages
374 .push(Message::text(Role::Assistant, "this turn's reply"));
375 let built = judge.build_request(&State::default(), &judged);
376
377 assert_eq!(built.llm_request.instructions.len(), 1);
379 assert_eq!(built.llm_request.instructions[0].role, Role::System);
380 assert_eq!(built.llm_request.messages.len(), 1);
381 assert_eq!(built.llm_request.messages[0].role, Role::User);
382 assert!(
383 built.llm_request.messages[0]
384 .text_content("")
385 .is_some_and(|text| text.contains("Conversation turn 4"))
386 );
387 assert_eq!(
389 built.llm_request.output.max_output_tokens,
390 Some(super::super::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
391 );
392 assert!(built.llm_request.output.response_format.is_some());
393 Ok(())
394 }
395
396 #[test]
397 fn judge_request_uses_the_configured_completion_cap() -> Result<()> {
398 let judge = escalation_judge(512)?;
399
400 let built = judge.build_request(&State::default(), &request_at_turn(None, 1));
401
402 assert_eq!(built.llm_request.output.max_output_tokens, Some(512));
403 Ok(())
404 }
405
406 #[test]
407 fn conversation_turn_counts_assistant_replies() {
408 for turn in [1, 5] {
411 let mut judged = request_at_turn(None, turn);
412 judged
413 .llm_request
414 .messages
415 .push(Message::text(Role::Assistant, "this turn's reply"));
416 assert_eq!(conversation_turn(&judged), turn);
417 }
418 }
419
420 #[test]
421 fn message_text_keeps_tool_calls_and_results() {
422 let call = Message {
423 role: Role::Assistant,
424 content: vec![
425 ContentBlock::Text {
426 text: "running it".to_string(),
427 },
428 ContentBlock::ToolCall(ToolCall {
429 id: "call-1".to_string(),
430 name: "bash".to_string(),
431 arguments: json!({"cmd": "ls"}),
432 }),
433 ],
434 };
435 let text = message_text(&call);
436 assert!(text.contains("running it"), "{text}");
437 assert!(text.contains(r#"tool_call bash({"cmd":"ls"})"#), "{text}");
438
439 let result = Message {
440 role: Role::Tool,
441 content: vec![ContentBlock::ToolResult(ToolResult {
442 tool_call_id: "call-1".to_string(),
443 content: vec![ContentBlock::Text {
444 text: "no such file".to_string(),
445 }],
446 is_error: Some(true),
447 })],
448 };
449 assert_eq!(message_text(&result), "no such file");
450 }
451
452 #[test]
453 fn truncate_middle_keeps_head_and_tail() {
454 let text = "a".repeat(40) + &"z".repeat(40);
455 let trimmed = truncate_middle(&text, 50);
456 assert!(trimmed.chars().count() <= 50, "{trimmed}");
457 assert!(trimmed.starts_with('a'));
458 assert!(trimmed.ends_with('z'));
459 assert!(trimmed.contains("[trimmed]"));
460
461 assert_eq!(truncate_middle("short", 50), "short");
463 }
464
465 #[test]
466 fn summary_keeps_anchors_and_the_recent_window() {
467 let mut messages = vec![
468 Message::text(Role::System, "you are a coding agent"),
469 Message::text(Role::User, "fix the failing test"),
470 ];
471 for i in 0..10 {
472 messages.push(Message::text(Role::Assistant, format!("step {i}")));
473 }
474 let config = EscalationJudgeConfig {
475 recent_turn_window: 3,
476 ..EscalationJudgeConfig::default()
477 };
478
479 let summary = summarize_for_judge(&messages, 11, &config);
480
481 assert!(
482 summary.contains("[system] you are a coding agent"),
483 "{summary}"
484 );
485 assert!(
486 summary.contains("[user (task)] fix the failing test"),
487 "{summary}"
488 );
489 assert!(summary.contains("Conversation turn 11; showing the last 3 of 12 messages"));
490 assert!(summary.contains("step 9"), "{summary}");
492 assert!(summary.contains("step 7"), "{summary}");
493 assert!(!summary.contains("step 6"), "{summary}");
494 }
495
496 #[test]
497 fn summary_drops_oldest_window_lines_under_the_char_cap() {
498 let mut messages = vec![
502 Message::text(Role::System, "framing"),
503 Message::text(Role::User, "task"),
504 ];
505 for i in 0..20 {
506 messages.push(Message::text(
507 Role::Assistant,
508 format!("{i} {}", "x".repeat(2_000)),
509 ));
510 }
511 let config = EscalationJudgeConfig {
512 window_message_chars: 2_000,
513 ..EscalationJudgeConfig::default()
514 };
515
516 let summary = summarize_for_judge(&messages, 21, &config);
517
518 assert!(
519 summary.chars().count() <= MAX_REQUEST_CHARS,
520 "{}",
521 summary.chars().count()
522 );
523 assert!(summary.contains("[system] framing"), "{summary}");
525 assert!(summary.contains("[user (task)] task"), "{summary}");
526 assert!(summary.contains("19 xxx"), "{summary}");
527 assert!(!summary.contains("0 xxx"), "{summary}");
528 }
529}