wehub-resource-sync
276 lines
9.0 KiB
Python
276 lines
9.0 KiB
Python
"""Plan-and-execute agent with replan on failure, plan diffs, and dual budgets.
|
|
|
|
Conceptual references:
|
|
- ./docs/en.md (this lesson)
|
|
- Phase 14 lesson 01 (agent loop fundamentals)
|
|
- Phase 13 lesson 02 (tool protocols overview)
|
|
|
|
Stdlib only. Run: python3 code/main.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable
|
|
|
|
|
|
@dataclass
|
|
class Step:
|
|
id: int
|
|
tool_name: str
|
|
args: dict
|
|
expected_outcome: str
|
|
result: Any | None = None
|
|
error: str | None = None
|
|
|
|
def signature(self) -> tuple:
|
|
return (self.tool_name, json.dumps(self.args, sort_keys=True))
|
|
|
|
|
|
@dataclass
|
|
class PlanDiff:
|
|
revision: int
|
|
removed: list[int]
|
|
added: list[int]
|
|
revised: list[int]
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"revision": self.revision,
|
|
"removed": list(self.removed),
|
|
"added": list(self.added),
|
|
"revised": list(self.revised),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class Event:
|
|
type: str
|
|
payload: dict
|
|
ts: float = field(default_factory=time.time)
|
|
|
|
|
|
@dataclass
|
|
class SessionResult:
|
|
status: str
|
|
reason: str
|
|
history: list[Step]
|
|
revisions: list[PlanDiff]
|
|
events: list[Event]
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"status": self.status,
|
|
"reason": self.reason,
|
|
"history": [
|
|
{"id": s.id, "tool": s.tool_name, "args": s.args,
|
|
"result": s.result, "error": s.error}
|
|
for s in self.history
|
|
],
|
|
"revisions": [r.to_dict() for r in self.revisions],
|
|
"events": [{"type": e.type, "payload": e.payload, "ts": e.ts} for e in self.events],
|
|
}
|
|
|
|
|
|
Planner = Callable[[str, list[Step], str | None], list[Step]]
|
|
ToolExecutor = Callable[[str, dict], Any]
|
|
|
|
|
|
class ToolFailure(Exception):
|
|
pass
|
|
|
|
|
|
def _diff_plans(old: list[Step], new: list[Step], revision: int) -> PlanDiff:
|
|
old_ids = {s.id for s in old}
|
|
new_ids = {s.id for s in new}
|
|
removed = sorted(old_ids - new_ids)
|
|
added = sorted(new_ids - old_ids)
|
|
revised: list[int] = []
|
|
old_by_id = {s.id: s for s in old}
|
|
for s in new:
|
|
if s.id in old_ids and old_by_id[s.id].signature() != s.signature():
|
|
revised.append(s.id)
|
|
return PlanDiff(revision=revision, removed=removed, added=added, revised=revised)
|
|
|
|
|
|
class PlanExecuteAgent:
|
|
"""Sequential plan executor with replan on failure."""
|
|
|
|
def __init__(
|
|
self,
|
|
planner: Planner,
|
|
executor: ToolExecutor,
|
|
*,
|
|
max_steps: int = 12,
|
|
max_replans: int = 5,
|
|
) -> None:
|
|
self._planner = planner
|
|
self._executor = executor
|
|
self.max_steps = max_steps
|
|
self.max_replans = max_replans
|
|
self._events: list[Event] = []
|
|
|
|
def _emit(self, etype: str, payload: dict) -> None:
|
|
self._events.append(Event(type=etype, payload=payload))
|
|
|
|
def run(self, goal: str) -> SessionResult:
|
|
self._events = []
|
|
history: list[Step] = []
|
|
revisions: list[PlanDiff] = []
|
|
steps_taken = 0
|
|
replans_used = 0
|
|
last_error: str | None = None
|
|
|
|
plan = self._planner(goal, history, None)
|
|
self._emit("plan.commit", {"revision": 0, "steps": _summarize(plan)})
|
|
|
|
if not plan:
|
|
self._emit("session.complete", {"reason": "no_plan"})
|
|
return SessionResult(
|
|
status="failed", reason="no_plan",
|
|
history=history, revisions=revisions, events=list(self._events),
|
|
)
|
|
|
|
cursor = 0
|
|
revision = 0
|
|
|
|
while cursor < len(plan):
|
|
if steps_taken >= self.max_steps:
|
|
self._emit("session.complete", {"reason": "step_budget"})
|
|
return SessionResult(
|
|
status="failed", reason="step_budget",
|
|
history=history, revisions=revisions, events=list(self._events),
|
|
)
|
|
|
|
step = plan[cursor]
|
|
self._emit("step.start", {"step_id": step.id, "tool": step.tool_name})
|
|
try:
|
|
step.result = self._executor(step.tool_name, step.args)
|
|
self._emit("step.end", {"step_id": step.id, "outcome": "ok"})
|
|
history.append(step)
|
|
cursor += 1
|
|
steps_taken += 1
|
|
continue
|
|
except Exception as exc:
|
|
step.error = f"{type(exc).__name__}: {exc}"
|
|
self._emit("step.end", {"step_id": step.id, "outcome": "error", "error": step.error})
|
|
history.append(step)
|
|
steps_taken += 1
|
|
last_error = step.error
|
|
|
|
if replans_used >= self.max_replans:
|
|
self._emit("session.complete", {"reason": "replan_budget"})
|
|
return SessionResult(
|
|
status="failed", reason="replan_budget",
|
|
history=history, revisions=revisions, events=list(self._events),
|
|
)
|
|
|
|
replans_used += 1
|
|
revision += 1
|
|
new_plan = self._planner(goal, history, last_error)
|
|
self._emit("plan.draft", {"revision": revision, "steps": _summarize(new_plan)})
|
|
if not new_plan:
|
|
self._emit("session.complete", {"reason": "no_plan"})
|
|
return SessionResult(
|
|
status="failed", reason="no_plan",
|
|
history=history, revisions=revisions, events=list(self._events),
|
|
)
|
|
diff = _diff_plans(plan[cursor:], new_plan, revision)
|
|
revisions.append(diff)
|
|
self._emit("plan.diff", diff.to_dict())
|
|
plan = new_plan
|
|
cursor = 0
|
|
self._emit("plan.commit", {"revision": revision, "steps": _summarize(plan)})
|
|
|
|
self._emit("session.complete", {"reason": "goal_met"})
|
|
return SessionResult(
|
|
status="completed", reason="goal_met",
|
|
history=history, revisions=revisions, events=list(self._events),
|
|
)
|
|
|
|
|
|
def _summarize(plan: list[Step]) -> list[dict]:
|
|
return [{"id": s.id, "tool": s.tool_name, "outcome": s.expected_outcome} for s in plan]
|
|
|
|
|
|
def make_deterministic_planner(fail_step_id: int | None, recovery: str = "route_around") -> Planner:
|
|
"""Planner used in the demo and tests.
|
|
|
|
When ``fail_step_id`` is given, the planner inserts a ``_force_fail`` marker
|
|
into that step's args on the initial plan. Executors that honor the marker
|
|
raise on that step, exercising the replan path. The marker is removed on the
|
|
revised plan so the route-around can succeed.
|
|
"""
|
|
|
|
def planner(goal: str, history: list[Step], last_error: str | None) -> list[Step]:
|
|
if last_error is None:
|
|
initial = [
|
|
Step(1, "fetch", {"key": "input"}, "loaded user input"),
|
|
Step(2, "transform", {"mode": "v1"}, "computed v1 form"),
|
|
Step(3, "render", {}, "rendered output"),
|
|
Step(4, "submit", {}, "submitted to backend"),
|
|
]
|
|
if fail_step_id is not None:
|
|
for s in initial:
|
|
if s.id == fail_step_id:
|
|
s.args = {**s.args, "_force_fail": True}
|
|
return initial
|
|
if recovery == "route_around" and "transform" in last_error:
|
|
return [
|
|
Step(2, "transform", {"mode": "v2"}, "computed via fallback"),
|
|
Step(3, "render", {}, "rendered output"),
|
|
Step(4, "submit", {}, "submitted to backend"),
|
|
]
|
|
if recovery == "give_up":
|
|
return [
|
|
Step(98, "log_failure", {"why": last_error or ""}, "logged failure"),
|
|
Step(99, "notify_user", {}, "told the user"),
|
|
]
|
|
return []
|
|
|
|
return planner
|
|
|
|
|
|
def _demo() -> None:
|
|
counters = {"transform_v1_calls": 0}
|
|
|
|
def executor(tool: str, args: dict) -> Any:
|
|
if args.get("_force_fail"):
|
|
counters["transform_v1_calls"] += 1
|
|
raise ToolFailure(f"{tool} marker-forced failure")
|
|
if tool == "fetch":
|
|
return {"k": "v"}
|
|
if tool == "transform":
|
|
if args.get("mode") == "v1":
|
|
counters["transform_v1_calls"] += 1
|
|
raise ToolFailure("transform v1 backend down")
|
|
return {"ok": True}
|
|
if tool == "render":
|
|
return "html"
|
|
if tool == "submit":
|
|
return {"id": 1}
|
|
if tool in ("log_failure", "notify_user"):
|
|
return "logged"
|
|
raise ToolFailure(f"unknown tool {tool}")
|
|
|
|
agent = PlanExecuteAgent(
|
|
planner=make_deterministic_planner(fail_step_id=2, recovery="route_around"),
|
|
executor=executor,
|
|
max_steps=12, max_replans=5,
|
|
)
|
|
res = agent.run("ship the report")
|
|
print(json.dumps({
|
|
"status": res.status,
|
|
"reason": res.reason,
|
|
"history": [(s.id, s.tool_name, bool(s.error)) for s in res.history],
|
|
"revisions": [r.to_dict() for r in res.revisions],
|
|
"events": [e.type for e in res.events],
|
|
"transform_v1_calls": counters["transform_v1_calls"],
|
|
}, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_demo()
|