« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Async execution helpers for the Plan-and-Execute strategy.
3Provides module-level async functions that handle all I/O-bound operations:
4- Calling the LLM with timeout and error recovery.
5- Running tools with timeout and error recovery.
6- Executing a single plan step (tool-based or reasoning-based).
7- Replanning after step failures.
8- Synthesizing a final answer from all completed steps.
9"""
11from __future__ import annotations
13import asyncio
14import time
15from typing import TYPE_CHECKING, Any, cast
17from lexigram.ai.agents.strategies.parsing import (
18 build_chat_messages_from_dict,
19 extract_tool_call,
20)
21from lexigram.ai.agents.strategies.plan_execute_planner import (
22 extract_final_answer,
23 extract_step_result,
24 format_completed_steps,
25 format_plan,
26 parse_plan,
27)
28from lexigram.ai.agents.strategies.plan_execute_types import (
29 EXECUTION_PROMPT,
30 REPLAN_PROMPT,
31 SYNTHESIS_PROMPT,
32 PlanStep,
33 PlanStepStatus,
34)
35from lexigram.ai.agents.strategies.token_utils import TokenAccumulator
36from lexigram.ai.agents.types import ToolExecutionRecord
37from lexigram.logging import (
38 get_logger,
39)
41if TYPE_CHECKING:
42 from lexigram.contracts.ai.agents import ToolProtocol
43 from lexigram.contracts.ai.llm import ChatMessage, LLMClientProtocol
45logger = get_logger(__name__)
48# ---------------------------------------------------------------------------
49# LLM interaction
50# ---------------------------------------------------------------------------
53async def call_llm(
54 llm: LLMClientProtocol,
55 messages: list[ChatMessage],
56 *,
57 timeout: float,
58 usage: TokenAccumulator | None = None,
59) -> str | None:
60 """Call the LLM and return text content, or ``None`` on failure.
62 Args:
63 llm: LLM client implementing ``LLMClientProtocol``.
64 messages: Chat messages to send.
65 timeout: Maximum seconds to wait for a response.
66 usage: Optional accumulator to count tokens from the completion.
68 Returns:
69 Text content of the completion, or ``None`` if the call failed.
70 """
71 try:
72 result = await asyncio.wait_for(
73 llm.complete(cast("list[Any]", messages)),
74 timeout=timeout,
75 )
76 except TimeoutError:
77 logger.warning("plan_execute_llm_timeout", timeout=timeout)
78 return None
79 except (OSError, ConnectionError, RuntimeError, ValueError) as exc:
80 logger.warning("plan_execute_llm_error", error=str(exc))
81 return None
83 if not result.is_ok():
84 logger.warning("plan_execute_llm_err", error=str(result.unwrap_err()))
85 return None
87 completion = result.unwrap()
88 if usage is not None:
89 usage.add(completion)
90 return completion.content if hasattr(completion, "content") else str(completion)
93# ---------------------------------------------------------------------------
94# Tool execution
95# ---------------------------------------------------------------------------
98async def run_tool(
99 tool_name: str,
100 tool_args: dict[str, Any],
101 tool_map: dict[str, ToolProtocol],
102 *,
103 timeout: float,
104) -> ToolExecutionRecord:
105 """Execute a named tool with timeout and error recovery.
107 Args:
108 tool_name: Name of the tool to execute.
109 tool_args: Keyword arguments forwarded to ``tool.execute()``.
110 tool_map: Mapping of tool names to tool instances.
111 timeout: Maximum seconds to allow the tool to run.
113 Returns:
114 A ``ToolExecutionRecord`` capturing the outcome (success or failure).
115 """
116 if tool_name not in tool_map:
117 return ToolExecutionRecord(
118 tool_name=tool_name,
119 arguments=tool_args,
120 error=f"Unknown tool: {tool_name}. Available: {list(tool_map)}",
121 )
123 tool = tool_map[tool_name]
124 start = time.monotonic()
126 try:
127 output = await asyncio.wait_for(
128 tool.execute(**tool_args),
129 timeout=timeout,
130 )
131 duration = (time.monotonic() - start) * 1000
132 return ToolExecutionRecord(
133 tool_name=tool_name,
134 arguments=tool_args,
135 result=output,
136 duration_ms=duration,
137 )
138 except TimeoutError:
139 duration = (time.monotonic() - start) * 1000
140 return ToolExecutionRecord(
141 tool_name=tool_name,
142 arguments=tool_args,
143 error=f"Tool '{tool_name}' timed out after {timeout}s",
144 duration_ms=duration,
145 )
146 except (RuntimeError, TypeError, ValueError, OSError, LookupError) as exc:
147 duration = (time.monotonic() - start) * 1000
148 return ToolExecutionRecord(
149 tool_name=tool_name,
150 arguments=tool_args,
151 error=f"Tool '{tool_name}' failed: {exc}",
152 duration_ms=duration,
153 )
156# ---------------------------------------------------------------------------
157# Step execution
158# ---------------------------------------------------------------------------
161async def execute_tool_step(
162 llm: LLMClientProtocol,
163 plan_step: PlanStep,
164 plan: list[PlanStep],
165 tool_map: dict[str, ToolProtocol],
166 original_message: str,
167 history: list[dict[str, Any]],
168 system_prompt: str,
169 *,
170 tool_timeout: float,
171 llm_timeout: float,
172 observation_max_chars: int,
173 usage: TokenAccumulator | None = None,
174 guard_pipeline: Any = None,
175) -> tuple[str, ToolExecutionRecord | None]:
176 """Execute a tool-based plan step.
178 Uses the LLM to determine tool arguments, then executes the tool.
180 Args:
181 usage: Optional accumulator to count tokens from the LLM call.
183 Returns:
184 A ``(result_text, tool_record)`` tuple. ``tool_record`` is ``None``
185 if the tool could not be invoked.
186 """
187 completed = format_completed_steps(plan)
188 exec_prompt = EXECUTION_PROMPT.format(
189 step_number=plan_step.number,
190 plan_text=format_plan(plan),
191 completed_steps=completed or "(none yet)",
192 step_description=plan_step.description,
193 )
194 messages = build_chat_messages_from_dict(
195 original_message,
196 history,
197 system_prompt + exec_prompt,
198 )
200 llm_text = await call_llm(llm, messages, timeout=llm_timeout, usage=usage)
201 if llm_text is None:
202 plan_step.status = PlanStepStatus.FAILED
203 return "LLM returned empty response", None
205 # Parse tool call from LLM response
206 tool_name, tool_args = extract_tool_call(llm_text)
207 if tool_name is None:
208 tool_name = plan_step.tool_name
209 if not tool_args:
210 tool_args = plan_step.tool_args or {}
212 # Execute the tool
213 if tool_name and tool_name in tool_map:
214 record = await run_tool(tool_name, tool_args, tool_map, timeout=tool_timeout)
215 if record.succeeded:
216 result_text = str(record.result)
217 plan_step.status = PlanStepStatus.COMPLETED
218 else:
219 result_text = f"Error: {record.error}"
220 plan_step.status = PlanStepStatus.FAILED
222 # Guard before truncation so detectors see the full content
223 from lexigram.ai.agents.strategies.guard_hook import guard_observation
225 result_text = await guard_observation(
226 guard_pipeline, result_text, tool_name=tool_name
227 )
228 if len(result_text) > observation_max_chars:
229 result_text = result_text[:observation_max_chars] + "\n[TRUNCATED]"
230 return result_text, record
232 return f"Tool '{tool_name}' not found", None
235async def execute_reasoning_step(
236 llm: LLMClientProtocol,
237 plan_step: PlanStep,
238 plan: list[PlanStep],
239 original_message: str,
240 history: list[dict[str, Any]],
241 system_prompt: str,
242 *,
243 llm_timeout: float,
244 observation_max_chars: int,
245 usage: TokenAccumulator | None = None,
246 guard_pipeline: Any = None,
247) -> str:
248 """Execute a reasoning-only plan step via the LLM.
250 Returns:
251 The guarded observation text — the extracted ``STEP_RESULT:``
252 value, or the raw LLM response when no marker is found — truncated
253 to ``observation_max_chars``. Args are the same as
254 :func:`execute_tool_step` plus ``usage``.
255 """
256 completed = format_completed_steps(plan)
257 exec_prompt = EXECUTION_PROMPT.format(
258 step_number=plan_step.number,
259 plan_text=format_plan(plan),
260 completed_steps=completed or "(none yet)",
261 step_description=plan_step.description,
262 )
263 messages = build_chat_messages_from_dict(
264 original_message,
265 history,
266 system_prompt + exec_prompt,
267 )
269 llm_text = await call_llm(llm, messages, timeout=llm_timeout, usage=usage)
270 if llm_text is None:
271 return "(LLM returned empty response)"
273 result = extract_step_result(llm_text)
274 result_text = result if result else llm_text
276 # Guard before truncation so detectors see the full content
277 from lexigram.ai.agents.strategies.guard_hook import guard_observation
279 result_text = await guard_observation(
280 guard_pipeline, result_text, tool_name="reasoning"
281 )
282 if len(result_text) > observation_max_chars:
283 result_text = result_text[:observation_max_chars]
284 return result_text
287# ---------------------------------------------------------------------------
288# Replanning
289# ---------------------------------------------------------------------------
292async def replan(
293 llm: LLMClientProtocol,
294 plan: list[PlanStep],
295 failed_step: PlanStep,
296 error: str,
297 original_message: str,
298 history: list[dict[str, Any]],
299 system_prompt: str,
300 *,
301 llm_timeout: float,
302 usage: TokenAccumulator | None = None,
303) -> list[PlanStep]:
304 """Ask the LLM to replan after a step failure.
306 Args:
307 llm: LLM client.
308 plan: The current plan (may include already-completed steps).
309 failed_step: The step that failed.
310 error: The error message from the failed step.
311 original_message: The original user message.
312 history: Conversation history.
313 system_prompt: System prompt to prepend.
314 llm_timeout: Timeout for the LLM call.
315 usage: Optional accumulator to count tokens from the LLM call.
317 Returns:
318 A list of new ``PlanStep`` objects renumbered from after ``failed_step``.
319 Empty list if replanning failed.
320 """
321 completed = format_completed_steps(plan)
322 remaining = "\n".join(
323 f"{s.number}. {s.description}"
324 for s in plan
325 if s.number > failed_step.number and s.status == PlanStepStatus.PENDING
326 )
328 replan_text = REPLAN_PROMPT.format(
329 failed_step=failed_step.number,
330 error=error,
331 plan_text=format_plan(plan),
332 completed_steps=completed or "(none)",
333 remaining_steps=remaining or "(none)",
334 )
335 messages = build_chat_messages_from_dict(
336 original_message,
337 history,
338 system_prompt + replan_text,
339 )
341 llm_text = await call_llm(llm, messages, timeout=llm_timeout, usage=usage)
342 if llm_text is None:
343 return []
345 new_steps = parse_plan(llm_text)
346 # Re-number from after the failed step
347 offset = failed_step.number
348 for i, step in enumerate(new_steps):
349 step.number = offset + i + 1
351 logger.info(
352 "plan_execute_replan",
353 failed_step=failed_step.number,
354 new_steps=len(new_steps),
355 )
356 return new_steps
359# ---------------------------------------------------------------------------
360# Synthesis
361# ---------------------------------------------------------------------------
364async def synthesize(
365 llm: LLMClientProtocol,
366 original_message: str,
367 plan: list[PlanStep],
368 history: list[dict[str, Any]],
369 system_prompt: str,
370 *,
371 llm_timeout: float,
372 usage: TokenAccumulator | None = None,
373) -> str:
374 """Synthesize a final answer from all completed step results.
376 Calls the LLM with a synthesis prompt. Falls back to concatenating
377 step results if the LLM call fails. ``usage`` is an optional
378 accumulator for counting tokens from the LLM call.
379 """
380 all_results = "\n".join(
381 f"Step {s.number} ({s.status}): {s.description}\n Result: {s.result or '(no result)'}"
382 for s in plan
383 if s.status != "skipped"
384 )
386 synthesis_prompt = SYNTHESIS_PROMPT.format(
387 original_task=original_message,
388 all_results=all_results,
389 )
390 messages = build_chat_messages_from_dict(
391 original_message,
392 history,
393 system_prompt + synthesis_prompt,
394 )
396 llm_text = await call_llm(llm, messages, timeout=llm_timeout, usage=usage)
397 if llm_text is None:
398 return (
399 "\n".join(
400 s.result
401 for s in plan
402 if s.result and s.status == PlanStepStatus.COMPLETED
403 )
404 or "Unable to synthesize a final answer."
405 )
407 final = extract_final_answer(llm_text)
408 return final if final else llm_text
411__all__ = [
412 "call_llm",
413 "execute_reasoning_step",
414 "execute_tool_step",
415 "replan",
416 "run_tool",
417 "synthesize",
418]