tokenjam-bench 0.1.0__py3-none-any.whl

tjbench/__init__.py

@@ -0,0 +1,6 @@
+"""tokenjam-benchmark — the package.
+
+All modules live here so the repo root stays clean and the installed `tjbench`
+console script (`tjbench.cli:cli`) does not collide with same-named PyPI
+packages (e.g. `cli`). Run from source via the root `run.py`.
+"""

tjbench/agent_pipeline.py

@@ -0,0 +1,117 @@
+"""Agent proof pipeline.
+
+Runs an agent benchmark on the ORIGINAL model and the CANDIDATE model (TokenJam's
+recommendation), each via the multi-turn AgentRunner, scores each run on its
+trace (answer correctness + tool-call validation incl. the safety gate), prices
+the summed multi-turn token usage, and feeds the per-task outcomes into the SAME
+assembler the single-shot path uses — so Wilson CIs, McNemar, and cost
+validation apply unchanged.
+
+This is the payoff of the keystone: agent benchmarks inherit all the statistical
+rigor for free, and a candidate that takes an unsafe action fails the task even
+when its answer text is correct.
+"""
+from __future__ import annotations
+
+from tjbench.agents.runner import AgentRunner
+from tjbench.benchmarks import get_agent_benchmark
+from tjbench.cost import price_completion
+from tjbench.models.anthropic_agent_client import get_tool_calling_client
+from tjbench.models.registry import parse_spec
+from tjbench.pipeline import assemble_proof, resolve_candidate
+from tjbench.report import TaskOutcome
+
+
+def _run_agent_samples(client, provider, model, benchmark, task, registry,
+                       samples, max_turns, temperature, max_tokens):
+    passes = 0
+    out_tok = 0
+    cost = 0.0
+    last_detail = ""
+    for _ in range(samples):
+        runner = AgentRunner(client, registry, max_turns=max_turns,
+                             temperature=temperature, max_tokens=max_tokens)
+        trace = runner.run(task.task_id, task.prompt)
+        score = benchmark.score(task, trace)
+        seq = trace.tool_sequence()
+        last_detail = (
+            f"tools={seq} stopped={trace.stopped_reason} "
+            f"turns={trace.num_turns} -> {score.detail}"
+        )
+        if score.passed:
+            passes += 1
+        out_tok += trace.total_output_tokens
+        cost += price_completion(provider, model, trace.as_completion())
+    return passes, out_tok, round(cost, 8), last_detail
+
+
+def run_agent_proof(
+    *,
+    benchmark_name: str,
+    original_spec: str,
+    candidate_spec: str | None = None,
+    limit: int | None = None,
+    samples: int = 1,
+    temperature: float = 0.0,
+    max_turns: int = 8,
+    max_tokens: int = 1024,
+    mock: bool = False,
+    candidate_behavior: str = "ok",
+    alpha: float = 0.05,
+):
+    """Run an agent-benchmark proof (original vs TokenJam's candidate)."""
+    if samples < 1:
+        raise ValueError("samples must be >= 1")
+
+    recommended_by = "tokenjam.DOWNGRADE_CANDIDATES"
+    if candidate_spec is None:
+        candidate_spec = resolve_candidate(original_spec)
+        if candidate_spec is None:
+            raise ValueError(
+                f"TokenJam has no downgrade candidate for '{original_spec}'. "
+                f"Pass --candidate explicitly to override."
+            )
+    else:
+        recommended_by = "explicit --candidate override"
+
+    orig_provider, orig_model = parse_spec(original_spec)
+    cand_provider, cand_model = parse_spec(candidate_spec)
+
+    # Offline: original behaves correctly, candidate's behavior is configurable
+    # (ok | wrong | unsafe) to exercise the answer + safety gates.
+    original = get_tool_calling_client(original_spec, mock=mock, behavior="ok")
+    candidate = get_tool_calling_client(candidate_spec, mock=mock, behavior=candidate_behavior)
+
+    benchmark = get_agent_benchmark(benchmark_name)
+    registry = benchmark.tools()
+    tasks = benchmark.tasks(limit=limit)
+
+    outcomes: list[TaskOutcome] = []
+    tot_o = tot_c = 0
+    for task in tasks:
+        o_pass, o_out, o_cost, o_detail = _run_agent_samples(
+            original, orig_provider, orig_model, benchmark, task, registry,
+            samples, max_turns, temperature, max_tokens)
+        c_pass, c_out, c_cost, c_detail = _run_agent_samples(
+            candidate, cand_provider, cand_model, benchmark, task, registry,
+            samples, max_turns, temperature, max_tokens)
+        tot_o += o_pass
+        tot_c += c_pass
+        outcomes.append(TaskOutcome(
+            task_id=task.task_id, samples=samples,
+            original_passes=o_pass, candidate_passes=c_pass,
+            original_cost_usd=o_cost, candidate_cost_usd=c_cost,
+            original_output_tokens=o_out, candidate_output_tokens=c_out,
+            original_detail=o_detail, candidate_detail=c_detail,
+        ))
+
+    return assemble_proof(
+        outcomes,
+        benchmark_name=benchmark_name,
+        original_spec=original_spec, candidate_spec=candidate_spec,
+        recommended_by=recommended_by, samples=samples, mock=mock,
+        orig_provider=orig_provider, orig_model=orig_model,
+        cand_provider=cand_provider, cand_model=cand_model,
+        sample_pass_totals=(tot_o, tot_c),
+        alpha=alpha,
+    )

