atman_runtime/tools/
llm_generate_branches.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_list_from_text};
5use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
6use crate::value::Value;
7
8pub struct LlmGenerateBranchesTool;
9
10impl Tool for LlmGenerateBranchesTool {
11    fn name(&self) -> &str {
12        "llm.generate_branches"
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.generate_branches: 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.generate_branches: prompt".into()))?;
34
35            if prompt.trim().is_empty() {
36                return Err(RuntimeError::ToolFailed(
37                    "llm.generate_branches: prompt is required".into(),
38                ));
39            }
40
41            // Parse optional count
42            let count = args.named("count").and_then(|v| match v {
43                Value::Int(n) if *n > 0 => Some(*n as usize),
44                _ => None,
45            });
46
47            // Construct prompt
48            let count_str = count
49                .map(|c| c.to_string())
50                .unwrap_or_else(|| "3-8".to_string());
51            let tool_prompt = format!(
52                "Decompose the following task into {count_str} independent parallel subtasks.\n\
53                 Each subtask should be a concise one-line description.\n\
54                 Respond as a JSON array of strings, e.g. [\"subtask 1\", \"subtask 2\", ...].\n\
55                 No markdown, no explanation, just the JSON array.\n\n{prompt}"
56            );
57
58            let parse_retry_count = extract_retry(&args);
59
60            for attempt in 0..=parse_retry_count {
61                let mut llm_args = parse_llm_args_from_toolargs(&args, registry)?;
62                llm_args.retry_count = 0;
63                llm_args.messages_override = None;
64                llm_args.fallback_value = None;
65                llm_args.prompt = Some(if attempt == 0 {
66                    tool_prompt.clone()
67                } else {
68                    format!(
69                        "{tool_prompt}\n\nYour previous response could not be parsed. Please respond with ONLY a JSON array of strings."
70                    )
71                });
72
73                let result = dispatch_llm(llm_args, ctx).await;
74                if let Value::Err(e) = result {
75                    return Err(e);
76                }
77
78                match parse_branches_result(&result) {
79                    Ok(list) => {
80                        let mut list = list;
81                        if let Some(c) = count {
82                            if list.len() > c {
83                                crate::notify!(
84                                    warn,
85                                    "llm.generate_branches: generated {} branches, truncating to {c}",
86                                    list.len()
87                                );
88                                list.truncate(c);
89                            }
90                        }
91                        return Ok(Value::List(list));
92                    }
93                    Err(_) if attempt < parse_retry_count => continue,
94                    Err(e) => {
95                        return Err(RuntimeError::ToolFailed(format!(
96                            "llm.generate_branches: {e}"
97                        )));
98                    }
99                }
100            }
101
102            Err(RuntimeError::ToolFailed(
103                "llm.generate_branches: unreachable".into(),
104            ))
105        })
106    }
107}
108
109fn extract_retry(args: &ToolArgs) -> u32 {
110    args.named("retry")
111        .and_then(|v| match v {
112            Value::Int(n) if *n >= 0 => Some(*n as u32),
113            _ => None,
114        })
115        .unwrap_or(0)
116}
117
118fn value_to_string(v: &Value) -> String {
119    match v {
120        Value::Str(s) => s.clone(),
121        Value::Int(n) => n.to_string(),
122        Value::Float(f) => f.to_string(),
123        Value::Bool(b) => b.to_string(),
124        Value::Unit => String::new(),
125        other => format!("{:?}", other),
126    }
127}
128
129fn parse_branches_result(result: &Value) -> Result<Vec<Value>, String> {
130    // If already List (assistant_message_to_value pre-parsed JSON array)
131    if let Value::List(items) = result {
132        return Ok(items
133            .iter()
134            .map(|v| match v {
135                Value::Str(_) => v.clone(),
136                other => Value::Str(value_to_string(other)),
137            })
138            .collect());
139    }
140
141    // Extract text and parse
142    let text = match llm_result_to_text(result) {
143        Some(t) => t,
144        None => return Err("could not extract text from LLM response".into()),
145    };
146
147    let items = parse_list_from_text(&text);
148    if items.is_empty() {
149        return Err("could not parse any branches from response".into());
150    }
151
152    Ok(items.into_iter().map(Value::Str).collect())
153}
atman_runtime/tools/llm_generate_branches.rs

atman_runtime/tools/
llm_generate_branches.rs
