agentbox-rl 0.2.0__py3-none-any.whl

agentbox/__init__.py

@@ -0,0 +1,52 @@
+"""AgentBox — local agentic virtual environment for trajectory building."""
+
+from agentbox.agent.agent import Agent
+from agentbox.benchmark import (
+    BenchmarkReport,
+    BenchmarkRunConfig,
+    BenchmarkRunner,
+    BenchmarkSuite,
+    ModelUnderTest,
+    compare_reports,
+    create_suite_from_tasks,
+    load_suite,
+)
+from agentbox.config import AgentConfig, ModelConfig, RolloutConfig, SandboxConfig
+from agentbox.metrics import AggregateMetrics, aggregate_trajectories
+from agentbox.runner.parallel import ParallelResult, ParallelRunner
+from agentbox.runner.rollout import Rollout
+from agentbox.sandbox.manager import SandboxManager
+from agentbox.tasks.schema import Task, VerifierSpec
+from agentbox.tools.base import BaseTool
+from agentbox.tools.decorator import tool
+from agentbox.trajectory.schema import Trajectory
+from agentbox.version import __version__
+
+__all__ = [
+    "__version__",
+    "Agent",
+    "AgentConfig",
+    "AggregateMetrics",
+    "BaseTool",
+    "BenchmarkReport",
+    "BenchmarkRunConfig",
+    "BenchmarkRunner",
+    "BenchmarkSuite",
+    "ModelConfig",
+    "ModelUnderTest",
+    "ParallelResult",
+    "ParallelRunner",
+    "Rollout",
+    "RolloutConfig",
+    "SandboxConfig",
+    "SandboxManager",
+    "Task",
+    "Trajectory",
+    "VerifierSpec",
+    "aggregate_trajectories",
+    "compare_reports",
+    "create_suite_from_tasks",
+    "load_suite",
+    "tool",
+]
+

agentbox/agent/__init__.py

@@ -0,0 +1,6 @@
+"""Agent loop and facade."""
+
+from agentbox.agent.agent import Agent
+from agentbox.agent.loop import AgentLoop, AgentLoopResult
+
+__all__ = ["Agent", "AgentLoop", "AgentLoopResult"]

agentbox/agent/agent.py

