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        self.run_with_config(objective, None).await
110    }
111
112    /// Runs the full task with a [`lc_core::runnables::RunnableConfig`] so the
113    /// planning/replanning
114    /// LLM calls emit `on_llm_start/end` to the configured callbacks/OTel
115    /// backend (T6, v0.23.0). The per-step executor was already traced in 0.22.4
116    /// A18; this closes the gap on the planner side.
117    pub async fn run_with_config(
118        &self,
119        objective: &str,
120        config: Option<&lc_core::runnables::RunnableConfig>,
121    ) -> Result<String, PlanExecuteError> {
122        let planner = Planner::new(self.llm.clone());
123        let mut plan = planner
124            .plan(objective, config)
125            .await
126            .map_err(|e| PlanExecuteError::PlanningError(e.to_string()))?;
127
128        for replan_count in 0..=self.max_replans {
129            let pending_ids: Vec<usize> = plan
130                .steps
131                .iter()
132                .filter(|s| s.status == StepStatus::Pending)
133                .map(|s| s.id)
134                .collect();
135
136            let mut failed = false;
137            for step_id in pending_ids {
138                let step = plan.steps.iter_mut().find(|s| s.id == step_id);
139                let step_desc = match step {
140                    Some(s) => {
141                        s.status = StepStatus::Running;
142                        s.description.clone()
143                    }
144                    None => {
145                        // step_id does not correspond to any step; skip it
146                        continue;
147                    }
148                };
149
150                match self.execute_step(&step_desc, config).await {
151                    Ok(result) => plan.mark_completed(step_id, result),
152                    Err(e) => {
153                        let error_msg = e.to_string();
154                        plan.mark_failed(step_id, error_msg.clone());
155                        if replan_count < self.max_replans {
156                            // 0.22.0 C4 fix: carry the already-completed steps
157                            // (and their results) into the replan — the prompt
158                            // asks for the *remaining* work, and the completed
159                            // steps are spliced back so the summary keeps the
160                            // history instead of re-executing everything.
161                            let completed_steps: Vec<(String, String)> = plan
162                                .steps
163                                .iter()
164                                .filter(|s| s.status == StepStatus::Completed)
165                                .filter_map(|s| {
166                                    s.result
167                                        .as_ref()
168                                        .map(|r| (s.description.clone(), r.clone()))
169                                })
170                                .collect();
171                            let completed_block = if completed_steps.is_empty() {
172                                String::new()
173                            } else {
174                                completed_steps
175                                    .iter()
176                                    .map(|(d, r)| format!("- {d} → {r}"))
177                                    .collect::<Vec<_>>()
178                                    .join("\n")
179                            };
180                            plan = planner
181                                .replan(objective, &step_desc, &error_msg, &completed_block, config)
182                                .await
183                                .map_err(|e| PlanExecuteError::PlanningError(e.to_string()))?;
184                            // Splice the completed steps (with results) back at
185                            // the front so they show up in the summary and are
186                            // not re-executed (only Pending steps run).
187                            let mut merged: Vec<PlanStep> = Vec::new();
188                            let mut next_id = 0usize;
189                            for (desc, result) in &completed_steps {
190                                let mut st = PlanStep::new(next_id, desc.clone());
191                                st.status = StepStatus::Completed;
192                                st.result = Some(result.clone());
193                                merged.push(st);
194                                next_id += 1;
195                            }
196                            for s in plan.steps {
197                                if s.status != StepStatus::Completed {
198                                    merged.push(PlanStep::new(next_id, s.description));
199                                    next_id += 1;
200                                }
201                            }
202                            plan = Plan::new(objective, merged);
203                            failed = true;
204                            break;
205                        } else {
206                            return Err(PlanExecuteError::MaxReplansReached {
207                                step: step_desc,
208                                reason: error_msg,
209                            });
210                        }
211                    }
212                }
213            }
214
215            if !failed && plan.is_complete() {
216                let summary: Vec<String> = plan
217                    .steps
218                    .iter()
219                    .map(|s| {
220                        format!(
221                            "{}. {}: {}",
222                            s.id + 1,
223                            s.description,
224                            s.result.as_deref().unwrap_or("无结果")
225                        )
226                    })
227                    .collect();
228                return Ok(summary.join("\n"));
229            }
230        }
231        Err(PlanExecuteError::PlanIncomplete)
232    }
233
234    /// Executes a single step: prefers the agent from `agent_factory`, falling back to FunctionCallingAgent (P1-2)
235    ///
236    /// T6 (v0.23.0): `config` is forwarded to the executor so per-step LLM calls
237    /// share the same callbacks/OTel trace as the planning calls.
238    async fn execute_step(
239        &self,
240        step: &str,
241        config: Option<&lc_core::runnables::RunnableConfig>,
242    ) -> Result<String, PlanExecuteError> {
243        let agent: Arc<dyn BaseAgent> = match &self.agent_factory {
244            Some(factory) => factory(),
245            None => Arc::new(FunctionCallingAgent::from_arc(
246                self.llm.clone(),
247                self.tools.clone(),
248                None,
249            )) as Arc<dyn BaseAgent>,
250        };
251        let executor = AgentExecutor::new(agent, self.tools.clone()).with_max_iterations(5);
252        executor
253            .invoke_with_config(step.to_string(), config.cloned())
254            .await
255            .map_err(|e| PlanExecuteError::StepExecutionError(e.to_string()))
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262    use crate::types::{AgentFinish, AgentOutput, AgentStep};
263    use crate::AgentError;
264    use async_trait::async_trait;
265    use lc_providers::{OpenAIChat, OpenAIConfig};
266    use std::collections::HashMap;
267
268    /// Offline-capable mock agent: returns Finish directly, to verify that agent_factory takes effect (P1-2).
269    struct FakeAgent;
270
271    #[async_trait]
272    impl BaseAgent for FakeAgent {
273        async fn plan(
274            &self,
275            _intermediate_steps: &[AgentStep],
276            _inputs: &HashMap<String, String>,
277            _config: Option<&lc_core::runnables::RunnableConfig>,
278        ) -> Result<AgentOutput, AgentError> {
279            Ok(AgentOutput::Finish(AgentFinish::new(
280                "executed by factory".to_string(),
281                String::new(),
282            )))
283        }
284    }
285
286    /// P1-2: execute_step should use the agent from agent_factory, not a hardcoded FunctionCallingAgent.
287    #[tokio::test]
288    async fn test_execute_step_uses_agent_factory() {
289        let agent = PlanExecuteAgent::new(OpenAIChat::new(OpenAIConfig::default()), vec![])
290            .with_agent_factory(Arc::new(|| Arc::new(FakeAgent) as Arc<dyn BaseAgent>));
291        let result = agent.execute_step("step 1", None).await.unwrap();
292        assert_eq!(result, "executed by factory");
293    }
294}