tjbench/agents/__init__.py

@@ -0,0 +1,25 @@
+"""Agent execution: the multi-turn AgentRunner, tools, trace, and validation.
+
+This package is the keystone that turns the bench from an LLM benchmark into an
+agent benchmark. It feeds the SAME proof machinery (stats + cost) as the
+single-shot path — a trace yields a per-task pass/fail and a measured
+multi-turn cost.
+"""
+from __future__ import annotations
+
+from tjbench.agents.runner import AgentRunner
+from tjbench.agents.tools import Tool, ToolRegistry, ToolResult
+from tjbench.agents.trace import AgentTrace, ToolCallRecord, TurnRecord
+from tjbench.agents.validation import ToolValidation, validate_tools
+
+__all__ = [
+    "AgentRunner",
+    "Tool",
+    "ToolRegistry",
+    "ToolResult",
+    "AgentTrace",
+    "TurnRecord",
+    "ToolCallRecord",
+    "ToolValidation",
+    "validate_tools",
+]

tjbench/agents/runner.py

@@ -0,0 +1,66 @@
+"""AgentRunner — the multi-turn loop. THE keystone.
+
+It drives a tool-calling model against a tool registry until the model returns a
+final answer or hits `max_turns`, recording every turn into an AgentTrace. This
+one component is what unlocks agent benchmarks, multi-turn evaluation, tool-call
+validation, and side-effect safety — all of which were blocked on the
+single-shot `complete(prompt)` interface.
+"""
+from __future__ import annotations
+
+from tjbench.agents.tools import ToolRegistry
+from tjbench.agents.trace import AgentTrace, ToolCallRecord, TurnRecord
+from tjbench.models.tool_calling import ToolCallingClient
+
+
+class AgentRunner:
+    def __init__(self, client: ToolCallingClient, tools: ToolRegistry,
+                 max_turns: int = 8, temperature: float = 0.0,
+                 max_tokens: int = 1024) -> None:
+        self.client = client
+        self.tools = tools
+        self.max_turns = max_turns
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+
+    def run(self, task_id: str, prompt: str) -> AgentTrace:
+        messages: list[dict] = [{"role": "user", "content": prompt}]
+        trace = AgentTrace(task_id=task_id)
+        specs = self.tools.specs()
+
+        for i in range(self.max_turns):
+            turn = self.client.chat(messages, specs, self.temperature, self.max_tokens)
+
+            if not turn.wants_tools:
+                trace.turns.append(TurnRecord(
+                    index=i, assistant_text=turn.text, tool_calls=[],
+                    input_tokens=turn.input_tokens, output_tokens=turn.output_tokens,
+                    cache_tokens=turn.cache_tokens,
+                ))
+                trace.final_text = turn.text
+                trace.stopped_reason = "final"
+                return trace
+
+            # The model asked for tool(s): execute each, record, feed results back.
+            messages.append({
+                "role": "assistant", "content": turn.text, "tool_calls": turn.tool_calls,
+            })
+            call_records: list[ToolCallRecord] = []
+            for tc in turn.tool_calls:
+                result = self.tools.execute(tc.name, tc.arguments)
+                call_records.append(ToolCallRecord(
+                    name=tc.name, arguments=tc.arguments,
+                    result=result.output, is_error=result.is_error,
+                ))
+                messages.append({
+                    "role": "tool", "tool_call_id": tc.id, "name": tc.name,
+                    "content": result.output,
+                })
+            trace.turns.append(TurnRecord(
+                index=i, assistant_text=turn.text, tool_calls=call_records,
+                input_tokens=turn.input_tokens, output_tokens=turn.output_tokens,
+                cache_tokens=turn.cache_tokens,
+            ))
+
+        trace.stopped_reason = "max_turns"
+        return trace

