Skip to main content

sim_lib_agent/
planning.rs

1use sim_kernel::{Cx, Error, Expr, Result, Symbol};
2use sim_lib_agent_runner_core::{ModelRequest, ModelResponse, ModelRunner};
3
4/// Unit of work produced by planning combinators.
5#[derive(Clone, Debug, PartialEq, Eq)]
6pub struct PlanningTask {
7    /// Stable task identifier.
8    pub id: String,
9    /// Prompt that drives the task.
10    pub prompt: String,
11}
12
13impl PlanningTask {
14    /// Builds a planning task from an id and prompt.
15    pub fn new(id: impl Into<String>, prompt: impl Into<String>) -> Self {
16        Self {
17            id: id.into(),
18            prompt: prompt.into(),
19        }
20    }
21}
22
23/// Output captured after running one planning task.
24#[derive(Clone, Debug, PartialEq, Eq)]
25pub struct PlanningOutput {
26    /// Task that produced this output.
27    pub task: PlanningTask,
28    /// Text content captured from the run.
29    pub content: String,
30}
31
32impl PlanningOutput {
33    /// Builds an output from a task and its captured content.
34    pub fn new(task: PlanningTask, content: impl Into<String>) -> Self {
35        Self {
36            task,
37            content: content.into(),
38        }
39    }
40}
41
42/// Result of decomposing a goal and running the ordered subtasks.
43#[derive(Clone, Debug, PartialEq, Eq)]
44pub struct Decomposition {
45    /// Original goal that was decomposed.
46    pub goal: PlanningTask,
47    /// Ordered subtasks produced by decomposition.
48    pub subtasks: Vec<PlanningTask>,
49    /// Output captured for each executed subtask.
50    pub outputs: Vec<PlanningOutput>,
51    /// Remaining step budget after execution.
52    pub budget_left: u32,
53}
54
55/// Critique result, including a bounded retry when the runner asks for one.
56#[derive(Clone, Debug, PartialEq, Eq)]
57pub struct Reflection {
58    /// Whether the runner accepted the output.
59    pub accept: bool,
60    /// Critique text explaining the verdict.
61    pub critique: String,
62    /// Retry task requested by the runner, when any.
63    pub retry: Option<PlanningTask>,
64    /// Output captured from the executed retry, when one ran.
65    pub retry_output: Option<PlanningOutput>,
66    /// Remaining retry budget after reflection.
67    pub budget_left: u32,
68}
69
70/// Asks a runner to split a goal into deterministic ordered subtasks.
71pub fn decompose(
72    cx: &mut Cx,
73    goal: &PlanningTask,
74    runner: &dyn ModelRunner,
75    max_steps: u32,
76) -> Result<Vec<PlanningTask>> {
77    if max_steps == 0 {
78        return Err(Error::Eval(
79            "decompose budget exhausted before planning".to_owned(),
80        ));
81    }
82    let response = runner.infer(
83        cx,
84        ModelRequest::new(
85            operation_expr("decompose", vec![(Symbol::new("goal"), task_expr(goal))]),
86            Vec::new(),
87        ),
88    )?;
89    let subtasks = parse_tasks(&response, &goal.id)?;
90    if subtasks.is_empty() {
91        return Err(Error::Eval(
92            "decompose runner returned no subtasks".to_owned(),
93        ));
94    }
95    if subtasks.len() as u32 > max_steps {
96        return Err(Error::Eval(format!(
97            "decompose produced {} subtasks over budget {max_steps}",
98            subtasks.len()
99        )));
100    }
101    Ok(subtasks)
102}
103
104/// Decomposes a goal, then executes each returned subtask in order.
105pub fn decompose_and_run(
106    cx: &mut Cx,
107    goal: &PlanningTask,
108    runner: &dyn ModelRunner,
109    max_steps: u32,
110) -> Result<Decomposition> {
111    let subtasks = decompose(cx, goal, runner, max_steps)?;
112    let mut outputs = Vec::with_capacity(subtasks.len());
113    for task in &subtasks {
114        outputs.push(run_task(cx, task, runner)?);
115    }
116    Ok(Decomposition {
117        goal: goal.clone(),
118        budget_left: max_steps.saturating_sub(outputs.len() as u32),
119        subtasks,
120        outputs,
121    })
122}
123
124/// Critiques an output and runs at most one retry requested by the runner.
125pub fn reflect(
126    cx: &mut Cx,
127    output: &PlanningOutput,
128    runner: &dyn ModelRunner,
129    retry_budget: u32,
130) -> Result<Reflection> {
131    let response = runner.infer(
132        cx,
133        ModelRequest::new(
134            operation_expr(
135                "reflect",
136                vec![(Symbol::new("output"), output_expr(output))],
137            ),
138            Vec::new(),
139        ),
140    )?;
141    let mut reflection = parse_reflection(&response, retry_budget)?;
142    if !reflection.accept && reflection.retry.is_some() {
143        if retry_budget == 0 {
144            return Err(Error::Eval(
145                "reflect retry budget exhausted before re-run".to_owned(),
146            ));
147        }
148        let retry = reflection.retry.clone().expect("retry checked above");
149        reflection.retry_output = Some(run_task(cx, &retry, runner)?);
150        reflection.budget_left = retry_budget - 1;
151    }
152    Ok(reflection)
153}
154
155fn run_task(cx: &mut Cx, task: &PlanningTask, runner: &dyn ModelRunner) -> Result<PlanningOutput> {
156    let response = runner.infer(
157        cx,
158        ModelRequest::new(
159            operation_expr("execute", vec![(Symbol::new("task"), task_expr(task))]),
160            Vec::new(),
161        ),
162    )?;
163    Ok(PlanningOutput::new(task.clone(), response_text(&response)))
164}
165
166fn parse_tasks(response: &ModelResponse, goal_id: &str) -> Result<Vec<PlanningTask>> {
167    let payloads = response_payloads(response);
168    let mut tasks = Vec::new();
169    for (index, payload) in payloads.iter().enumerate() {
170        match payload {
171            Expr::List(items) | Expr::Vector(items) => {
172                for (item_index, item) in items.iter().enumerate() {
173                    tasks.push(parse_task(item, &step_id(goal_id, item_index + 1))?);
174                }
175            }
176            Expr::String(text) => {
177                tasks.extend(parse_text_tasks(goal_id, text));
178            }
179            Expr::Map(entries) => {
180                if let Some(text) = text_part(entries) {
181                    tasks.extend(parse_text_tasks(goal_id, &text));
182                } else {
183                    tasks.push(parse_task(payload, &step_id(goal_id, index + 1))?);
184                }
185            }
186            expr => tasks.push(parse_task(expr, &step_id(goal_id, index + 1))?),
187        }
188    }
189    Ok(tasks)
190}
191
192fn parse_reflection(response: &ModelResponse, retry_budget: u32) -> Result<Reflection> {
193    let payloads = response_payloads(response);
194    if let Some(Expr::Map(entries)) = payloads.first() {
195        let accept = bool_field(entries, "accept").unwrap_or(false);
196        let critique = string_field(entries, "critique")
197            .or_else(|| string_field(entries, "reason"))
198            .unwrap_or_else(|| response_text(response));
199        let retry = expr_field(entries, "retry")
200            .or_else(|| expr_field(entries, "retry-task"))
201            .filter(|expr| !matches!(expr, Expr::Nil))
202            .map(|expr| parse_task(expr, "retry-1"))
203            .transpose()?;
204        return Ok(Reflection {
205            accept,
206            critique,
207            retry,
208            retry_output: None,
209            budget_left: retry_budget,
210        });
211    }
212
213    let text = response_text(response);
214    let lower = text.to_ascii_lowercase();
215    let accept = lower.starts_with("accept") || lower.contains("accept: true");
216    let retry = retry_text(&text).map(|prompt| PlanningTask::new("retry-1", prompt));
217    Ok(Reflection {
218        accept,
219        critique: text,
220        retry,
221        retry_output: None,
222        budget_left: retry_budget,
223    })
224}
225
226fn parse_task(expr: &Expr, default_id: &str) -> Result<PlanningTask> {
227    match expr {
228        Expr::Map(entries) => {
229            let id = string_field(entries, "id").unwrap_or_else(|| default_id.to_owned());
230            let prompt = string_field(entries, "prompt")
231                .or_else(|| string_field(entries, "instruction"))
232                .or_else(|| string_field(entries, "task"))
233                .ok_or_else(|| {
234                    Error::Eval(format!(
235                        "planning task {id} requires prompt, instruction, or task"
236                    ))
237                })?;
238            Ok(PlanningTask::new(id, prompt))
239        }
240        Expr::String(text) => Ok(PlanningTask::new(default_id.to_owned(), text.clone())),
241        expr => Ok(PlanningTask::new(default_id.to_owned(), expr_text(expr))),
242    }
243}
244
245fn parse_text_tasks(goal_id: &str, text: &str) -> Vec<PlanningTask> {
246    text.lines()
247        .map(str::trim)
248        .filter(|line| !line.is_empty())
249        .enumerate()
250        .map(|(index, line)| {
251            PlanningTask::new(step_id(goal_id, index + 1), clean_list_marker(line))
252        })
253        .collect()
254}
255
256fn operation_expr(name: &str, fields: Vec<(Symbol, Expr)>) -> Expr {
257    let mut entries = vec![(
258        Expr::Symbol(Symbol::new("op")),
259        Expr::Symbol(Symbol::new(name)),
260    )];
261    entries.extend(
262        fields
263            .into_iter()
264            .map(|(key, value)| (Expr::Symbol(key), value)),
265    );
266    Expr::Map(entries)
267}
268
269fn task_expr(task: &PlanningTask) -> Expr {
270    Expr::Map(vec![
271        (
272            Expr::Symbol(Symbol::new("id")),
273            Expr::String(task.id.clone()),
274        ),
275        (
276            Expr::Symbol(Symbol::new("prompt")),
277            Expr::String(task.prompt.clone()),
278        ),
279    ])
280}
281
282fn output_expr(output: &PlanningOutput) -> Expr {
283    Expr::Map(vec![
284        (Expr::Symbol(Symbol::new("task")), task_expr(&output.task)),
285        (
286            Expr::Symbol(Symbol::new("content")),
287            Expr::String(output.content.clone()),
288        ),
289    ])
290}
291
292fn response_payloads(response: &ModelResponse) -> Vec<Expr> {
293    if response.content.len() == 1
294        && let Expr::Map(entries) = &response.content[0]
295        && let Some(text) = text_part(entries)
296    {
297        return vec![Expr::String(text)];
298    }
299    response.content.clone()
300}
301
302fn response_text(response: &ModelResponse) -> String {
303    response
304        .content
305        .iter()
306        .map(expr_text)
307        .collect::<Vec<_>>()
308        .join("\n")
309}
310
311fn expr_text(expr: &Expr) -> String {
312    match expr {
313        Expr::Nil => "nil".to_owned(),
314        Expr::Bool(value) => value.to_string(),
315        Expr::Symbol(symbol) => symbol.to_string(),
316        Expr::String(text) => text.clone(),
317        Expr::Map(entries) => text_part(entries).unwrap_or_else(|| format!("{expr:?}")),
318        _ => format!("{expr:?}"),
319    }
320}
321
322fn text_part(entries: &[(Expr, Expr)]) -> Option<String> {
323    match string_field(entries, "text") {
324        Some(text) if matches!(symbol_field(entries, "type").as_deref(), Some("text")) => {
325            Some(text)
326        }
327        _ => None,
328    }
329}
330
331fn string_field(entries: &[(Expr, Expr)], field: &str) -> Option<String> {
332    expr_field(entries, field).and_then(|expr| match expr {
333        Expr::String(text) => Some(text.clone()),
334        Expr::Symbol(symbol) => Some(symbol.to_string()),
335        _ => None,
336    })
337}
338
339fn symbol_field(entries: &[(Expr, Expr)], field: &str) -> Option<String> {
340    expr_field(entries, field).and_then(|expr| match expr {
341        Expr::Symbol(symbol) => Some(symbol.to_string()),
342        _ => None,
343    })
344}
345
346fn bool_field(entries: &[(Expr, Expr)], field: &str) -> Option<bool> {
347    expr_field(entries, field).and_then(|expr| match expr {
348        Expr::Bool(value) => Some(*value),
349        _ => None,
350    })
351}
352
353fn expr_field<'a>(entries: &'a [(Expr, Expr)], field: &str) -> Option<&'a Expr> {
354    entries.iter().find_map(|(key, value)| {
355        if matches!(key, Expr::Symbol(symbol) if symbol.name.as_ref() == field) {
356            Some(value)
357        } else {
358            None
359        }
360    })
361}
362
363fn retry_text(text: &str) -> Option<String> {
364    let lower = text.to_ascii_lowercase();
365    lower.find("retry:").map(|index| {
366        text[index + "retry:".len()..]
367            .trim()
368            .trim_matches('"')
369            .to_owned()
370    })
371}
372
373fn clean_list_marker(line: &str) -> String {
374    let trimmed = line.trim_start_matches(['-', '*', ' ']);
375    let without_number = trimmed
376        .split_once('.')
377        .and_then(|(head, tail)| head.chars().all(|ch| ch.is_ascii_digit()).then_some(tail))
378        .unwrap_or(trimmed);
379    without_number.trim().to_owned()
380}
381
382fn step_id(goal_id: &str, one_based_index: usize) -> String {
383    format!("{goal_id}.{one_based_index}")
384}