1"""Synchronous planning helpers for the Plan-and-Execute strategy.
2
3Provides pure-function utilities for:
4- Parsing a numbered plan from raw LLM text.
5- Formatting a plan and its completed steps for LLM prompts.
6- Extracting structured markers (``STEP_RESULT:``, ``FINAL_ANSWER:``) from
7 LLM responses.
8
9All functions in this module are synchronous and side-effect free — they
10operate on plain data types and do not call the LLM or any I/O.
11"""
12
13from __future__ import annotations
14
15import re
16
17from lexigram.ai.agents.strategies.plan_execute_types import PlanStep, PlanStepStatus
18
19# ---------------------------------------------------------------------------
20# Plan parsing
21# ---------------------------------------------------------------------------
22
23
24def parse_plan(text: str) -> list[PlanStep]:
25 """Parse a numbered plan from LLM output.
26
27 Expected format::
28
29 PLAN:
30 1. [TOOL:search] Search for revenue data
31 2. [REASON] Analyze the search results
32
33 Returns at most 10 steps (hard cap for safety).
34 """
35 plan: list[PlanStep] = []
36 in_plan = False
37
38 for line in text.split("\n"):
39 stripped = line.strip()
40 if stripped.upper().startswith("PLAN:"):
41 in_plan = True
42 continue
43
44 if not in_plan:
45 continue
46
47 # Match numbered steps: "1. [TOOL:name] description" or "1. [REASON] description"
48 match = re.match(
49 r"(\d+)\.\s*(?:\[TOOL:(\w+)\])?\s*(?:\[REASON\])?\s*(.*)",
50 stripped,
51 )
52 if match:
53 step_num = int(match.group(1))
54 tool_name = match.group(2)
55 description = match.group(3).strip()
56 plan.append(
57 PlanStep(
58 number=step_num,
59 description=description,
60 tool_name=tool_name,
61 )
62 )
63 elif stripped and plan:
64 # Non-matching line after plan started — plan is done
65 if not stripped[0].isdigit():
66 break
67
68 return plan[:10] # Cap at max steps for safety
69
70
71# ---------------------------------------------------------------------------
72# Plan formatting
73# ---------------------------------------------------------------------------
74
75
76def format_plan(plan: list[PlanStep]) -> str:
77 """Format plan steps as a numbered list for LLM prompts."""
78 lines = []
79 for step in plan:
80 prefix = f"[TOOL:{step.tool_name}]" if step.tool_name else "[REASON]"
81 status = (
82 f" ({step.status.value if isinstance(step.status, PlanStepStatus) else str(step.status)})"
83 if step.status != PlanStepStatus.PENDING
84 else ""
85 )
86 lines.append(f"{step.number}. {prefix} {step.description}{status}")
87 return "\n".join(lines)
88
89
90def format_completed_steps(plan: list[PlanStep]) -> str:
91 """Format completed steps with their results for LLM prompts."""
92 completed = [s for s in plan if s.status == PlanStepStatus.COMPLETED and s.result]
93 if not completed:
94 return ""
95 return "\n".join(
96 f"Step {s.number}: {s.description}\n → {s.result}" for s in completed
97 )
98
99
100# ---------------------------------------------------------------------------
101# Marker extraction
102# ---------------------------------------------------------------------------
103
104
105def extract_step_result(text: str) -> str | None:
106 """Extract the ``STEP_RESULT:`` marker from an LLM response.
107
108 Returns the text following the marker, or ``None`` if not present.
109 """
110 marker = "STEP_RESULT:"
111 idx = text.upper().find(marker)
112 if idx == -1:
113 return None
114 return text[idx + len(marker) :].strip()
115
116
117def extract_final_answer(text: str) -> str:
118 """Extract the ``FINAL_ANSWER:`` marker, falling back to raw text.
119
120 Returns the text following the marker if present, otherwise the full
121 stripped text.
122 """
123 marker = "FINAL_ANSWER:"
124 idx = text.upper().find(marker)
125 if idx == -1:
126 return text.strip()
127 return text[idx + len(marker) :].strip()
128
129
130__all__ = [
131 "extract_final_answer",
132 "extract_step_result",
133 "format_completed_steps",
134 "format_plan",
135 "parse_plan",
136]