tjbench/agents/swe_bench_tools.py

@@ -0,0 +1,296 @@
+"""SWE-Bench tool implementations that operate on a real workspace.
+
+These tools are designed to be bound to a specific task's file workspace
+at runtime. They provide the core developer operations needed for
+SWE-Bench: reading files, editing files, and running commands.
+"""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from tjbench.agents.tools import ToolResult
+
+
+class SWEBenchToolSet:
+    """Collection of tools for SWE-Bench agent evaluation.
+    
+    Each tool operates on a specific workspace directory. The toolset
+    is instantiated per-task and bound to the task's workspace.
+    """
+
+    def __init__(self, workspace: Path) -> None:
+        self.workspace = workspace
+        self._files: dict[str, str] = {}  # Cache of file contents
+
+    def _resolve_path(self, path: str) -> Path:
+        """Resolve a path relative to the workspace."""
+        # Prevent directory traversal outside workspace
+        resolved = (self.workspace / path).resolve()
+        if not str(resolved).startswith(str(self.workspace.resolve())):
+            raise ValueError(f"Path {path} escapes workspace")
+        return resolved
+
+    def _read_file(self, path: Path) -> str:
+        """Read a file, caching the result."""
+        str_path = str(path)
+        if str_path not in self._files:
+            if path.exists():
+                self._files[str_path] = path.read_text(encoding="utf-8")
+            else:
+                self._files[str_path] = ""
+        return self._files[str_path]
+
+    def _write_file(self, path: Path, content: str) -> None:
+        """Write a file and update cache."""
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(content, encoding="utf-8")
+        self._files[str(path)] = content
+
+    # --- Tool implementations ---
+
+    def view(self, args: dict[str, Any]) -> ToolResult:
+        """View the contents of a file."""
+        try:
+            path = self._resolve_path(args["path"])
+            if not path.exists():
+                return ToolResult(
+                    output=f"Error: File '{args['path']}' does not exist.",
+                    is_error=True,
+                )
+            content = self._read_file(path)
+            # Add line numbers for readability
+            lines = content.split("\n")
+            numbered = "\n".join(f"{i+1:4d} | {line}" for i, line in enumerate(lines))
+            return ToolResult(output=f"File: {args['path']}\n{numbered}")
+        except Exception as e:
+            return ToolResult(output=f"Error: {e}", is_error=True)
+
+    def view_range(self, args: dict[str, Any]) -> ToolResult:
+        """View a specific range of lines in a file."""
+        try:
+            path = self._resolve_path(args["path"])
+            start = args["start"]
+            end = args["end"]
+            
+            if not path.exists():
+                return ToolResult(
+                    output=f"Error: File '{args['path']}' does not exist.",
+                    is_error=True,
+                )
+            
+            content = self._read_file(path)
+            lines = content.split("\n")
+            
+            # Clamp to valid range
+            start = max(1, start)
+            end = min(len(lines), end)
+            
+            selected = lines[start - 1:end]
+            numbered = "\n".join(f"{i+start:4d} | {line}" for i, line in enumerate(selected))
+            return ToolResult(
+                output=f"File: {args['path']} (lines {start}-{end})\n{numbered}"
+            )
+        except Exception as e:
+            return ToolResult(output=f"Error: {e}", is_error=True)
+
+    def str_replace(self, args: dict[str, Any]) -> ToolResult:
+        """Replace an exact string in a file."""
+        try:
+            path = self._resolve_path(args["path"])
+            old_str = args["old_str"]
+            new_str = args["new_str"]
+            
+            if not path.exists():
+                return ToolResult(
+                    output=f"Error: File '{args['path']}' does not exist.",
+                    is_error=True,
+                )
+            
+            content = self._read_file(path)
+            
+            if old_str not in content:
+                return ToolResult(
+                    output=f"Error: Could not find the exact string in {args['path']}. "
+                           "Make sure the old_str matches exactly (including whitespace).",
+                    is_error=True,
+                )
+            
+            # Count occurrences
+            count = content.count(old_str)
+            if count > 1:
+                return ToolResult(
+                    output=f"Error: Found {count} occurrences of the string. "
+                           "Please use a more specific old_str that matches exactly once.",
+                    is_error=True,
+                )
+            
+            new_content = content.replace(old_str, new_str, 1)
+            self._write_file(path, new_content)
+            
+            return ToolResult(
+                output=f"Successfully replaced in {args['path']}. "
+                       f"Changed {len(old_str)} chars to {len(new_str)} chars."
+            )
+        except Exception as e:
+            return ToolResult(output=f"Error: {e}", is_error=True)
+
+    def create(self, args: dict[str, Any]) -> ToolResult:
+        """Create a new file with the given content."""
+        try:
+            path = self._resolve_path(args["path"])
+            content = args["content"]
+            
+            if path.exists():
+                return ToolResult(
+                    output=f"Error: File '{args['path']}' already exists. Use str_replace to modify it.",
+                    is_error=True,
+                )
+            
+            self._write_file(path, content)
+            return ToolResult(output=f"Created file: {args['path']}")
+        except Exception as e:
+            return ToolResult(output=f"Error: {e}", is_error=True)
+
+    def insert(self, args: dict[str, Any]) -> ToolResult:
+        """Insert text after a specific line."""
+        try:
+            path = self._resolve_path(args["path"])
+            line = args["line"]
+            new_str = args["new_str"]
+            
+            if not path.exists():
+                return ToolResult(
+                    output=f"Error: File '{args['path']}' does not exist.",
+                    is_error=True,
+                )
+            
+            content = self._read_file(path)
+            lines = content.split("\n")
+            
+            if line < 0 or line > len(lines):
+                return ToolResult(
+                    output=f"Error: Line {line} is out of range (file has {len(lines)} lines).",
+                    is_error=True,
+                )
+            
+            lines.insert(line, new_str)
+            self._write_file(path, "\n".join(lines))
+            
+            return ToolResult(output=f"Inserted after line {line} in {args['path']}")
+        except Exception as e:
+            return ToolResult(output=f"Error: {e}", is_error=True)
+
+    def bash(self, args: dict[str, Any]) -> ToolResult:
+        """Run a shell command in the workspace."""
+        try:
+            command = args["command"]
+            timeout = args.get("timeout", 30)
+            
+            result = subprocess.run(
+                command,
+                shell=True,
+                cwd=self.workspace,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+            
+            output = f"Exit code: {result.returncode}\n"
+            if result.stdout:
+                output += f"STDOUT:\n{result.stdout}\n"
+            if result.stderr:
+                output += f"STDERR:\n{result.stderr}\n"
+            
+            return ToolResult(
+                output=output,
+                is_error=result.returncode != 0,
+            )
+        except subprocess.TimeoutExpired:
+            return ToolResult(
+                output=f"Error: Command timed out after {timeout}s.",
+                is_error=True,
+            )
+        except Exception as e:
+            return ToolResult(output=f"Error: {e}", is_error=True)
+
+    def get_tool_specs(self) -> list[dict[str, Any]]:
+        """Return tool specifications for the agent."""
+        return [
+            {
+                "name": "view",
+                "description": "View the contents of a file. Shows line numbers.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string", "description": "Path to the file (relative to workspace)"},
+                    },
+                    "required": ["path"],
+                },
+            },
+            {
+                "name": "view_range",
+                "description": "View a specific range of lines in a file.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string"},
+                        "start": {"type": "integer", "description": "Start line (1-indexed)"},
+                        "end": {"type": "integer", "description": "End line (1-indexed)"},
+                    },
+                    "required": ["path", "start", "end"],
+                },
+            },
+            {
+                "name": "str_replace",
+                "description": "Replace an exact string in a file. The old_str must match exactly once.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string"},
+                        "old_str": {"type": "string", "description": "Exact text to replace (must match exactly once)"},
+                        "new_str": {"type": "string", "description": "Replacement text"},
+                    },
+                    "required": ["path", "old_str", "new_str"],
+                },
+            },
+            {
+                "name": "create",
+                "description": "Create a new file with the given content.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string"},
+                        "content": {"type": "string"},
+                    },
+                    "required": ["path", "content"],
+                },
+            },
+            {
+                "name": "insert",
+                "description": "Insert text after a specific line in a file.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string"},
+                        "line": {"type": "integer", "description": "Line after which to insert (1-indexed)"},
+                        "new_str": {"type": "string"},
+                    },
+                    "required": ["path", "line", "new_str"],
+                },
+            },
+            {
+                "name": "bash",
+                "description": "Run a shell command in the workspace. Use for running tests, git, etc.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "command": {"type": "string", "description": "Shell command to run"},
+                        "timeout": {"type": "integer", "default": 30, "description": "Timeout in seconds"},
+                    },
+                    "required": ["command"],
+                },
+            },
+        ]

