Spaces:
Sleeping
Sleeping
File size: 5,453 Bytes
d3d0e0e | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | from __future__ import annotations
import json
import re
from dataclasses import dataclass
from data_agent_baseline.agents.model import ModelAdapter, ModelMessage, ModelStep
from data_agent_baseline.agents.prompt import (
REACT_SYSTEM_PROMPT,
build_observation_prompt,
build_system_prompt,
build_task_prompt,
)
from data_agent_baseline.agents.runtime import AgentRunResult, AgentRuntimeState, StepRecord
from data_agent_baseline.benchmark.schema import PublicTask
from data_agent_baseline.tools.registry import ToolRegistry
@dataclass(frozen=True, slots=True)
class ReActAgentConfig:
max_steps: int = 16
def _strip_json_fence(raw_response: str) -> str:
text = raw_response.strip()
fence_match = re.search(r"```json\s*(.*?)\s*```", text, flags=re.IGNORECASE | re.DOTALL)
if fence_match is not None:
return fence_match.group(1).strip()
generic_fence_match = re.search(r"```\s*(.*?)\s*```", text, flags=re.DOTALL)
if generic_fence_match is not None:
return generic_fence_match.group(1).strip()
return text
def _load_single_json_object(text: str) -> dict[str, object]:
payload, end = json.JSONDecoder().raw_decode(text)
remainder = text[end:].strip()
if remainder:
cleaned_remainder = re.sub(r"(?:\\[nrt])+", "", remainder).strip()
if cleaned_remainder:
raise ValueError("Model response must contain only one JSON object.")
if not isinstance(payload, dict):
raise ValueError("Model response must be a JSON object.")
return payload
def parse_model_step(raw_response: str) -> ModelStep:
normalized = _strip_json_fence(raw_response)
payload = _load_single_json_object(normalized)
thought = payload.get("thought", "")
action = payload.get("action")
action_input = payload.get("action_input", {})
if not isinstance(thought, str):
raise ValueError("thought must be a string.")
if not isinstance(action, str) or not action:
raise ValueError("action must be a non-empty string.")
if not isinstance(action_input, dict):
raise ValueError("action_input must be a JSON object.")
return ModelStep(
thought=thought,
action=action,
action_input=action_input,
raw_response=raw_response,
)
class ReActAgent:
def __init__(
self,
*,
model: ModelAdapter,
tools: ToolRegistry,
config: ReActAgentConfig | None = None,
system_prompt: str | None = None,
) -> None:
self.model = model
self.tools = tools
self.config = config or ReActAgentConfig()
self.system_prompt = system_prompt or REACT_SYSTEM_PROMPT
def _build_messages(self, task: PublicTask, state: AgentRuntimeState) -> list[ModelMessage]:
system_content = build_system_prompt(
self.tools.describe_for_prompt(),
system_prompt=self.system_prompt,
)
messages = [ModelMessage(role="system", content=system_content)]
messages.append(ModelMessage(role="user", content=build_task_prompt(task)))
for step in state.steps:
messages.append(ModelMessage(role="assistant", content=step.raw_response))
messages.append(
ModelMessage(role="user", content=build_observation_prompt(step.observation))
)
return messages
def run(self, task: PublicTask) -> AgentRunResult:
state = AgentRuntimeState()
for step_index in range(1, self.config.max_steps + 1):
raw_response = self.model.complete(self._build_messages(task, state))
try:
model_step = parse_model_step(raw_response)
tool_result = self.tools.execute(task, model_step.action, model_step.action_input)
observation = {
"ok": tool_result.ok,
"tool": model_step.action,
"content": tool_result.content,
}
step_record = StepRecord(
step_index=step_index,
thought=model_step.thought,
action=model_step.action,
action_input=model_step.action_input,
raw_response=raw_response,
observation=observation,
ok=tool_result.ok,
)
state.steps.append(step_record)
if tool_result.is_terminal:
state.answer = tool_result.answer
break
except Exception as exc:
observation = {
"ok": False,
"error": str(exc),
}
state.steps.append(
StepRecord(
step_index=step_index,
thought="",
action="__error__",
action_input={},
raw_response=raw_response,
observation=observation,
ok=False,
)
)
if state.answer is None and state.failure_reason is None:
state.failure_reason = "Agent did not submit an answer within max_steps."
return AgentRunResult(
task_id=task.task_id,
answer=state.answer,
steps=list(state.steps),
failure_reason=state.failure_reason,
)
|