import os
import re
import json
from abc import ABC, abstractmethod
from typing import Dict, Any, List
from concurrent.futures import ThreadPoolExecutor, as_completed
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from .utils import parse_tag
class BaseExplorerAgent(ABC):
item_id_tag: str = "id"
system_prompt: str = "You are a helpful assistant."
def __init__(self, client, model: str, max_iterations: int = 20):
self.client = client
self.model = model
self.max_iterations = max_iterations
@abstractmethod
def get_tools(self) -> List[Dict[str, Any]]:
"""Return the list of tools available to this agent."""
pass
@abstractmethod
def execute_tool(self, tool_name: str, tool_args: Dict[str, Any], iteration: int, context: Dict[str, Any]) -> str:
"""Execute a tool and return the output string."""
pass
@abstractmethod
def format_initial_prompt(self, **kwargs) -> str:
"""Format the initial prompt for the agent."""
pass
@abstractmethod
def get_force_output_message(self) -> str:
"""Return the message to force output when max iterations reached."""
pass
@abstractmethod
def get_item_content(self, item_id: str, context: Dict[str, Any]) -> str:
"""Fetch the content for a supporting item by its ID."""
pass
def normalize_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
"""Convert domain-specific item format to normalized format with 'id' key.
Override this in subclasses that use different identifiers."""
return item
def denormalize_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
"""Convert normalized format with 'id' key back to domain-specific format.
Override this in subclasses that use different identifiers."""
return item
def parse_supporting_items(self, content: str) -> List[Dict[str, Any]]:
"""Parse supporting items from the agent output using domain-specific tag."""
items = []
outer_match = re.search(r'(.*?)', content, re.DOTALL)
if outer_match:
items_content = outer_match.group(1)
item_matches = re.findall(r'- (.*?)
', items_content, re.DOTALL)
for item_match in item_matches:
id_match = re.search(rf'<{self.item_id_tag}>(.*?)', item_match, re.DOTALL)
reasoning_match = re.search(r'(.*?)', item_match, re.DOTALL)
if id_match:
items.append({
self.item_id_tag: id_match.group(1).strip(),
'reasoning': reasoning_match.group(1).strip() if reasoning_match else ''
})
return items
def parse_structured_output(self, content: str) -> Dict[str, Any]:
"""Parse the structured output from the agent."""
return {
"clues": parse_tag(content, "clues"),
"question": parse_tag(content, "question"),
"truth": parse_tag(content, "truth"),
"supporting_items": self.parse_supporting_items(content)
}
def on_tool_result(self, tool_name: str, tool_args: Dict[str, Any], output: str, context: Dict[str, Any]) -> None:
"""Hook called after each tool execution. Override to track visited items, etc."""
pass
def build_tool_trajectory_entry(self, tool_name: str, tool_args: Dict[str, Any], output: str, context: Dict[str, Any]) -> Dict[str, Any]:
"""Build a trajectory entry for a tool call. Override to customize."""
return {
"type": "tool_call",
"tool_name": tool_name,
"arguments": tool_args,
"output": output
}
def run_agent_loop(self, input_messages: List, trajectory: List, context: Dict[str, Any]) -> Dict[str, Any] | None:
"""Run the main agent loop with tools."""
request_body = {
"model": self.model,
"system": self.system_prompt,
"max_tokens": 20000,
"tools": self.get_tools(),
"tool_choice": {"type": "auto"},
"thinking": {"type": "enabled", "budget_tokens": 2000}
}
parsed = None
for i in range(self.max_iterations):
request_body["messages"] = input_messages
response = self.client.messages.create(**request_body)
tool_use_items = [item for item in response.content if getattr(item, 'type', None) == 'tool_use']
thinking_items = [item for item in response.content if getattr(item, 'type', None) == 'thinking']
text_items = [item for item in response.content if getattr(item, 'type', None) == 'text']
if thinking_items:
for thinking_item in thinking_items:
trajectory.append({
"type": "thinking",
"output": thinking_item.thinking
})
if not tool_use_items:
for item in text_items:
if item.type == "text":
content = item.text
parsed = self.parse_structured_output(content)
trajectory.append({
"type": "output_text",
"output": content,
**parsed
})
break
serialized_items = []
for item in response.content:
serialized_item = item.model_dump(mode="python")
if 'status' in serialized_item:
del serialized_item['status']
serialized_items.append(serialized_item)
input_messages.append({"role": "assistant", "content": serialized_items})
tool_results = []
for tool_call in tool_use_items:
tool_args = tool_call.input
tool_name = tool_call.name
output = self.execute_tool(tool_name, tool_args, i, context)
self.on_tool_result(tool_name, tool_args, output, context)
tool_results.append({
"type": "tool_result",
"tool_use_id": tool_call.id,
"content": output
})
trajectory_entry = self.build_tool_trajectory_entry(tool_name, tool_args, output, context)
trajectory.append(trajectory_entry)
input_messages.append({"role": "user", "content": tool_results})
return parsed
def force_output(self, input_messages: List, trajectory: List) -> Dict[str, Any] | None:
"""Force the agent to produce output when max iterations reached."""
input_messages.append({"role": "user", "content": self.get_force_output_message()})
response = self.client.messages.create(
model=self.model,
system=self.system_prompt,
max_tokens=20000,
messages=input_messages,
thinking={"type": "enabled", "budget_tokens": 2000}
)
thinking_items = [item for item in response.content if getattr(item, 'type', None) == 'thinking']
text_items = [item for item in response.content if getattr(item, 'type', None) == 'text']
if thinking_items:
for thinking_item in thinking_items:
trajectory.append({
"type": "thinking",
"output": thinking_item.thinking
})
parsed = None
for item in text_items:
if item.type == "text":
content = item.text
parsed = self.parse_structured_output(content)
trajectory.append({
"type": "forced_output",
"output": content,
**parsed
})
return parsed
def is_processed(self, filepath: str) -> bool:
try:
with open(filepath, "r") as f:
data = json.load(f)
if not data.get("tasks") or len(data["tasks"]) == 0:
return False
task = data["tasks"][0]
if task.get("clues") is None or task.get("question") is None or task.get("truth") is None:
return False
supporting_items = task.get("supporting_items", [])
if len(supporting_items) != 3:
return False
unique_ids = set([item['id'] for item in supporting_items])
if len(unique_ids) != 3:
return False
items_and_contents = task.get("items_and_contents", {})
for item in supporting_items:
item_id = item.get("id", "")
if item_id and item_id not in items_and_contents:
return False
return True
except (json.JSONDecodeError, KeyError, TypeError):
return False
@abstractmethod
def run_single(self, **kwargs) -> Dict[str, Any]:
"""Run a single exploration. Implementation varies by domain."""
pass
def build_result(self, parsed: Dict[str, Any] | None, context: Dict[str, Any], extra_fields: Dict[str, Any] = None) -> Dict[str, Any]:
"""Build the final result dict with normalized items."""
clues = parsed["clues"] if parsed else None
question = parsed["question"] if parsed else None
truth = parsed["truth"] if parsed else None
supporting_items = parsed["supporting_items"] if parsed else []
items_and_contents = {}
normalized_items = []
for item in supporting_items:
item_id = item.get(self.item_id_tag, "")
normalized_item = self.normalize_item(item)
normalized_items.append(normalized_item)
if not item_id:
continue
content = self.get_item_content(item_id, context)
if content:
normalized_id = normalized_item.get("id", item_id)
items_and_contents[normalized_id] = content
task = {
"clues": clues,
"question": question,
"truth": truth,
"supporting_items": normalized_items,
"items_and_contents": items_and_contents
}
if extra_fields:
task.update(extra_fields)
return {"tasks": [task]}
def run_batch(self, seeds: List[Any], output_dir: str, max_workers: int = 8, get_seed_id=None, get_output_path=None) -> Dict[str, Any]:
"""Run batch processing with parallel workers."""
results = []
errors = []
os.makedirs(output_dir, exist_ok=True)
if get_seed_id is None:
get_seed_id = lambda s: str(s)
if get_output_path is None:
get_output_path = lambda s, d: os.path.join(d, f"{get_seed_id(s)}.json")
seeds_to_process = []
for seed in seeds:
output_path = get_output_path(seed, output_dir)
if not os.path.exists(output_path) or not self.is_processed(output_path):
seeds_to_process.append(seed)
skipped = len(seeds) - len(seeds_to_process)
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
) as progress:
task = progress.add_task(
f"Processing {len(seeds_to_process)}/{len(seeds)} explorations",
total=len(seeds_to_process)
)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_seed = {
executor.submit(self.run_single, seed=seed, output_dir=output_dir): seed
for seed in seeds_to_process
}
for future in as_completed(future_to_seed):
seed = future_to_seed[future]
try:
result = future.result()
results.append({"seed": get_seed_id(seed), "status": "success", "result": result})
except Exception as e:
errors.append({"seed": get_seed_id(seed), "error": str(e)})
progress.advance(task)
return {
"total": len(seeds),
"successful": len(results),
"failed": len(errors),
"skipped": skipped,
"errors": errors
}