tjbench/agents/tools.py

@@ -0,0 +1,63 @@
+"""Tools an agent can call, and the registry that executes them.
+
+A `Tool` carries a JSON-schema for its arguments (so it can be advertised to a
+tool-calling model) and a `dangerous` flag. The flag is load-bearing for the
+safety story from the review: a cheaper model can produce a correct-looking
+final answer while taking a catastrophic action (delete vs read). Marking the
+tool dangerous lets validation fail the task on the *action*, not the text.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+
+@dataclass
+class ToolResult:
+    output: str
+    is_error: bool = False
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str
+    parameters: dict[str, Any]          # JSON Schema for the arguments object
+    run: Callable[[dict[str, Any]], ToolResult]
+    dangerous: bool = False             # has side effects worth flagging on misuse
+
+    def spec(self) -> dict[str, Any]:
+        """Provider-agnostic advertisement of this tool to a model."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": self.parameters,
+        }
+
+
+class ToolRegistry:
+    def __init__(self, tools: list[Tool] | None = None) -> None:
+        self._tools: dict[str, Tool] = {}
+        for t in tools or []:
+            self.register(t)
+
+    def register(self, tool: Tool) -> None:
+        self._tools[tool.name] = tool
+
+    def names(self) -> list[str]:
+        return list(self._tools.keys())
+
+    def dangerous_names(self) -> set[str]:
+        return {n for n, t in self._tools.items() if t.dangerous}
+
+    def specs(self) -> list[dict[str, Any]]:
+        return [t.spec() for t in self._tools.values()]
+
+    def execute(self, name: str, arguments: dict[str, Any]) -> ToolResult:
+        tool = self._tools.get(name)
+        if tool is None:
+            return ToolResult(output=f"unknown tool: {name}", is_error=True)
+        try:
+            return tool.run(arguments or {})
+        except Exception as exc:  # a tool raising is a tool error, not a crash
+            return ToolResult(output=f"tool error: {exc}", is_error=True)