@@ -0,0 +1,59 @@
+"""High-level Agent facade."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from agentbox.config import AgentConfig, ModelConfig
+from agentbox.model.base import ModelClient
+from agentbox.model.openai_compat import OpenAICompatClient
+from agentbox.tools.base import BaseTool
+from agentbox.types import ToolMode
+
+
+class Agent:
+    """Ergonomic agent configuration for rollouts."""
+
+    def __init__(
+        self,
+        model: str | ModelConfig | ModelClient,
+        *,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        tools: ToolMode | list[str] | str = ToolMode.STRUCTURED,
+        custom_tools: Sequence[BaseTool] | None = None,
+        builtins: bool = True,
+        max_steps: int = 40,
+        temperature: float = 0.7,
+        **kwargs: Any,
+    ) -> None:
+        if isinstance(model, str):
+            self.model_config = ModelConfig(
+                model=model,
+                base_url=base_url,
+                api_key=api_key,
+                temperature=temperature,
+            )
+            self.model_client: ModelClient = OpenAICompatClient(self.model_config)
+        elif isinstance(model, ModelConfig):
+            self.model_config = model
+            self.model_client = OpenAICompatClient(model)
+        else:
+            self.model_config = ModelConfig(model=getattr(model, "model", "unknown"))
+            self.model_client = model
+
+        tool_mode: ToolMode | list[str]
+        if isinstance(tools, str) and tools in {m.value for m in ToolMode}:
+            tool_mode = ToolMode(tools)
+        else:
+            tool_mode = tools  # type: ignore[assignment]
+
+        agent_kwargs = {
+            "tools": tool_mode,
+            "custom_tools": list(custom_tools or []),
+            "builtins": builtins,
+            "max_steps": max_steps,
+            **kwargs,
+        }
+        self.config = AgentConfig(**agent_kwargs)

agentbox/agent/loop.py

@@ -0,0 +1,201 @@
+"""Multi-turn agent loop with OpenAI tool calling."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from agentbox.config import AgentConfig
+from agentbox.model.base import ModelClient, ModelResponse
+from agentbox.tools.base import ToolContext
+from agentbox.tools.executor import ToolExecutor
+from agentbox.trajectory.schema import Message, ToolCallRecord
+from agentbox.types import FinalStatus, MessageRole
+
+
+class AgentLoopResult(BaseModel):
+    messages: list[Message]
+    tool_call_records: list[ToolCallRecord]
+    final_status: FinalStatus
+    steps: int
+    model_calls: int
+    stop_reason: str
+    error: str | None = None
+    usage_totals: dict[str, int] = Field(default_factory=dict)
+
+
+class AgentLoop:
+    def __init__(
+        self,
+        model: ModelClient,
+        tools: ToolExecutor,
+        config: AgentConfig,
+        *,
+        on_step: Callable[[int, Message], Awaitable[None] | None] | None = None,
+    ) -> None:
+        self.model = model
+        self.tools = tools
+        self.config = config
+        self.on_step = on_step
+
+    async def run(
+        self,
+        *,
+        system_prompt: str,
+        user_prompt: str,
+        ctx: ToolContext,
+        episode_deadline: float | None = None,
+    ) -> AgentLoopResult:
+        messages: list[Message] = [
+            Message(role=MessageRole.SYSTEM, content=system_prompt),
+            Message(role=MessageRole.USER, content=user_prompt),
+        ]
+        tool_records: list[ToolCallRecord] = []
+        usage_totals = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        model_calls = 0
+        step = 0
+        schemas = self.tools.schemas()
+
+        try:
+            while step < self.config.max_steps:
+                if episode_deadline is not None and time.monotonic() >= episode_deadline:
+                    return AgentLoopResult(
+                        messages=messages,
+                        tool_call_records=tool_records,
+                        final_status=FinalStatus.TIMEOUT,
+                        steps=step,
+                        model_calls=model_calls,
+                        stop_reason="timeout",
+                        usage_totals=usage_totals,
+                    )
+
+                openai_messages = [m.to_openai_dict() for m in messages]
+                try:
+                    response = await asyncio.wait_for(
+                        self.model.complete(
+                            openai_messages,
+                            tools=schemas or None,
+                            tool_choice="auto",
+                        ),
+                        timeout=self.config.step_timeout_s,
+                    )
+                except TimeoutError:
+                    return AgentLoopResult(
+                        messages=messages,
+                        tool_call_records=tool_records,
+                        final_status=FinalStatus.TIMEOUT,
+                        steps=step,
+                        model_calls=model_calls,
+                        stop_reason="timeout",
+                        error=f"step timed out after {self.config.step_timeout_s}s",
+                        usage_totals=usage_totals,
+                    )
+                model_calls += 1
+                self._accumulate_usage(usage_totals, response)
+
+                assistant = self._to_assistant_message(response)
+                messages.append(assistant)
+
+                if self.on_step:
+                    maybe = self.on_step(step, assistant)
+                    if maybe is not None:
+                        await maybe
+
+                if not response.tool_calls:
+                    return AgentLoopResult(
+                        messages=messages,
+                        tool_call_records=tool_records,
+                        final_status=FinalStatus.SUCCESS,  # provisional; verifier decides reward
+                        steps=step,
+                        model_calls=model_calls,
+                        stop_reason="final_answer",
+                        usage_totals=usage_totals,
+                    )
+
+                ctx.step = step
+                try:
+                    tool_messages = await asyncio.wait_for(
+                        self.tools.execute_many(
+                            response.tool_calls,
+                            ctx,
+                            parallel=self.config.parallel_tool_calls,
+                        ),
+                        timeout=self.config.step_timeout_s,
+                    )
+                except TimeoutError:
+                    return AgentLoopResult(
+                        messages=messages,
+                        tool_call_records=tool_records,
+                        final_status=FinalStatus.TIMEOUT,
+                        steps=step,
+                        model_calls=model_calls,
+                        stop_reason="timeout",
+                        error=f"tool batch timed out after {self.config.step_timeout_s}s",
+                        usage_totals=usage_totals,
+                    )
+                for tc, tm in zip(response.tool_calls, tool_messages, strict=True):
+                    try:
+                        args = json.loads(tc.function.arguments or "{}")
+                    except json.JSONDecodeError:
+                        args = {"_raw": tc.function.arguments}
+                    is_error = bool(
+                        tm.content and str(tm.content).startswith("ERROR:")
+                    )
+                    tool_records.append(
+                        ToolCallRecord(
+                            step=step,
+                            tool_call_id=tc.id,
+                            name=tc.function.name,
+                            arguments=args if isinstance(args, dict) else {},
+                            result=tm.content or "",
+                            is_error=is_error,
+                        )
+                    )
+                messages.extend(tool_messages)
+                step += 1
+
+            return AgentLoopResult(
+                messages=messages,
+                tool_call_records=tool_records,
+                final_status=FinalStatus.MAX_STEPS,
+                steps=step,
+                model_calls=model_calls,
+                stop_reason="max_steps",
+                usage_totals=usage_totals,
+            )
+        except Exception as exc:
+            return AgentLoopResult(
+                messages=messages,
+                tool_call_records=tool_records,
+                final_status=FinalStatus.ERROR,
+                steps=step,
+                model_calls=model_calls,
+                stop_reason="error",
+                error=str(exc),
+                usage_totals=usage_totals,
+            )
+
+    def _to_assistant_message(self, response: ModelResponse) -> Message:
+        return Message(
+            role=MessageRole.ASSISTANT,
+            content=response.content,
+            tool_calls=response.tool_calls or None,
+            reasoning_content=response.reasoning_content
+            if self.config.include_thinking
+            else None,
+            raw=response.raw_choice,
+        )
+
+    @staticmethod
+    def _accumulate_usage(
+        totals: dict[str, int], response: ModelResponse
+    ) -> None:
+        if not response.usage:
+            return
+        for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
+            totals[key] = totals.get(key, 0) + int(response.usage.get(key, 0))

