158 lines
5.6 KiB
Python
158 lines
5.6 KiB
Python
"""Agent evaluation driver: a MESSAGE PUMP, not a thinking framework.
|
|
|
|
We are EVALUATING agents, not being one: the model under test does whatever
|
|
thinking it wants; this loop only (1) hands tool declarations over, (2)
|
|
executes the model's tool_calls against an Environment plugin, (3) feeds
|
|
observations back, (4) stops at max-turns / no-more-calls / env-done.
|
|
Everything is recorded as a trajectory for env_reward scorers.
|
|
"""
|
|
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from ..data.sample import ChatMessage, Sample
|
|
from ..model.adapter import ModelAdapter
|
|
from ..model.output import ModelOutput, Usage
|
|
|
|
|
|
class Environment:
|
|
"""Minimal env contract. One instance per sample.
|
|
|
|
Two execution models, auto-selected by the runner:
|
|
- MESSAGE PUMP (default): implement reset/step/final_state; the runner
|
|
drives the model turn-by-turn (bfcl_mock style).
|
|
- SELF-RUNNING (official engines): implement run_task(adapter, sample);
|
|
the whole simulation happens inside (official tau2/swe bundles), and
|
|
a prediction dict is returned directly. Falls back to the pump when
|
|
run_task is not overridden.
|
|
"""
|
|
|
|
name = 'base'
|
|
needs_adapter = False # self-running engines set True; the runner passes
|
|
# the adapter to the constructor (no name hardcoding)
|
|
|
|
def __init__(self, adapter=None):
|
|
self.adapter = adapter
|
|
|
|
def reset(self, sample: Sample) -> List[ChatMessage]:
|
|
return []
|
|
|
|
async def step(self, tool_calls: List[Any], text: str, sample: Sample
|
|
) -> List[ChatMessage]:
|
|
raise NotImplementedError
|
|
|
|
def final_state(self) -> Dict[str, Any]:
|
|
return {}
|
|
|
|
async def run_task(self, adapter, sample: Sample, **kw) -> Optional[Dict[str, Any]]:
|
|
"""Self-running hook. Return a prediction dict (raw/trajectory/
|
|
env_state/usage/group_key) or None to fall back to the pump."""
|
|
return None
|
|
|
|
|
|
# ---------------- environment registry (万物皆可插件: envs too) ----------------
|
|
|
|
from ..eval.registry import EvalRegistry # noqa: E402
|
|
|
|
ENV_REGISTRY = EvalRegistry('environment')
|
|
|
|
|
|
def register_env(name: str):
|
|
"""Class decorator: @register_env('bfcl_mock'). One env per benchmark
|
|
family; drop a module in agent/envs/ and it auto-registers."""
|
|
|
|
def decorator(cls):
|
|
ENV_REGISTRY.register(name, cls)
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
|
|
def get_env(name: str, adapter=None):
|
|
cls = ENV_REGISTRY.get(name)
|
|
return cls(adapter=adapter)
|
|
|
|
|
|
class Trajectory:
|
|
"""Recorded turns: role-tagged messages + per-turn usage."""
|
|
|
|
def __init__(self):
|
|
self.messages: List[Dict[str, Any]] = []
|
|
self.turns = 0
|
|
self.usage = Usage()
|
|
self.env_state: Dict[str, Any] = {}
|
|
|
|
def add(self, role: str, content: str, **extra) -> None:
|
|
entry: Dict[str, Any] = {'role': role, 'content': content, 'turn': self.turns}
|
|
entry.update(extra)
|
|
self.messages.append(entry)
|
|
|
|
|
|
async def drive(
|
|
adapter: ModelAdapter,
|
|
sample: Sample,
|
|
env: Optional[Environment] = None,
|
|
max_turns: int = 8,
|
|
system: str = '',
|
|
) -> Trajectory:
|
|
"""Run one sample through the model (+env if present).
|
|
|
|
No env -> single-turn fast path (one generate, done).
|
|
With env-> pump until the model stops calling tools / env says done /
|
|
max_turns reached.
|
|
"""
|
|
traj = Trajectory()
|
|
base: List[ChatMessage] = []
|
|
if system:
|
|
base.append(ChatMessage(role='system', content=system))
|
|
base += ([ChatMessage(role='user', content=sample.input)] if isinstance(sample.input, str)
|
|
else list(sample.input))
|
|
if env is not None:
|
|
base += env.reset(sample)
|
|
if getattr(adapter, 'name', '') == 'mock' \
|
|
and adapter.extra.get('mode') in ('boxed', 'oracle', 'fc') \
|
|
and sample.target not in ('', None):
|
|
base = base + [ChatMessage(role='user', content=f'MOCKTARGET::{sample.target}')]
|
|
|
|
tools = [{'name': t.name, 'description': t.description or '',
|
|
'parameters': t.parameters} for t in (sample.tools or [])] or None
|
|
|
|
messages = list(base)
|
|
t0 = time.time()
|
|
for turn in range(max_turns + 1):
|
|
out: ModelOutput = await adapter.generate(messages, tools=tools)
|
|
traj.turns = turn + 1
|
|
traj.usage = traj.usage + out.usage
|
|
traj.add('assistant', out.text, tool_calls=[c.model_dump() for c in out.tool_calls])
|
|
messages.append(ChatMessage(role='assistant',
|
|
content=out.text or _calls_text(out.tool_calls)))
|
|
if not out.tool_calls or env is None:
|
|
break # plain answer or single-turn: done
|
|
observations = await env.step(out.tool_calls, out.text, sample)
|
|
for obs in observations:
|
|
traj.add(obs.role, obs.content)
|
|
messages.append(obs)
|
|
if env is not None:
|
|
traj.env_state = env.final_state()
|
|
traj.messages.insert(0, {'role': 'meta', 'content': f'turns={traj.turns} '
|
|
f'latency={time.time() - t0:.1f}s'})
|
|
return traj
|
|
|
|
|
|
def _calls_text(calls) -> str:
|
|
import json
|
|
|
|
return json.dumps([c.to_openai()['function'] for c in calls], ensure_ascii=False)
|
|
|
|
|
|
def trajectory_to_prediction(traj: Trajectory) -> Dict[str, Any]:
|
|
"""Fold a trajectory into the runner's prediction-dict shape."""
|
|
last_assistant = next((m for m in reversed(traj.messages) if m['role'] == 'assistant'), {})
|
|
return {
|
|
'raw': last_assistant.get('content', ''),
|
|
'trajectory': traj.messages,
|
|
'env_state': traj.env_state or None,
|
|
'usage': traj.usage.model_dump(),
|
|
'group_key': str(last_assistant.get('turn', '')),
|
|
}
|