from __future__ import annotations import json import re import time from dataclasses import asdict from pathlib import Path from .config import AgentConfig, ensure_workspace from .executor import WorkspaceExecutor from .llm import OpenAICompatibleChatClient from .prompts import SYSTEM_PROMPT, build_task_prompt from .registry import ResourceRegistry from .retriever import ResourceRetriever from .schema import ReActStep, Resource, RunResult, TaskSpec _EXEC_RE = re.compile(r"python|shell)[\"']\s*>(?P.*?)", re.S) _SOLUTION_RE = re.compile(r"(?P.*?)", re.S) class BiomniReActAgent: """Biomni-inspired resource retrieval plus ReAct execution loop.""" def __init__( self, config: AgentConfig | None = None, registry: ResourceRegistry | None = None, ) -> None: self.config = config or AgentConfig() self.registry = registry or ResourceRegistry() def run(self, task: TaskSpec) -> RunResult: run_started = time.perf_counter() workspace = ensure_workspace(task.workspace) all_resources = self.registry.all() selected = self._retrieve(task, all_resources) artifact_paths = { "retrieval_plan": workspace / "retrieval_plan.json", "trace": workspace / "trace.jsonl", "final_answer": workspace / "final_answer.txt", "summary": workspace / "run_summary.json", } self._write_json(artifact_paths["retrieval_plan"], [asdict(resource) for resource in selected]) initial_prompt = build_task_prompt(task, selected) context_tokens = count_tokens(SYSTEM_PROMPT + "\n" + initial_prompt, self.config.model) metrics = { "selected_tools": len(selected), "selected_tool_names": [resource.name for resource in selected], "available_tools": len(all_resources), "context_tokens": context_tokens, "planning_latency_s": 0.0, "total_runtime_s": 0.0, "overhead_planning_ratio": 0.0, } if not self.config.api_key: final = ( "No BIOMNI_REACT_API_KEY was configured. " "Resource retrieval completed; configure an API key to run the ReAct loop." ) result = RunResult( task_name=task.name, success=False, final_answer=final, steps=[], selected_resources=selected, artifact_paths=artifact_paths, error="missing_api_key", ) metrics["total_runtime_s"] = time.perf_counter() - run_started self._persist_result(result, metrics) return result client = OpenAICompatibleChatClient(self.config.api_key, self.config.base_url, self.config.model) executor = WorkspaceExecutor(workspace, timeout_s=self.config.command_timeout_s) messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": initial_prompt}, ] steps: list[ReActStep] = [] final_answer = "" error = None for iteration in range(1, self.config.max_iterations + 1): planning_started = time.perf_counter() response = client.complete(messages) metrics["planning_latency_s"] += time.perf_counter() - planning_started solution = _SOLUTION_RE.search(response) action = _EXEC_RE.search(response) step = ReActStep(iteration=iteration, thought=response.strip()) if solution: final_answer = solution.group("body").strip() steps.append(step) self._append_trace(artifact_paths["trace"], step) break if action: kind = action.group("kind") body = action.group("body").strip() observation = executor.run_python(body) if kind == "python" else executor.run_shell(body) step.action_type = kind step.action = body step.observation = observation messages.append({"role": "assistant", "content": response}) messages.append({"role": "user", "content": f"Observation:\n{observation}"}) else: messages.append({"role": "assistant", "content": response}) messages.append( { "role": "user", "content": "Continue. Use an execute block for an action or a solution block to finish.", } ) steps.append(step) self._append_trace(artifact_paths["trace"], step) if task.expected_outputs and all(path.exists() and path.stat().st_size > 0 for path in task.expected_outputs): final_answer = ( "Expected output file(s) were generated: " + ", ".join(str(path) for path in task.expected_outputs) ) break if iteration >= max(1, self.config.max_iterations - 2): messages.append( { "role": "user", "content": ( "You are near the iteration limit. Stop exploring. " "Write the expected output file now, verify it exists, and finish with ...." ), } ) else: error = "max_iterations_reached" final_answer = "The agent stopped after reaching the iteration limit." result = RunResult( task_name=task.name, success=error is None and bool(final_answer), final_answer=final_answer, steps=steps, selected_resources=selected, artifact_paths=artifact_paths, error=error, ) metrics["total_runtime_s"] = time.perf_counter() - run_started if metrics["total_runtime_s"] > 0: metrics["overhead_planning_ratio"] = metrics["planning_latency_s"] / metrics["total_runtime_s"] self._persist_result(result, metrics) return result def _retrieve(self, task: TaskSpec, resources: list[Resource]) -> list[Resource]: query = " ".join([task.name, task.objective, " ".join(task.constraints), " ".join(task.metadata.values())]) retriever = ResourceRetriever(resources) return retriever.retrieve(query, self.config.retrieval_top_k) def _persist_result(self, result: RunResult, metrics: dict[str, object] | None = None) -> None: result.artifact_paths["final_answer"].write_text(result.final_answer + "\n", encoding="utf-8") summary = { "task_name": result.task_name, "success": result.success, "error": result.error, "steps": len(result.steps), "selected_resources": [resource.name for resource in result.selected_resources], "artifacts": {key: str(path) for key, path in result.artifact_paths.items()}, "metrics": metrics or {}, } self._write_json(result.artifact_paths["summary"], summary) @staticmethod def _write_json(path: Path, payload: object) -> None: path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") @staticmethod def _append_trace(path: Path, step: ReActStep) -> None: with path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(asdict(step), ensure_ascii=False) + "\n") def count_tokens(text: str, model: str | None = None) -> int: try: import tiktoken try: encoding = tiktoken.encoding_for_model(model or "gpt-4o-mini") except Exception: encoding = tiktoken.get_encoding("cl100k_base") return len(encoding.encode(text)) except Exception: return len(re.findall(r"\S+", text))