agentbox/agent/prompts.py

@@ -0,0 +1,44 @@
+"""Default system prompts."""
+
+from __future__ import annotations
+
+from agentbox.types import ToolMode
+
+
+STRUCTURED_SYSTEM_PROMPT = """\
+You are a coding agent working inside an isolated Linux container.
+Your workspace is {workspace_dir}.
+
+You solve the user's task using the provided tools. Prefer small, correct edits.
+When you are done, respond with a brief final message and NO tool calls.
+Do not invent file paths. Always verify your work with tests or commands when possible.
+"""
+
+SHELL_SYSTEM_PROMPT = """\
+You are a coding agent working inside an isolated Linux container.
+Your workspace is {workspace_dir}.
+
+You only have the run_command tool. Use standard Unix tools to inspect and edit files.
+When you are done, respond with a brief final message and NO tool calls.
+"""
+
+
+def render_system_prompt(
+    *,
+    workspace_dir: str = "/workspace",
+    mode: ToolMode | str = ToolMode.STRUCTURED,
+    extra: str | None = None,
+    override: str | None = None,
+) -> str:
+    if override:
+        base = override.format(workspace_dir=workspace_dir)
+    else:
+        template = (
+            SHELL_SYSTEM_PROMPT
+            if mode in (ToolMode.SHELL, "shell")
+            else STRUCTURED_SYSTEM_PROMPT
+        )
+        base = template.format(workspace_dir=workspace_dir)
+    if extra:
+        base = base.rstrip() + "\n\n" + extra.strip()
+    return base

agentbox/benchmark/__init__.py

@@ -0,0 +1,29 @@
+"""Native real-rollout benchmark suites."""
+
+from agentbox.benchmark.compare import compare_reports
+from agentbox.benchmark.loader import create_suite_from_tasks, load_suite, save_suite
+from agentbox.benchmark.runner import BenchmarkRunner
+from agentbox.benchmark.schema import (
+    BenchmarkReport,
+    BenchmarkRunConfig,
+    BenchmarkSuite,
+    BenchmarkSuiteManifest,
+    ModelUnderTest,
+    ScoringConfig,
+    SetupCheckSpec,
+)
+
+__all__ = [
+    "BenchmarkReport",
+    "BenchmarkRunConfig",
+    "BenchmarkRunner",
+    "BenchmarkSuite",
+    "BenchmarkSuiteManifest",
+    "ModelUnderTest",
+    "ScoringConfig",
+    "SetupCheckSpec",
+    "compare_reports",
+    "create_suite_from_tasks",
+    "load_suite",
+    "save_suite",
+]

agentbox/benchmark/compare.py

