1use sim_kernel::{Cx, Error, Expr, Result, Symbol};
2use sim_lib_agent_runner_core::{ModelRequest, ModelResponse, ModelRunner};
3
4#[derive(Clone, Debug, PartialEq, Eq)]
6pub struct PlanningTask {
7 pub id: String,
9 pub prompt: String,
11}
12
13impl PlanningTask {
14 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#[derive(Clone, Debug, PartialEq, Eq)]
25pub struct PlanningOutput {
26 pub task: PlanningTask,
28 pub content: String,
30}
31
32impl PlanningOutput {
33 pub fn new(task: PlanningTask, content: impl Into<String>) -> Self {
35 Self {
36 task,
37 content: content.into(),
38 }
39 }
40}
41
42#[derive(Clone, Debug, PartialEq, Eq)]
44pub struct Decomposition {
45 pub goal: PlanningTask,
47 pub subtasks: Vec<PlanningTask>,
49 pub outputs: Vec<PlanningOutput>,
51 pub budget_left: u32,
53}
54
55#[derive(Clone, Debug, PartialEq, Eq)]
57pub struct Reflection {
58 pub accept: bool,
60 pub critique: String,
62 pub retry: Option<PlanningTask>,
64 pub retry_output: Option<PlanningOutput>,
66 pub budget_left: u32,
68}
69
70pub 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
104pub 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
124pub 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}