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.call_purpose = crate::context_plan::ContextCallPurpose::BranchGeneration;
66                llm_args.prompt = Some(if attempt == 0 {
67                    tool_prompt.clone()
68                } else {
69                    format!(
70                        "{tool_prompt}\n\nYour previous response could not be parsed. Please respond with ONLY a JSON array of strings."
71                    )
72                });
73
74                let result = dispatch_llm(llm_args, ctx).await;
75                if let Value::Err(e) = result {
76                    return Err(e);
77                }
78
79                match parse_branches_result(&result) {
80                    Ok(list) => {
81                        let mut list = list;
82                        if let Some(c) = count {
83                            if list.len() > c {
84                                crate::notify!(
85                                    warn,
86                                    "llm.generate_branches: generated {} branches, truncating to {c}",
87                                    list.len()
88                                );
89                                list.truncate(c);
90                            }
91                        }
92                        return Ok(Value::List(list));
93                    }
94                    Err(_) if attempt < parse_retry_count => continue,
95                    Err(e) => {
96                        return Err(RuntimeError::ToolFailed(format!(
97                            "llm.generate_branches: {e}"
98                        )));
99                    }
100                }
101            }
102
103            Err(RuntimeError::ToolFailed(
104                "llm.generate_branches: unreachable".into(),
105            ))
106        })
107    }
108}
109
110fn extract_retry(args: &ToolArgs) -> u32 {
111    args.named("retry")
112        .and_then(|v| match v {
113            Value::Int(n) if *n >= 0 => Some(*n as u32),
114            _ => None,
115        })
116        .unwrap_or(0)
117}
118
119fn value_to_string(v: &Value) -> String {
120    match v {
121        Value::Str(s) => s.clone(),
122        Value::Int(n) => n.to_string(),
123        Value::Float(f) => f.to_string(),
124        Value::Bool(b) => b.to_string(),
125        Value::Unit => String::new(),
126        other => format!("{:?}", other),
127    }
128}
129
130fn parse_branches_result(result: &Value) -> Result<Vec<Value>, String> {
131    // If already List (assistant_message_to_value pre-parsed JSON array)
132    if let Value::List(items) = result {
133        return Ok(items
134            .iter()
135            .map(|v| match v {
136                Value::Str(_) => v.clone(),
137                other => Value::Str(value_to_string(other)),
138            })
139            .collect());
140    }
141
142    // Extract text and parse
143    let text = match llm_result_to_text(result) {
144        Some(t) => t,
145        None => return Err("could not extract text from LLM response".into()),
146    };
147
148    let items = parse_list_from_text(&text);
149    if items.is_empty() {
150        return Err("could not parse any branches from response".into());
151    }
152
153    Ok(items.into_iter().map(Value::Str).collect())
154}
atman_runtime/tools/llm_generate_branches.rs

atman_runtime/tools/
llm_generate_branches.rs
