1"""Plan and Execute strategy for agent reasoning.
2
3Implements a two-phase reasoning approach:
4 1. **PLAN** — The LLM creates a numbered step-by-step plan.
5 2. **EXECUTE** — Each step is executed in order (tool call or LLM reasoning).
6 3. **REPLAN** — On step failure, the LLM re-plans remaining steps.
7 4. **SYNTHESIZE** — After execution, the LLM synthesizes a final answer.
8
9This strategy excels at complex, multi-step tasks where the full plan
10can be decomposed upfront. For tasks requiring reactive adaptation,
11use ``ReActStrategy`` instead.
12
13Reference:
14 Wang et al., "Plan-and-Solve Prompting" (2023)
15"""
16
17from __future__ import annotations
18
19import time
20from typing import TYPE_CHECKING, Any
21
22from lexigram.ai.agents.strategies.base import AbstractStrategy
23from lexigram.ai.agents.strategies.plan_execute_executor import (
24 call_llm,
25 execute_reasoning_step,
26 execute_tool_step,
27 replan,
28 run_tool,
29 synthesize,
30)
31from lexigram.ai.agents.strategies.plan_execute_planner import (
32 extract_final_answer,
33 extract_step_result,
34 format_completed_steps,
35 format_plan,
36 parse_plan,
37)
38from lexigram.ai.agents.strategies.plan_execute_types import (
39 PLANNING_PROMPT,
40 PlanStep,
41 PlanStepStatus,
42)
43from lexigram.ai.agents.strategies.token_utils import TokenAccumulator
44from lexigram.ai.agents.types import ReasoningStep, ToolExecutionRecord
45from lexigram.contracts.ai.agents import AgentError, AgentResponse
46from lexigram.logging import (
47 get_logger,
48)
49from lexigram.result import Err, Ok, Result
50
51if TYPE_CHECKING:
52 from lexigram.contracts.ai.agents import ToolProtocol
53 from lexigram.contracts.ai.llm import LLMClientProtocol
54
55logger = get_logger(__name__)
56
57
58class PlanAndExecuteStrategy(AbstractStrategy):
59 """Plan and Execute reasoning strategy.
60
61 Decomposes complex tasks into explicit plans, executes each step
62 sequentially, and synthesizes the results into a final answer.
63
64 This strategy is best for:
65 - Multi-step research tasks
66 - Tasks requiring ordered operations
67 - Complex queries that benefit from explicit decomposition
68
69 Example::
70
71 strategy = PlanAndExecuteStrategy(max_steps=8, max_replans=2)
72 result = await strategy.execute(
73 message="Compare the revenue of Apple and Google",
74 tools=[search_tool, calculator_tool],
75 history=[],
76 llm=llm_client,
77 )
78 """
79
80 def __init__(
81 self,
82 max_steps: int = 10,
83 max_replans: int = 2,
84 tool_timeout: float = 30.0,
85 observation_max_chars: int = 10_000,
86 llm_timeout: float = 120.0,
87 ) -> None:
88 """Initialize the Plan and Execute strategy.
89
90 Args:
91 max_steps: Maximum number of plan steps.
92 max_replans: Maximum replanning attempts on step failure.
93 tool_timeout: Per-tool execution timeout in seconds.
94 observation_max_chars: Max characters for tool output before truncation.
95 llm_timeout: Per-LLM-call timeout in seconds.
96 """
97 self.max_steps = max_steps
98 self.max_replans = max_replans
99 self.tool_timeout = tool_timeout
100 self.observation_max_chars = observation_max_chars
101 self.llm_timeout = llm_timeout
102
103 async def execute(
104 self,
105 message: str,
106 tools: list[ToolProtocol],
107 history: list[dict[str, Any]],
108 llm: LLMClientProtocol,
109 **kwargs: Any,
110 ) -> Result[AgentResponse, AgentError]:
111 """Execute the Plan-and-Execute reasoning loop.
112
113 Args:
114 message: The user's input message.
115 tools: Tools available to the agent.
116 history: Conversation history as list of message dicts.
117 llm: LLM client implementing ``LLMClientProtocol``.
118 **kwargs: Additional parameters (system_prompt, etc.).
119
120 Returns:
121 ``Ok(AgentResponse)`` with the final synthesized answer and
122 full reasoning trace. ``Err(AgentError)`` on failure.
123 """
124 system_prompt: str = kwargs.get("system_prompt", "")
125 guard_pipeline = kwargs.get("guard_pipeline")
126 steps: list[ReasoningStep] = []
127 tool_calls: list[ToolExecutionRecord] = []
128 usage = TokenAccumulator()
129 start_time = time.monotonic()
130
131 tool_map: dict[str, ToolProtocol] = {t.name: t for t in tools}
132 tool_descriptions = "\n".join(f"- **{t.name}**: {t.description}" for t in tools)
133
134 # ---- PHASE 1: PLANNING ----
135 planning_prompt = PLANNING_PROMPT.format(
136 max_steps=self.max_steps,
137 tool_descriptions=tool_descriptions or "(no tools available)",
138 )
139 plan_text = await self._call_llm(
140 llm,
141 message,
142 history,
143 system_prompt + planning_prompt,
144 usage=usage,
145 )
146 if plan_text is None:
147 return Err(AgentError("LLM returned empty response during planning phase"))
148
149 plan = self._parse_plan(plan_text)
150 if not plan:
151 # No structured plan — fall back to direct synthesis
152 return await self._direct_synthesis(
153 llm,
154 message,
155 plan_text,
156 history,
157 system_prompt,
158 steps,
159 tool_calls,
160 start_time,
161 usage=usage,
162 )
163
164 steps.append(
165 ReasoningStep(
166 step_number=0,
167 thought=f"Created plan with {len(plan)} steps",
168 action="plan",
169 observation=plan_text,
170 )
171 )
172
173 logger.info(
174 "plan_execute_plan_created",
175 step_count=len(plan),
176 tool_steps=sum(1 for s in plan if s.tool_name),
177 )
178
179 # ---- PHASE 2: EXECUTION ----
180 replan_count = 0
181
182 for plan_step in plan:
183 if plan_step.status == PlanStepStatus.SKIPPED:
184 continue
185
186 if plan_step.tool_name and plan_step.tool_name in tool_map:
187 # Tool execution step
188 step_result, tool_record = await self._execute_tool_step(
189 llm,
190 plan_step,
191 plan,
192 tool_map,
193 message,
194 history,
195 system_prompt,
196 usage=usage,
197 guard_pipeline=guard_pipeline,
198 )
199 if tool_record:
200 tool_calls.append(tool_record)
201
202 steps.append(
203 ReasoningStep(
204 step_number=plan_step.number,
205 thought=f"Executing: {plan_step.description}",
206 action=plan_step.tool_name,
207 tool_call=tool_record,
208 observation=step_result,
209 )
210 )
211
212 if (
213 plan_step.status == PlanStepStatus.FAILED
214 and replan_count < self.max_replans
215 ):
216 # ---- PHASE 3: REPLAN ----
217 replan_count += 1
218 new_plan = await self._replan(
219 llm,
220 plan,
221 plan_step,
222 step_result,
223 message,
224 history,
225 system_prompt,
226 usage=usage,
227 )
228 if new_plan:
229 # Replace remaining steps
230 for old_step in plan:
231 if old_step.number > plan_step.number:
232 old_step.status = PlanStepStatus.SKIPPED
233 plan.extend(new_plan)
234 steps.append(
235 ReasoningStep(
236 step_number=plan_step.number,
237 thought=f"Replanning after failure (attempt {replan_count})",
238 action="replan",
239 observation=f"Created {len(new_plan)} new steps",
240 )
241 )
242 continue
243 else:
244 # LLM reasoning step
245 step_result = await self._execute_reasoning_step(
246 llm,
247 plan_step,
248 plan,
249 message,
250 history,
251 system_prompt,
252 usage=usage,
253 guard_pipeline=guard_pipeline,
254 )
255
256 steps.append(
257 ReasoningStep(
258 step_number=plan_step.number,
259 thought=f"Reasoning: {plan_step.description}",
260 action="reason",
261 observation=step_result,
262 )
263 )
264
265 plan_step.result = step_result
266 plan_step.status = PlanStepStatus.COMPLETED
267
268 # ---- PHASE 4: SYNTHESIS ----
269 final_answer = await self._synthesize(
270 llm,
271 message,
272 plan,
273 history,
274 system_prompt,
275 usage=usage,
276 )
277
278 elapsed = (time.monotonic() - start_time) * 1000
279 logger.info(
280 "plan_execute_complete",
281 steps=len(steps),
282 tool_calls=len(tool_calls),
283 replans=replan_count,
284 duration_ms=round(elapsed, 2),
285 )
286
287 return Ok(
288 AgentResponse(
289 message=final_answer,
290 steps=steps,
291 tool_calls=tool_calls,
292 total_tokens=usage.total_tokens,
293 prompt_tokens=usage.prompt_tokens,
294 completion_tokens=usage.completion_tokens,
295 duration_ms=elapsed,
296 metadata={
297 "strategy": "plan_and_execute",
298 "plan_steps": len(plan),
299 "replans": replan_count,
300 },
301 )
302 )
303
304 # ------------------------------------------------------------------
305 # Plan Parsing
306 # ------------------------------------------------------------------
307
308 @staticmethod
309 def _parse_plan(text: str) -> list[PlanStep]:
310 """Parse a numbered plan from LLM output."""
311 return parse_plan(text)
312
313 # ------------------------------------------------------------------
314 # Step Execution
315 # ------------------------------------------------------------------
316
317 async def _execute_tool_step(
318 self,
319 llm: LLMClientProtocol,
320 plan_step: PlanStep,
321 plan: list[PlanStep],
322 tool_map: dict[str, ToolProtocol],
323 original_message: str,
324 history: list[dict[str, Any]],
325 system_prompt: str,
326 usage: TokenAccumulator | None = None,
327 guard_pipeline: Any = None,
328 ) -> tuple[str, ToolExecutionRecord | None]:
329 """Execute a tool-based plan step."""
330 return await execute_tool_step(
331 llm=llm,
332 plan_step=plan_step,
333 plan=plan,
334 tool_map=tool_map,
335 original_message=original_message,
336 history=history,
337 system_prompt=system_prompt,
338 tool_timeout=self.tool_timeout,
339 llm_timeout=self.llm_timeout,
340 observation_max_chars=self.observation_max_chars,
341 usage=usage,
342 guard_pipeline=guard_pipeline,
343 )
344
345 async def _execute_reasoning_step(
346 self,
347 llm: LLMClientProtocol,
348 plan_step: PlanStep,
349 plan: list[PlanStep],
350 original_message: str,
351 history: list[dict[str, Any]],
352 system_prompt: str,
353 usage: TokenAccumulator | None = None,
354 guard_pipeline: Any = None,
355 ) -> str:
356 """Execute a reasoning-only plan step via LLM."""
357 return await execute_reasoning_step(
358 llm=llm,
359 plan_step=plan_step,
360 plan=plan,
361 original_message=original_message,
362 history=history,
363 system_prompt=system_prompt,
364 llm_timeout=self.llm_timeout,
365 observation_max_chars=self.observation_max_chars,
366 usage=usage,
367 guard_pipeline=guard_pipeline,
368 )
369
370 # ------------------------------------------------------------------
371 # Replanning
372 # ------------------------------------------------------------------
373
374 async def _replan(
375 self,
376 llm: LLMClientProtocol,
377 plan: list[PlanStep],
378 failed_step: PlanStep,
379 error: str,
380 original_message: str,
381 history: list[dict[str, Any]],
382 system_prompt: str,
383 usage: TokenAccumulator | None = None,
384 ) -> list[PlanStep]:
385 """Ask LLM to replan after a step failure."""
386 return await replan(
387 llm=llm,
388 plan=plan,
389 failed_step=failed_step,
390 error=error,
391 original_message=original_message,
392 history=history,
393 system_prompt=system_prompt,
394 llm_timeout=self.llm_timeout,
395 usage=usage,
396 )
397
398 # ------------------------------------------------------------------
399 # Synthesis
400 # ------------------------------------------------------------------
401
402 async def _synthesize(
403 self,
404 llm: LLMClientProtocol,
405 original_message: str,
406 plan: list[PlanStep],
407 history: list[dict[str, Any]],
408 system_prompt: str,
409 usage: TokenAccumulator | None = None,
410 ) -> str:
411 """Synthesize final answer from all completed step results."""
412 return await synthesize(
413 llm=llm,
414 original_message=original_message,
415 plan=plan,
416 history=history,
417 system_prompt=system_prompt,
418 llm_timeout=self.llm_timeout,
419 usage=usage,
420 )
421
422 async def _direct_synthesis(
423 self,
424 llm: LLMClientProtocol,
425 message: str,
426 initial_response: str,
427 history: list[dict[str, Any]],
428 system_prompt: str,
429 steps: list[ReasoningStep],
430 tool_calls: list[ToolExecutionRecord],
431 start_time: float,
432 usage: TokenAccumulator,
433 ) -> Result[AgentResponse, AgentError]:
434 """Fallback when no plan could be parsed — treat as direct response."""
435 final = self._extract_final_answer(initial_response)
436 answer = final if final else initial_response
437
438 steps.append(
439 ReasoningStep(
440 step_number=1,
441 thought="No structured plan generated — using direct response",
442 action="direct",
443 observation=answer,
444 )
445 )
446
447 elapsed = (time.monotonic() - start_time) * 1000
448 return Ok(
449 AgentResponse(
450 message=answer,
451 steps=steps,
452 tool_calls=tool_calls,
453 total_tokens=usage.total_tokens,
454 prompt_tokens=usage.prompt_tokens,
455 completion_tokens=usage.completion_tokens,
456 duration_ms=elapsed,
457 metadata={"strategy": "plan_and_execute", "direct_response": True},
458 )
459 )
460
461 # ------------------------------------------------------------------
462 # Tool Execution
463 # ------------------------------------------------------------------
464
465 async def _run_tool(
466 self,
467 tool_name: str,
468 tool_args: dict[str, Any],
469 tool_map: dict[str, ToolProtocol],
470 ) -> ToolExecutionRecord:
471 """Execute a tool with timeout and error recovery."""
472 return await run_tool(
473 tool_name=tool_name,
474 tool_args=tool_args,
475 tool_map=tool_map,
476 timeout=self.tool_timeout,
477 )
478
479 # ------------------------------------------------------------------
480 # LLM Interaction
481 # ------------------------------------------------------------------
482
483 async def _call_llm(
484 self,
485 llm: LLMClientProtocol,
486 original_message: str,
487 history: list[dict[str, Any]],
488 prompt: str,
489 usage: TokenAccumulator | None = None,
490 ) -> str | None:
491 """Call the LLM and return text content, or ``None`` on failure."""
492 from lexigram.ai.agents.strategies.parsing import build_chat_messages_from_dict
493
494 messages = build_chat_messages_from_dict(original_message, history, prompt)
495 return await call_llm(llm, messages, timeout=self.llm_timeout, usage=usage)
496
497 # ------------------------------------------------------------------
498 # Parsing Helpers
499 # ------------------------------------------------------------------
500
501 @staticmethod
502 def _extract_step_result(text: str) -> str | None:
503 """Extract STEP_RESULT from LLM response."""
504 return extract_step_result(text)
505
506 @staticmethod
507 def _extract_final_answer(text: str) -> str:
508 """Extract FINAL_ANSWER marker, falling back to raw text."""
509 return extract_final_answer(text)
510
511 # ------------------------------------------------------------------
512 # Message / Plan Formatting
513 # ------------------------------------------------------------------
514
515 @staticmethod
516 def _format_plan(plan: list[PlanStep]) -> str:
517 """Format plan steps as numbered list."""
518 return format_plan(plan)
519
520 @staticmethod
521 def _format_completed_steps(plan: list[PlanStep]) -> str:
522 """Format completed steps with their results."""
523 return format_completed_steps(plan)
524
525
526__all__ = ["PlanAndExecuteStrategy", "PlanStep", "PlanStepStatus"]