agentlib_reasoning/

planner.rs

use crate::utils::{call_model, extract_text, parse_json};
use agentlib_core::{
    ModelMessage, ModelRequest, PlanTask, ReasoningContext, ReasoningEngine, ReasoningStep, Role,
};
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use serde::Deserialize;
use std::collections::HashMap;
use uuid::Uuid;

pub struct PlannerEngine {
    max_execution_steps: usize,
    allow_replan: bool,
    planner_prompt: String,
    executor_prompt: String,
}

impl PlannerEngine {
    pub fn new(
        max_execution_steps: usize,
        allow_replan: bool,
        planner_prompt: Option<String>,
        executor_prompt: Option<String>,
    ) -> Self {
        Self {
            max_execution_steps,
            allow_replan,
            planner_prompt: planner_prompt.unwrap_or_else(|| {
                "You are a planning assistant. Break the user's request into a clear, ordered list of subtasks.\n\nRespond with ONLY a JSON array of tasks in this exact format (no markdown, no preamble):\n[\n  { \"id\": \"t1\", \"description\": \"...\", \"dependsOn\": [] },\n  { \"id\": \"t2\", \"description\": \"...\", \"dependsOn\": [\"t1\"] }\n]\n\nRules:\n- Each task must be atomic and independently executable\n- dependsOn lists task ids that must complete first\n- Order tasks so dependencies come first\n- Be specific — the executor will act on each description".to_string()
            }),
            executor_prompt: executor_prompt.unwrap_or_else(|| {
                "You are an execution assistant. Complete the given subtask using available tools.\nFocus only on the current task. Be concise and direct.".to_string()
            }),
        }
    }

    async fn make_plan(&self, r_ctx: &mut ReasoningContext<'_>) -> Result<Vec<PlanTask>> {
        let available_tools = r_ctx
            .tools
            .list()
            .iter()
            .map(|t| format!("- {}: {:?}", t.name, t.description))
            .collect::<Vec<_>>()
            .join("\n");
        let tool_context = if available_tools.is_empty() {
            "".to_string()
        } else {
            format!("\n\nAvailable Tools:\n{}", available_tools)
        };

        let plan_messages = vec![
            ModelMessage {
                role: Role::System,
                content: format!("{}{}", self.planner_prompt, tool_context),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::User,
                content: r_ctx.ctx.input.clone(),
                tool_call_id: None,
                tool_calls: None,
            },
        ];

        let request = ModelRequest {
            messages: plan_messages,
            tools: None,
        };

        let response = r_ctx.model.complete(request).await?;

        // Accumulate usage
        if let Some(usage) = &response.usage {
            r_ctx.ctx.usage.prompt_tokens += usage.prompt_tokens;
            r_ctx.ctx.usage.completion_tokens += usage.completion_tokens;
            r_ctx.ctx.usage.total_tokens += usage.total_tokens;
        }

        #[derive(Deserialize)]
        struct RawTask {
            id: Option<String>,
            description: String,
            #[serde(rename = "dependsOn")]
            depends_on: Option<Vec<String>>,
        }

        match parse_json::<Vec<RawTask>>(&response.message.content) {
            Ok(raw_tasks) => Ok(raw_tasks
                .into_iter()
                .map(|t| PlanTask {
                    id: t.id.unwrap_or_else(|| Uuid::new_v4().to_string()),
                    description: t.description,
                    depends_on: t.depends_on,
                    status: "pending".to_string(),
                    result: None,
                })
                .collect()),
            Err(_) => Ok(vec![PlanTask {
                id: "t1".to_string(),
                description: r_ctx.ctx.input.clone(),
                depends_on: None,
                status: "pending".to_string(),
                result: None,
            }]),
        }
    }

    async fn execute_task(
        &self,
        r_ctx: &mut ReasoningContext<'_>,
        task: &PlanTask,
        previous_results: &HashMap<String, String>,
    ) -> Result<String> {
        let dep_context = if let Some(deps) = &task.depends_on {
            if deps.is_empty() {
                "".to_string()
            } else {
                format!(
                    "\n\nContext from previous tasks:\n{}",
                    deps.iter()
                        .map(|id| format!(
                            "[{}]: {}",
                            id,
                            previous_results.get(id).unwrap_or(&"N/A".to_string())
                        ))
                        .collect::<Vec<_>>()
                        .join("\n")
                )
            }
        } else {
            "".to_string()
        };

        let mut task_messages = vec![
            ModelMessage {
                role: Role::System,
                content: self.executor_prompt.clone(),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::User,
                content: format!(
                    "Original goal: {}\n\nCurrent task: {}{}",
                    r_ctx.ctx.input, task.description, dep_context
                ),
                tool_call_id: None,
                tool_calls: None,
            },
        ];

        let mut steps = 0;
        let max_task_steps = 5;

        while steps < max_task_steps {
            let response = call_model(r_ctx, task_messages.clone()).await?;
            task_messages.push(response.message.clone());

            if let Some(tool_calls) = &response.message.tool_calls {
                if tool_calls.is_empty() {
                    return Ok(extract_text(&response.message.content));
                }

                for tc in tool_calls {
                    let result = r_ctx
                        .call_tool(&tc.name, tc.arguments.clone(), tc.id.clone())
                        .await?;
                    task_messages.push(ModelMessage {
                        role: Role::Tool,
                        content: result.to_string(),
                        tool_call_id: Some(tc.id.clone()),
                        tool_calls: None,
                    });
                }
            } else {
                return Ok(extract_text(&response.message.content));
            }

            steps += 1;
        }

        Err(anyhow!(
            "[PlannerEngine] Task \"{}\" exceeded max steps.",
            task.id
        ))
    }