@@ -0,0 +1,51 @@
+"""Compare two benchmark reports."""
+
+from __future__ import annotations
+
+from agentbox.benchmark.schema import BenchmarkReport, ComparisonResult
+
+
+def compare_reports(
+    a: BenchmarkReport,
+    b: BenchmarkReport,
+    *,
+    require_hash_match: bool = False,
+) -> ComparisonResult:
+    warnings: list[str] = []
+    if a.suite_id != b.suite_id:
+        warnings.append(f"suite_id differs: {a.suite_id!r} vs {b.suite_id!r}")
+    hash_match = (a.suite_content_hash or "") == (b.suite_content_hash or "")
+    if not hash_match:
+        warnings.append(
+            f"content_hash differs: {a.suite_content_hash!r} vs {b.suite_content_hash!r}"
+        )
+    if require_hash_match and not hash_match:
+        warnings.append("require_hash_match=True: comparison may be unfair")
+
+    a_models = {m.model_id: m for m in a.models}
+    b_models = {m.model_id: m for m in b.models}
+    deltas: list[dict] = []
+
+    for mid in sorted(set(a_models) | set(b_models)):
+        ma, mb = a_models.get(mid), b_models.get(mid)
+        if ma is None or mb is None:
+            warnings.append(f"model {mid!r} missing in one report")
+            continue
+        deltas.append(
+            {
+                "model_id": mid,
+                "success_rate_a": ma.aggregate.success_rate,
+                "success_rate_b": mb.aggregate.success_rate,
+                "delta_success_rate": mb.aggregate.success_rate - ma.aggregate.success_rate,
+                "mean_steps_a": ma.aggregate.mean_steps,
+                "mean_steps_b": mb.aggregate.mean_steps,
+                "delta_mean_steps": mb.aggregate.mean_steps - ma.aggregate.mean_steps,
+            }
+        )
+
+    return ComparisonResult(
+        suite_id=a.suite_id,
+        hash_match=hash_match,
+        warnings=warnings,
+        deltas=deltas,
+    )

agentbox/benchmark/hash.py

