Spaces:
Runtime error
Runtime error
File size: 4,217 Bytes
44de0e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import json
import time
from collections.abc import Iterator
from _core.llm import LLMClient
from _core.models import estimate_cost
from _core.tools import ToolRegistry
from _core.tracer import Step
PLANNER_PROMPT = (
"You are a planner. Given a task, output a JSON array of 2-5 short step "
"strings that would accomplish it. Output ONLY the JSON array."
)
EXECUTOR_PROMPT = "You are an executor. Complete the given step concisely. Use a tool if it helps."
REFLECTOR_PROMPT = (
"You are a reviewer. Given the task and the results so far, decide if the task "
'is complete. Reply with JSON {"done": bool, "feedback": str}. Output ONLY JSON.'
)
class PlanExecuteAgent:
def __init__(self, llm: LLMClient, tools: ToolRegistry, max_rounds: int = 2):
self.llm = llm
self.tools = tools
self.max_rounds = max_rounds
def _ask(self, system: str, user: str, use_tools: bool = False):
start = time.monotonic()
resp = self.llm.chat(
[{"role": "system", "content": system}, {"role": "user", "content": user}],
tools=self.tools.to_openai_schema() if use_tools else None,
)
latency = int((time.monotonic() - start) * 1000)
cost = estimate_cost(self.llm.model, resp.prompt_tokens, resp.completion_tokens)
return resp, latency, cost
def run(self, task: str) -> Iterator[Step]:
feedback = ""
results: list[str] = []
for _ in range(self.max_rounds):
planner_input = f"Task: {task}" + (f"\nPrior feedback: {feedback}" if feedback else "")
resp, lat, cost = self._ask(PLANNER_PROMPT, planner_input)
try:
plan = json.loads(resp.content or "[]")
except json.JSONDecodeError:
plan = [resp.content or ""]
yield Step(
kind="thought",
content="Plan:\n" + "\n".join(f"{i + 1}. {s}" for i, s in enumerate(plan)),
tokens=resp.prompt_tokens + resp.completion_tokens,
cost_usd=cost,
latency_ms=lat,
)
results = []
for s in plan:
eresp, elat, ecost = self._ask(
EXECUTOR_PROMPT, f"Task: {task}\nStep: {s}", use_tools=True
)
if eresp.tool_calls:
for tc in eresp.tool_calls:
yield Step(kind="action", content=f"{tc['name']}({tc['arguments']})")
try:
args = json.loads(tc["arguments"]) if tc["arguments"] else {}
except json.JSONDecodeError as e:
obs = f"Error: invalid tool arguments JSON: {e}"
else:
obs = self.tools.execute(tc["name"], args)
yield Step(kind="observation", content=obs)
results.append(f"{s}: {obs}")
else:
yield Step(kind="action", content=f"Execute: {s}")
yield Step(
kind="observation",
content=eresp.content or "",
tokens=eresp.prompt_tokens + eresp.completion_tokens,
cost_usd=ecost,
latency_ms=elat,
)
results.append(f"{s}: {eresp.content}")
rresp, rlat, rcost = self._ask(
REFLECTOR_PROMPT, f"Task: {task}\nResults:\n" + "\n".join(results)
)
try:
verdict = json.loads(rresp.content or "{}")
except json.JSONDecodeError:
verdict = {"done": True, "feedback": ""}
yield Step(
kind="thought",
content=f"Reflection (done={verdict.get('done')}): {verdict.get('feedback', '')}",
tokens=rresp.prompt_tokens + rresp.completion_tokens,
cost_usd=rcost,
latency_ms=rlat,
)
if verdict.get("done"):
break
feedback = verdict.get("feedback", "")
yield Step(kind="final", content="\n".join(results))
|