    async fn synthesize(
        &self,
        r_ctx: &mut ReasoningContext<'_>,
        plan: &[PlanTask],
        results: &HashMap<String, String>,
    ) -> Result<String> {
        let summary_context = plan
            .iter()
            .filter(|t| t.status == "done")
            .map(|t| {
                format!(
                    "[{}] {}:\n{}",
                    t.id,
                    t.description,
                    results.get(&t.id).unwrap_or(&"no result".to_string())
                )
            })
            .collect::<Vec<_>>()
            .join("\n\n");

        let synth_messages = vec![
            ModelMessage {
                role: Role::System,
                content: "Synthesize the results of the completed tasks into a clear, direct answer to the original user request.".to_string(),
                tool_call_id: None,
                tool_calls: None,
            },
            ModelMessage {
                role: Role::User,
                content: format!("Original request: {}\n\nTask results:\n{}", r_ctx.ctx.input, summary_context),
                tool_call_id: None,
                tool_calls: None,
            }
        ];

        let request = ModelRequest {
            messages: synth_messages,
            tools: None,
        };

        let response = r_ctx.model.complete(request).await?;

        // Accumulate usage
        if let Some(usage) = &response.usage {
            r_ctx.ctx.usage.prompt_tokens += usage.prompt_tokens;
            r_ctx.ctx.usage.completion_tokens += usage.completion_tokens;
            r_ctx.ctx.usage.total_tokens += usage.total_tokens;
        }

        Ok(response.message.content)
    }
}

impl Default for PlannerEngine {
    fn default() -> Self {
        Self::new(20, false, None, None)
    }
}

#[async_trait]
impl ReasoningEngine for PlannerEngine {
    fn name(&self) -> &str {
        "planner"
    }

    async fn execute(&self, r_ctx: &mut ReasoningContext<'_>) -> Result<String> {
        // Phase 1: Planning
        let mut plan = self.make_plan(r_ctx).await?;

        r_ctx.push_step(ReasoningStep::Plan {
            tasks: plan.clone(),
            engine: self.name().to_string(),
        });

        // Phase 2: Execution
        let mut task_results = HashMap::new();
        let mut execution_steps = 0;

        for i in 0..plan.len() {
            if execution_steps >= self.max_execution_steps {
                return Err(anyhow!(
                    "[PlannerEngine] Max execution steps ({}) reached.",
                    self.max_execution_steps
                ));
            }

            let task = &mut plan[i];

            // Check dependencies
            if let Some(deps) = &task.depends_on {
                let unmet_deps: Vec<_> = deps
                    .iter()
                    .filter(|dep| !task_results.contains_key(*dep))
                    .collect();
                if !unmet_deps.is_empty() {
                    continue;
                }
            }

            task.status = "in_progress".to_string();
            r_ctx.push_step(ReasoningStep::Thought {
                content: format!("Executing task [{}]: {}", task.id, task.description),
                engine: self.name().to_string(),
            });

            match self.execute_task(r_ctx, task, &task_results).await {
                Ok(result) => {
                    task.status = "done".to_string();
                    task.result = Some(serde_json::Value::String(result.clone()));
                    task_results.insert(task.id.clone(), result);
                    execution_steps += 1;
                }
                Err(err) => {
                    task.status = "failed".to_string();
                    if !self.allow_replan {
                        return Err(anyhow!(
                            "[PlannerEngine] Task \"{}\" failed: {}",
                            task.id,
                            err
                        ));
                    }
                }
            }
        }

        // Phase 3: Synthesize
        let summary = self.synthesize(r_ctx, &plan, &task_results).await?;
        r_ctx.push_step(ReasoningStep::Response {
            content: summary.clone(),
            engine: self.name().to_string(),
        });

        Ok(summary)
    }
}