@@ -0,0 +1,59 @@
+"""Stable content hashing for benchmark suite integrity."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+from agentbox.benchmark.schema import BenchmarkSuiteManifest
+from agentbox.tasks.schema import Task
+
+
+def _canonical_json(obj: Any) -> str:
+    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
+
+
+def _task_payload(task: Task) -> dict[str, Any]:
+    return {
+        "task_id": task.task_id,
+        "description": task.description,
+        "starter_files": task.starter_files,
+        "setup_commands": task.setup_commands,
+        "verifier": task.verifier.model_dump(mode="json"),
+        "metadata": task.metadata,
+        "system_prompt_extra": task.system_prompt_extra,
+        "allowed_tools": task.allowed_tools,
+        "max_steps": task.max_steps,
+    }
+
+
+def _fairness_config_payload(manifest: BenchmarkSuiteManifest) -> dict[str, Any]:
+    """Config fields that affect evaluation fairness (exclude free-form notes)."""
+    agent = manifest.agent.model_dump(
+        mode="json",
+        exclude={"custom_tools"},  # non-serializable / not used in suite freeze
+    )
+    return {
+        "suite_id": manifest.suite_id,
+        "version": manifest.version,
+        "sandbox": manifest.sandbox.model_dump(mode="json"),
+        "agent": agent,
+        "n_per_task": manifest.n_per_task,
+        "seed": manifest.seed,
+        "scoring": manifest.scoring.model_dump(mode="json"),
+    }
+
+
+def compute_suite_content_hash(
+    tasks: list[Task],
+    manifest: BenchmarkSuiteManifest,
+) -> str:
+    """SHA-256 hex over tasks + fairness-critical suite config."""
+    ordered = sorted(tasks, key=lambda t: t.task_id)
+    payload = {
+        "config": _fairness_config_payload(manifest),
+        "tasks": [_task_payload(t) for t in ordered],
+    }
+    raw = _canonical_json(payload).encode("utf-8")
+    return hashlib.sha256(raw).hexdigest()

agentbox/benchmark/loader.py

@@ -0,0 +1,152 @@
+"""Load and save benchmark suite directory packs."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+
+from agentbox.benchmark.schema import BenchmarkSuite, BenchmarkSuiteManifest
+from agentbox.config import AgentConfig, SandboxConfig
+from agentbox.errors import BenchmarkError, TaskValidationError
+from agentbox.tasks.bulk import load_task_dataset
+from agentbox.tasks.schema import Task
+from agentbox.types import ToolMode
+
+
+def load_suite(path: str | Path) -> BenchmarkSuite:
+    """Load suite.json + tasks from a directory pack."""
+    root = Path(path).resolve()
+    manifest_path = root / "suite.json"
+    if not manifest_path.is_file():
+        raise BenchmarkError(f"suite.json not found in {root}")
+    try:
+        data = json.loads(manifest_path.read_text(encoding="utf-8"))
+        manifest = BenchmarkSuiteManifest.model_validate(data)
+    except Exception as exc:
+        raise BenchmarkError(f"Invalid suite.json: {exc}") from exc
+
+    tasks_dir = root / manifest.tasks_path
+    if not tasks_dir.exists():
+        raise BenchmarkError(f"tasks path not found: {tasks_dir}")
+
+    tasks = load_task_dataset(tasks_dir)
+    if not tasks:
+        # try single-level task.json dirs only already covered; try direct json
+        for p in sorted(tasks_dir.glob("*.json")):
+            try:
+                tasks.append(Task.from_json(p))
+            except TaskValidationError:
+                continue
+    if not tasks:
+        raise BenchmarkError(f"No tasks found under {tasks_dir}")
+
+    # Fairness: benchmarks should not carry custom_tools blobs
+    agent = manifest.agent.model_copy(update={"custom_tools": [], "drop_tools_prob": 0.0})
+    manifest = manifest.model_copy(update={"agent": agent})
+    return BenchmarkSuite(root=root, manifest=manifest, tasks=tasks)
+
+
+def save_suite(suite: BenchmarkSuite, path: str | Path) -> Path:
+    """Write suite pack to directory (suite.json + tasks/)."""
+    root = Path(path).resolve()
+    root.mkdir(parents=True, exist_ok=True)
+    tasks_dir = root / suite.manifest.tasks_path
+    tasks_dir.mkdir(parents=True, exist_ok=True)
+
+    for task in suite.tasks:
+        task_dir = tasks_dir / task.task_id
+        task_dir.mkdir(parents=True, exist_ok=True)
+        # Write starter files into files/ and strip from task.json for cleanliness
+        files_dir = task_dir / "files"
+        if task.starter_files:
+            files_dir.mkdir(parents=True, exist_ok=True)
+            for rel, content in task.starter_files.items():
+                fp = files_dir / rel
+                fp.parent.mkdir(parents=True, exist_ok=True)
+                fp.write_text(content, encoding="utf-8")
+        slim = task.model_copy(update={"starter_files": {}})
+        slim.save_json(task_dir / "task.json")
+
+    # Serialize agent without custom_tools
+    m = suite.manifest.model_copy(deep=True)
+    m.agent = m.agent.model_copy(update={"custom_tools": []})
+    (root / "suite.json").write_text(
+        m.model_dump_json(indent=2),
+        encoding="utf-8",
+    )
+    return root
+
+
+def create_suite_from_tasks(
+    tasks_src: str | Path,
+    dest: str | Path,
+    *,
+    suite_id: str,
+    name: str,
+    version: str = "1.0.0",
+    description: str = "",
+    sandbox: SandboxConfig | None = None,
+    agent: AgentConfig | None = None,
+    n_per_task: int = 1,
+    concurrency: int = 8,
+    freeze: bool = True,
+) -> BenchmarkSuite:
+    """Create a suite directory from a tasks folder or dataset."""
+    src = Path(tasks_src)
+    tasks = load_task_dataset(src)
+    if not tasks and src.is_dir():
+        for p in sorted(src.rglob("task.json")):
+            try:
+                tasks.append(Task.from_dir(p.parent))
+            except TaskValidationError:
+                continue
+        for p in sorted(src.glob("*.json")):
+            try:
+                tasks.append(Task.from_json(p))
+            except TaskValidationError:
+                continue
+    # de-dupe
+    seen: set[str] = set()
+    unique: list[Task] = []
+    for t in tasks:
+        if t.task_id in seen:
+            continue
+        seen.add(t.task_id)
+        unique.append(t)
+    if not unique:
+        raise BenchmarkError(f"No tasks found in {src}")
+
+    manifest = BenchmarkSuiteManifest(
+        suite_id=suite_id,
+        name=name,
+        description=description,
+        version=version,
+        sandbox=sandbox or SandboxConfig(),
+        agent=agent
+        or AgentConfig(
+            tools=ToolMode.STRUCTURED,
+            max_steps=40,
+            drop_tools_prob=0.0,
+            custom_tools=[],
+        ),
+        n_per_task=n_per_task,
+        concurrency=concurrency,
+    )
+    suite = BenchmarkSuite(manifest=manifest, tasks=unique)
+    if freeze:
+        suite = suite.freeze()
+    save_suite(suite, dest)
+    suite.root = Path(dest).resolve()
+    return suite
+
+
+def snapshot_suite(suite: BenchmarkSuite, dest: str | Path) -> Path:
+    """Copy frozen suite into results for audit."""
+    dest_p = Path(dest)
+    if suite.root and suite.root.is_dir():
+        if dest_p.exists():
+            shutil.rmtree(dest_p)
+        shutil.copytree(suite.root, dest_p)
+        return dest_p
+    return save_suite(suite, dest_p)