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.call_purpose = crate::context_plan::ContextCallPurpose::Classification;
70                llm_args.prompt = Some(if attempt == 0 {
71                    tool_prompt.clone()
72                } else {
73                    format!(
74                        "{tool_prompt}\n\nYour previous response could not be parsed. Please respond with ONLY the requested format."
75                    )
76                });
77
78                let result = dispatch_llm(llm_args, ctx).await;
79                if let Value::Err(e) = result {
80                    return Err(e);
81                }
82
83                match parse_classify_result(&result, is_binary, &categories) {
84                    Ok(v) => return Ok(v),
85                    Err(_) if attempt < parse_retry_count => continue,
86                    Err(e) => {
87                        return if is_binary {
88                            crate::notify!(warn, "llm.classify: parse failed, returning false");
89                            Ok(Value::Bool(false))
90                        } else {
91                            Err(RuntimeError::ToolFailed(format!("llm.classify: {e}")))
92                        };
93                    }
94                }
95            }
96
97            // Compiler can't prove loop executes at least once
98            Ok(Value::Bool(false))
99        })
100    }
101}
102
103fn extract_retry(args: &ToolArgs) -> u32 {
104    args.named("retry")
105        .and_then(|v| match v {
106            Value::Int(n) if *n >= 0 => Some(*n as u32),
107            _ => None,
108        })
109        .unwrap_or(0)
110}
111
112fn parse_classify_result(
113    result: &Value,
114    is_binary: bool,
115    categories: &[String],
116) -> Result<Value, String> {
117    // Handle pre-parsed types from assistant_message_to_value
118    match result {
119        Value::Bool(b) if is_binary => return Ok(Value::Bool(*b)),
120        Value::Int(n) if is_binary => return Ok(Value::Bool(*n != 0)),
121        Value::Float(f) if is_binary => return Ok(Value::Bool(*f != 0.0)),
122        _ => {}
123    }
124
125    // Extract text
126    let text = match llm_result_to_text(result) {
127        Some(t) => t,
128        None => return Err("could not extract text from LLM response".into()),
129    };
130
131    if is_binary {
132        match parse_bool_from_text(&text) {
133            Some(b) => Ok(Value::Bool(b)),
134            None => Err(format!("could not parse as yes/no: {text}")),
135        }
136    } else {
137        match parse_category_from_text(&text, categories) {
138            Some(cat) => Ok(Value::Str(cat)),
139            None => Err(format!("could not match any category: {text}")),
140        }
141    }
142}
atman_runtime/tools/llm_classify.rs

atman_runtime/tools/
llm_classify.rs
