Skip to main content

lc_agents/plan_execute/
agent.rs

1//! PlanExecuteAgent - plan - execute - replan
2
3use std::sync::Arc;
4
5use crate::{AgentError, AgentExecutor, BaseAgent, FunctionCallingAgent};
6use lc_core::language_models::BaseChatModel;
7use lc_core::tools::BaseTool;
8use lc_providers::ProviderError;
9
10use super::plan::{Plan, PlanStep, StepStatus};
11use super::planner::Planner;
12
13/// Plan-Execute Agent error type
14#[derive(Debug, thiserror::Error)]
15#[non_exhaustive]
16pub enum PlanExecuteError {
17    /// Planning failed
18    #[error("Planning failed: {0}")]
19    PlanningError(String),
20    /// Step execution failed
21    #[error("Step execution failed: {0}")]
22    StepExecutionError(String),
23    /// Max replan count reached
24    #[error("Max replans reached: step [{step}] failed: {reason}")]
25    MaxReplansReached {
26        /// The failed step
27        step: String,
28        /// Failure reason
29        reason: String,
30    },
31    /// The plan is incomplete
32    #[error("Plan incomplete after all replans")]
33    PlanIncomplete,
34}
35
36impl From<AgentError> for PlanExecuteError {
37    fn from(e: AgentError) -> Self {
38        PlanExecuteError::StepExecutionError(e.to_string())
39    }
40}
41
42/// Plan-Execute Agent: plans first, executes step by step, and replans on failure.
43///
44/// Supports any LLM provider that implements `BaseChatModel`.
45pub struct PlanExecuteAgent {
46    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
47    tools: Vec<Arc<dyn BaseTool>>,
48    max_replans: usize,
49    /// Factory for the per-step execution agent (P1-2). Defaults to `FunctionCallingAgent`.
50    agent_factory: Option<Arc<dyn Fn() -> Arc<dyn BaseAgent> + Send + Sync>>,
51}
52
53impl PlanExecuteAgent {
54    /// Creates a new Plan-Execute Agent
55    ///
56    /// # Parameters
57    /// * `llm` - LLM client (any type implementing `BaseChatModel`)
58    /// * `tools` - available tools
59    ///
60    /// # Backward compatibility
61    /// Legacy code `PlanExecuteAgent::new(openai_chat, tools)` still works.
62    pub fn new<L>(llm: L, tools: Vec<Arc<dyn BaseTool>>) -> Self
63    where
64        L: BaseChatModel + Send + Sync + 'static,
65        L::Error: Into<ProviderError>,
66    {
67        Self {
68            llm: lc_providers::wrap_chat_model(llm),
69            tools,
70            max_replans: 2,
71            agent_factory: None,
72        }
73    }
74
75    /// Creates an agent from an already-wrapped `Arc<dyn BaseChatModel>`
76    pub fn from_arc(
77        llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
78        tools: Vec<Arc<dyn BaseTool>>,
79    ) -> Self {
80        Self {
81            llm,
82            tools,
83            max_replans: 2,
84            agent_factory: None,
85        }
86    }
87
88    /// Sets the maximum number of replans.
89    pub fn with_max_replans(mut self, n: usize) -> Self {
90        self.max_replans = n;
91        self
92    }
93
94    /// Custom factory for the per-step execution agent (P1-2).
95    ///
96    /// The `BaseAgent` returned by the factory is used for each `execute_step`
97    /// of the PlanExecute. Defaults to `FunctionCallingAgent`. Useful for
98    /// injecting ReAct / Streaming / custom agents as the execution agent.
99    pub fn with_agent_factory(
100        mut self,
101        factory: Arc<dyn Fn() -> Arc<dyn BaseAgent> + Send + Sync>,
102    ) -> Self {
103        self.agent_factory = Some(factory);
104        self
105    }
106
107    /// Runs the full task: plan -> execute each step -> replan on failure -> summarize
108    pub async fn run(&self, objective: &str) -> Result<String, PlanExecuteError> {
109        let planner = Planner::new(self.llm.clone());
110        let mut plan = planner
111            .plan(objective)
112            .await
113            .map_err(|e| PlanExecuteError::PlanningError(e.to_string()))?;
114
115        for replan_count in 0..=self.max_replans {
116            let pending_ids: Vec<usize> = plan
117                .steps
118                .iter()
119                .filter(|s| s.status == StepStatus::Pending)
120                .map(|s| s.id)
121                .collect();
122
123            let mut failed = false;
124            for step_id in pending_ids {
125                let step = plan.steps.iter_mut().find(|s| s.id == step_id);
126                let step_desc = match step {
127                    Some(s) => {
128                        s.status = StepStatus::Running;
129                        s.description.clone()
130                    }
131                    None => {
132                        // step_id does not correspond to any step; skip it
133                        continue;
134                    }
135                };
136
137                match self.execute_step(&step_desc).await {
138                    Ok(result) => plan.mark_completed(step_id, result),
139                    Err(e) => {
140                        let error_msg = e.to_string();
141                        plan.mark_failed(step_id, error_msg.clone());
142                        if replan_count < self.max_replans {
143                            // 0.22.0 C4 fix: carry the already-completed steps
144                            // (and their results) into the replan — the prompt
145                            // asks for the *remaining* work, and the completed
146                            // steps are spliced back so the summary keeps the
147                            // history instead of re-executing everything.
148                            let completed_steps: Vec<(String, String)> = plan
149                                .steps
150                                .iter()
151                                .filter(|s| s.status == StepStatus::Completed)
152                                .filter_map(|s| {
153                                    s.result
154                                        .as_ref()
155                                        .map(|r| (s.description.clone(), r.clone()))
156                                })
157                                .collect();
158                            let completed_block = if completed_steps.is_empty() {
159                                String::new()
160                            } else {
161                                completed_steps
162                                    .iter()
163                                    .map(|(d, r)| format!("- {d} → {r}"))
164                                    .collect::<Vec<_>>()
165                                    .join("\n")
166                            };
167                            plan = planner
168                                .replan(objective, &step_desc, &error_msg, &completed_block)
169                                .await
170                                .map_err(|e| PlanExecuteError::PlanningError(e.to_string()))?;
171                            // Splice the completed steps (with results) back at
172                            // the front so they show up in the summary and are
173                            // not re-executed (only Pending steps run).
174                            let mut merged: Vec<PlanStep> = Vec::new();
175                            let mut next_id = 0usize;
176                            for (desc, result) in &completed_steps {
177                                let mut st = PlanStep::new(next_id, desc.clone());
178                                st.status = StepStatus::Completed;
179                                st.result = Some(result.clone());
180                                merged.push(st);
181                                next_id += 1;
182                            }
183                            for s in plan.steps {
184                                if s.status != StepStatus::Completed {
185                                    merged.push(PlanStep::new(next_id, s.description));
186                                    next_id += 1;
187                                }
188                            }
189                            plan = Plan::new(objective, merged);
190                            failed = true;
191                            break;
192                        } else {
193                            return Err(PlanExecuteError::MaxReplansReached {
194                                step: step_desc,
195                                reason: error_msg,
196                            });
197                        }
198                    }
199                }
200            }
201
202            if !failed && plan.is_complete() {
203                let summary: Vec<String> = plan
204                    .steps
205                    .iter()
206                    .map(|s| {
207                        format!(
208                            "{}. {}: {}",
209                            s.id + 1,
210                            s.description,
211                            s.result.as_deref().unwrap_or("无结果")
212                        )
213                    })
214                    .collect();
215                return Ok(summary.join("\n"));
216            }
217        }
218        Err(PlanExecuteError::PlanIncomplete)
219    }
220
221    /// Executes a single step: prefers the agent from `agent_factory`, falling back to FunctionCallingAgent (P1-2)
222    async fn execute_step(&self, step: &str) -> Result<String, PlanExecuteError> {
223        let agent: Arc<dyn BaseAgent> = match &self.agent_factory {
224            Some(factory) => factory(),
225            None => Arc::new(FunctionCallingAgent::from_arc(
226                self.llm.clone(),
227                self.tools.clone(),
228                None,
229            )) as Arc<dyn BaseAgent>,
230        };
231        let executor = AgentExecutor::new(agent, self.tools.clone()).with_max_iterations(5);
232        executor
233            .invoke(step.to_string())
234            .await
235            .map_err(|e| PlanExecuteError::StepExecutionError(e.to_string()))
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use crate::types::{AgentFinish, AgentOutput, AgentStep};
243    use crate::AgentError;
244    use async_trait::async_trait;
245    use lc_providers::{OpenAIChat, OpenAIConfig};
246    use std::collections::HashMap;
247
248    /// Offline-capable mock agent: returns Finish directly, to verify that agent_factory takes effect (P1-2).
249    struct FakeAgent;
250
251    #[async_trait]
252    impl BaseAgent for FakeAgent {
253        async fn plan(
254            &self,
255            _intermediate_steps: &[AgentStep],
256            _inputs: &HashMap<String, String>,
257            _config: Option<&lc_core::runnables::RunnableConfig>,
258        ) -> Result<AgentOutput, AgentError> {
259            Ok(AgentOutput::Finish(AgentFinish::new(
260                "executed by factory".to_string(),
261                String::new(),
262            )))
263        }
264    }
265
266    /// P1-2: execute_step should use the agent from agent_factory, not a hardcoded FunctionCallingAgent.
267    #[tokio::test]
268    async fn test_execute_step_uses_agent_factory() {
269        let agent = PlanExecuteAgent::new(OpenAIChat::new(OpenAIConfig::default()), vec![])
270            .with_agent_factory(Arc::new(|| Arc::new(FakeAgent) as Arc<dyn BaseAgent>));
271        let result = agent.execute_step("step 1").await.unwrap();
272        assert_eq!(result, "executed by factory");
273    }
274}