| 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"<execute\s+type=[\"'](?P<kind>python|shell)[\"']\s*>(?P<body>.*?)</execute>", re.S) |
| _SOLUTION_RE = re.compile(r"<solution>(?P<body>.*?)</solution>", 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 <solution>...</solution>." |
| ), |
| } |
| ) |
| 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)) |
|
|
