atman_runtime/tools/
llm_classify.rs

1use crate::error::RuntimeError;
2use crate::eval::llm_args::parse_llm_args_from_toolargs;
3use crate::eval::llm_dispatch::dispatch_llm;
4use crate::eval::llm_parse::{llm_result_to_text, parse_bool_from_text, parse_category_from_text};
5use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
6use crate::value::Value;
7
8pub struct LlmClassifyTool;
9
10impl Tool for LlmClassifyTool {
11    fn name(&self) -> &str {
12        "llm.classify"
13    }
14
15    fn tier(&self) -> Tier {
16        Tier::Zero
17    }
18
19    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
20        Box::pin(async move {
21            let Some(registry) = ctx.registry.as_deref() else {
22                return Err(RuntimeError::ToolFailed(
23                    "llm.classify: no tool registry available".into(),
24                ));
25            };
26
27            let prompt = args
28                .named("prompt")
29                .and_then(|v| match v {
30                    Value::Str(s) => Some(s.clone()),
31                    _ => None,
32                })
33                .ok_or_else(|| RuntimeError::MissingArg("llm.classify: prompt".into()))?;
34
35            if prompt.trim().is_empty() {
36                return Err(RuntimeError::ToolFailed(
37                    "llm.classify: prompt is required".into(),
38                ));
39            }
40
41            // Parse categories (default: binary yes/no)
42            let categories: Vec<String> = match args.named("categories") {
43                Some(Value::List(items)) => items
44                    .iter()
45                    .filter_map(|v| match v {
46                        Value::Str(s) => Some(s.clone()),
47                        _ => None,
48                    })
49                    .collect(),
50                _ => Vec::new(),
51            };
52            let is_binary = categories.is_empty();
53
54            // Construct tool-specific prompt
55            let tool_prompt = if is_binary {
56                format!("Answer with exactly one word: \"yes\" or \"no\".\n\n{prompt}")
57            } else {
58                let labels = categories.join(", ");
59                format!("Answer with exactly one of these labels: {labels}.\n\n{prompt}")
60            };
61
62            let parse_retry_count = extract_retry(&args);
63
64            for attempt in 0..=parse_retry_count {
65                let mut llm_args = parse_llm_args_from_toolargs(&args, registry)?;
66                llm_args.retry_count = 0;
67                llm_args.messages_override = None;
68                llm_args.fallback_value = None;
69                llm_args.prompt = Some(if attempt == 0 {
70                    tool_prompt.clone()
71                } else {
72                    format!(
73                        "{tool_prompt}\n\nYour previous response could not be parsed. Please respond with ONLY the requested format."
74                    )
75                });
76
77                let result = dispatch_llm(llm_args, ctx).await;
78                if let Value::Err(e) = result {
79                    return Err(e);
80                }
81
82                match parse_classify_result(&result, is_binary, &categories) {
83                    Ok(v) => return Ok(v),
84                    Err(_) if attempt < parse_retry_count => continue,
85                    Err(e) => {
86                        return if is_binary {
87                            crate::notify!(warn, "llm.classify: parse failed, returning false");
88                            Ok(Value::Bool(false))
89                        } else {
90                            Err(RuntimeError::ToolFailed(format!("llm.classify: {e}")))
91                        };
92                    }
93                }
94            }
95
96            // Compiler can't prove loop executes at least once
97            Ok(Value::Bool(false))
98        })
99    }
100}
101
102fn extract_retry(args: &ToolArgs) -> u32 {
103    args.named("retry")
104        .and_then(|v| match v {
105            Value::Int(n) if *n >= 0 => Some(*n as u32),
106            _ => None,
107        })
108        .unwrap_or(0)
109}
110
111fn parse_classify_result(
112    result: &Value,
113    is_binary: bool,
114    categories: &[String],
115) -> Result<Value, String> {
116    // Handle pre-parsed types from assistant_message_to_value
117    match result {
118        Value::Bool(b) if is_binary => return Ok(Value::Bool(*b)),
119        Value::Int(n) if is_binary => return Ok(Value::Bool(*n != 0)),
120        Value::Float(f) if is_binary => return Ok(Value::Bool(*f != 0.0)),
121        _ => {}
122    }
123
124    // Extract text
125    let text = match llm_result_to_text(result) {
126        Some(t) => t,
127        None => return Err("could not extract text from LLM response".into()),
128    };
129
130    if is_binary {
131        match parse_bool_from_text(&text) {
132            Some(b) => Ok(Value::Bool(b)),
133            None => Err(format!("could not parse as yes/no: {text}")),
134        }
135    } else {
136        match parse_category_from_text(&text, categories) {
137            Some(cat) => Ok(Value::Str(cat)),
138            None => Err(format!("could not match any category: {text}")),
139        }
140    }
141}
atman_runtime/tools/llm_classify.rs

atman_runtime/tools/
llm_classify.rs