tjbench/agents/trace.py

@@ -0,0 +1,72 @@
+"""The observable record of an agent run.
+
+A single-shot completion has nothing to validate beyond its text. A multi-turn
+agent run produces a *trace*: which tools were called, with what arguments, in
+what order, whether they errored, the per-turn token usage, and the final
+answer. The trace is what makes tool-call validation, multi-turn evaluation, and
+side-effect safety possible — and summing tokens across turns is what makes the
+cost number honest for agents (a model that loops 8× costs 8×).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from tjbench.models.base import Completion
+
+
+@dataclass
+class ToolCallRecord:
+    name: str
+    arguments: dict
+    result: str
+    is_error: bool
+
+
+@dataclass
+class TurnRecord:
+    index: int
+    assistant_text: str
+    tool_calls: list[ToolCallRecord]
+    input_tokens: int
+    output_tokens: int
+    cache_tokens: int = 0
+
+
+@dataclass
+class AgentTrace:
+    task_id: str
+    turns: list[TurnRecord] = field(default_factory=list)
+    final_text: str = ""
+    stopped_reason: str = "final"        # "final" | "max_turns"
+
+    def tool_sequence(self) -> list[str]:
+        """Ordered tool names across the whole run."""
+        return [tc.name for turn in self.turns for tc in turn.tool_calls]
+
+    def all_tool_calls(self) -> list[ToolCallRecord]:
+        return [tc for turn in self.turns for tc in turn.tool_calls]
+
+    @property
+    def num_turns(self) -> int:
+        return len(self.turns)
+
+    @property
+    def total_input_tokens(self) -> int:
+        return sum(t.input_tokens for t in self.turns)
+
+    @property
+    def total_output_tokens(self) -> int:
+        return sum(t.output_tokens for t in self.turns)
+
+    @property
+    def total_cache_tokens(self) -> int:
+        return sum(t.cache_tokens for t in self.turns)
+
+    def as_completion(self) -> Completion:
+        """Collapse the run's token usage so the existing pricing path applies."""
+        return Completion(
+            text=self.final_text,
+            input_tokens=self.total_input_tokens,
+            output_tokens=self.total_output_tokens,
+            cache_tokens=self.total_cache_tokens,
+        )