Spaces:
Runtime error
Runtime error
| 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)) | |