arbiter-cli 0.1.0__py3-none-any.whl

arbiter/core/runner.py

@@ -0,0 +1,257 @@
+"""Core runner -- executes prompts against multiple models in parallel."""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime, timezone
+from typing import AsyncIterator, Callable, Optional
+
+from arbiter.core.config import resolve_model
+from arbiter.core.metrics import ComparisonResult, ModelMetrics
+from arbiter.core.providers.factory import create_provider
+
+
+# Callback type for streaming progress updates
+# (model_name, token_text, metrics_snapshot)
+StreamCallback = Callable[[str, str, ModelMetrics], None]
+
+
+async def run_single_model(
+    model_spec: str,
+    prompt: str,
+    system: Optional[str] = None,
+    image_path: Optional[str] = None,
+    on_token: Optional[StreamCallback] = None,
+) -> ModelMetrics:
+    """Run a single model and collect metrics.
+
+    Args:
+        model_spec: Model specifier (e.g. "gemma4", "openai:gpt-4o")
+        prompt: The prompt to send
+        system: Optional system prompt
+        image_path: Optional image for multimodal models
+        on_token: Optional callback fired for each token
+
+    Returns:
+        ModelMetrics with all performance data and generated output
+    """
+    config = resolve_model(model_spec)
+    provider, model_name = create_provider(config)
+
+    metrics = ModelMetrics(
+        model=model_name,
+        provider=config.provider,
+    )
+
+    metrics.start()
+    last_meta: dict = {}
+
+    try:
+        async for chunk in provider.stream_generate(
+            model=model_name,
+            prompt=prompt,
+            system=system,
+            image_path=image_path,
+        ):
+            if chunk.text:
+                if metrics._first_token_time is None:
+                    metrics.record_first_token()
+                metrics.record_token(chunk.text)
+
+                if on_token:
+                    on_token(model_name, chunk.text, metrics)
+
+            if chunk.meta:
+                last_meta.update(chunk.meta)
+
+            if chunk.done:
+                break
+
+    except Exception as e:
+        metrics.output = f"[ERROR] {type(e).__name__}: {e}"
+
+    metrics.finish(provider_meta=last_meta)
+    return metrics
+
+
+async def run_comparison(
+    model_specs: list[str],
+    prompt: str,
+    system: Optional[str] = None,
+    image_path: Optional[str] = None,
+    on_token: Optional[StreamCallback] = None,
+    sequential: bool = False,
+) -> ComparisonResult:
+    """Run the same prompt through multiple models.
+
+    Args:
+        model_specs: List of model specifiers
+        prompt: The prompt to send to all models
+        system: Optional system prompt
+        image_path: Optional image path for multimodal
+        on_token: Optional callback for streaming updates
+        sequential: If True, run models one at a time (saves memory).
+                    If False, run all in parallel (faster but uses more RAM).
+
+    Returns:
+        ComparisonResult with metrics for all models
+    """
+    if sequential:
+        return await _run_sequential(model_specs, prompt, system, image_path, on_token)
+    else:
+        return await _run_parallel(model_specs, prompt, system, image_path, on_token)
+
+
+async def _run_parallel(
+    model_specs: list[str],
+    prompt: str,
+    system: Optional[str],
+    image_path: Optional[str],
+    on_token: Optional[StreamCallback],
+) -> ComparisonResult:
+    """Run all models in parallel (faster, more memory)."""
+    tasks = [
+        run_single_model(
+            model_spec=spec,
+            prompt=prompt,
+            system=system,
+            image_path=image_path,
+            on_token=on_token,
+        )
+        for spec in model_specs
+    ]
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    model_metrics = []
+    for i, result in enumerate(results):
+        if isinstance(result, Exception):
+            metrics = ModelMetrics(
+                model=model_specs[i],
+                provider="unknown",
+                output=f"[ERROR] {type(result).__name__}: {result}",
+            )
+            model_metrics.append(metrics)
+        else:
+            model_metrics.append(result)
+
+    return ComparisonResult(
+        prompt=prompt,
+        models=model_metrics,
+        timestamp=datetime.now(timezone.utc).isoformat(),
+    )
+
+
+async def _run_sequential(
+    model_specs: list[str],
+    prompt: str,
+    system: Optional[str],
+    image_path: Optional[str],
+    on_token: Optional[StreamCallback],
+) -> ComparisonResult:
+    """Run models one at a time. Each finishes and frees memory before the next starts.
+
+    Better for low-RAM machines (8GB). Timing is still fair because each model
+    gets the machine to itself with no contention.
+    """
+    model_metrics = []
+
+    for spec in model_specs:
+        try:
+            metrics = await run_single_model(
+                model_spec=spec,
+                prompt=prompt,
+                system=system,
+                image_path=image_path,
+                on_token=on_token,
+            )
+            model_metrics.append(metrics)
+        except Exception as e:
+            metrics = ModelMetrics(
+                model=spec,
+                provider="unknown",
+                output=f"[ERROR] {type(e).__name__}: {e}",
+            )
+            model_metrics.append(metrics)
+
+    return ComparisonResult(
+        prompt=prompt,
+        models=model_metrics,
+        timestamp=datetime.now(timezone.utc).isoformat(),
+    )
+
+
+async def stream_comparison(
+    model_specs: list[str],
+    prompt: str,
+    system: Optional[str] = None,
+    image_path: Optional[str] = None,
+) -> AsyncIterator[tuple[str, str, ModelMetrics | None]]:
+    """Stream a comparison, yielding (event_type, model_name, data) tuples.
+
+    Event types:
+        "start"    - model generation started
+        "token"    - a token was received
+        "done"     - model generation complete
+        "error"    - model encountered an error
+
+    This is used by the WebSocket server for real-time dashboard updates.
+    """
+    queues: dict[str, asyncio.Queue] = {}
+    tasks: list[asyncio.Task] = []
+
+    async def _run_model(spec: str, queue: asyncio.Queue):
+        config = resolve_model(spec)
+        provider, model_name = create_provider(config)
+        metrics = ModelMetrics(model=model_name, provider=config.provider)
+
+        await queue.put(("start", model_name, metrics))
+        metrics.start()
+        last_meta: dict = {}
+
+        try:
+            async for chunk in provider.stream_generate(
+                model=model_name, prompt=prompt, system=system, image_path=image_path
+            ):
+                if chunk.text:
+                    if metrics._first_token_time is None:
+                        metrics.record_first_token()
+                    metrics.record_token(chunk.text)
+                    await queue.put(("token", model_name, metrics))
+
+                if chunk.meta:
+                    last_meta.update(chunk.meta)
+
+                if chunk.done:
+                    break
+
+            metrics.finish(provider_meta=last_meta)
+            await queue.put(("done", model_name, metrics))
+
+        except Exception as e:
+            metrics.output = f"[ERROR] {type(e).__name__}: {e}"
+            metrics.finish()
+            await queue.put(("error", model_name, metrics))
+
+        await queue.put(None)  # sentinel
+
+    # Start all models
+    merged_queue: asyncio.Queue = asyncio.Queue()
+    for spec in model_specs:
+        task = asyncio.create_task(_run_model(spec, merged_queue))
+        tasks.append(task)
+
+    finished = 0
+    total = len(model_specs)
+
+    while finished < total:
+        item = await merged_queue.get()
+        if item is None:
+            finished += 1
+            continue
+        yield item
+
+    # Ensure all tasks are cleaned up
+    for task in tasks:
+        if not task.done():
+            task.cancel()

arbiter/core/swe/__init__.py

@@ -0,0 +1 @@
+"""SWE-bench style testing -- real code execution in Docker containers."""

arbiter/core/swe/container.py

@@ -0,0 +1,158 @@
+"""Docker container management for safe code execution."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+
+@dataclass
+class ContainerResult:
+    """Result from running tests in a Docker container."""
+    exit_code: int
+    stdout: str
+    stderr: str
+    tests_passed: int
+    tests_failed: int
+    tests_total: int
+    duration_s: float
+
+    @property
+    def all_passed(self) -> bool:
+        return self.tests_failed == 0 and self.tests_passed > 0
+
+
+def check_docker() -> bool:
+    """Check if Docker is available and running."""
+    try:
+        result = subprocess.run(
+            ["docker", "info"],
+            capture_output=True, text=True, timeout=10,
+        )
+        return result.returncode == 0
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return False
+
+
+def build_test_image(test_pack_dir: Path, tag: str = "arbiter-test") -> bool:
+    """Build a Docker image for a test pack.
+
+    The test pack directory must contain a Dockerfile.
+    """
+    dockerfile = test_pack_dir / "Dockerfile"
+    if not dockerfile.exists():
+        # Generate a default Python Dockerfile
+        dockerfile.write_text(
+            "FROM python:3.11-slim\n"
+            "WORKDIR /workspace\n"
+            "COPY requirements.txt* ./\n"
+            "RUN pip install --no-cache-dir -r requirements.txt 2>/dev/null || true\n"
+            "RUN pip install --no-cache-dir pytest\n"
+            "COPY . .\n"
+        )
+
+    result = subprocess.run(
+        ["docker", "build", "-t", tag, "."],
+        cwd=str(test_pack_dir),
+        capture_output=True, text=True, timeout=120,
+    )
+    return result.returncode == 0
+
+
+def run_in_container(
+    image: str,
+    fix_code: str,
+    test_code: str,
+    test_command: str = "python -m pytest /workspace/test_fix.py -v --tb=short",
+    timeout: int = 60,
+) -> ContainerResult:
+    """Run a model's fix inside a Docker container against a test suite.
+
+    1. Writes the model's fix to a temp dir as fix.py
+    2. Writes the test suite as test_fix.py
+    3. Mounts both into the container
+    4. Runs pytest
+    5. Returns results
+
+    Network is disabled. Memory limited to 512MB. CPU limited to 1 core.
+    """
+    import time
+
+    tmpdir = tempfile.mkdtemp(prefix="arbiter-swe-")
+    fix_path = Path(tmpdir) / "fix.py"
+    test_path = Path(tmpdir) / "test_fix.py"
+
+    fix_path.write_text(fix_code)
+    test_path.write_text(test_code)
+
+    start = time.perf_counter()
+
+    try:
+        cmd = [
+            "docker", "run", "--rm",
+            "--network=none",
+            "--memory=512m",
+            "--cpus=1",
+            "-v", f"{tmpdir}:/workspace",
+            image,
+            "sh", "-c", test_command,
+        ]
+
+        result = subprocess.run(
+            cmd, capture_output=True, text=True, timeout=timeout,
+        )
+
+        elapsed = time.perf_counter() - start
+        stdout = result.stdout
+        stderr = result.stderr
+
+        passed, failed, total = _parse_pytest_output(stdout + stderr)
+
+        return ContainerResult(
+            exit_code=result.returncode,
+            stdout=stdout[-2000:],
+            stderr=stderr[-1000:],
+            tests_passed=passed,
+            tests_failed=failed,
+            tests_total=total,
+            duration_s=elapsed,
+        )
+
+    except subprocess.TimeoutExpired:
+        return ContainerResult(
+            exit_code=-1, stdout="", stderr="Timeout",
+            tests_passed=0, tests_failed=0, tests_total=0,
+            duration_s=time.perf_counter() - start,
+        )
+    finally:
+        import shutil
+        shutil.rmtree(tmpdir, ignore_errors=True)
+
+
+def _parse_pytest_output(output: str) -> tuple[int, int, int]:
+    """Parse pytest output to extract pass/fail counts."""
+    import re
+
+    # Look for "X passed, Y failed" pattern
+    passed = 0
+    failed = 0
+
+    # pytest summary line: "5 passed, 2 failed in 1.23s"
+    match = re.search(r"(\d+) passed", output)
+    if match:
+        passed = int(match.group(1))
+
+    match = re.search(r"(\d+) failed", output)
+    if match:
+        failed = int(match.group(1))
+
+    match = re.search(r"(\d+) error", output)
+    if match:
+        failed += int(match.group(1))
+
+    total = passed + failed
+    return passed, failed, total

arbiter/core/swe/runner.py

@@ -0,0 +1,220 @@
+"""SWE test runner -- sends buggy code to models, verifies fixes in Docker."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+from arbiter.core.config import resolve_model
+from arbiter.core.providers.factory import create_provider
+from arbiter.core.swe.container import ContainerResult, check_docker, run_in_container
+from arbiter.core.swe.sandbox import run_in_sandbox, check_sandbox
+from arbiter.core.swe.test_packs import BUILT_IN_PACKS, TestPack, TestCase
+
+
+@dataclass
+class SWETestResult:
+    """Result from one model attempting one SWE test case."""
+    model: str
+    test_case: str
+    category: str
+    passed: bool
+    tests_passed: int
+    tests_total: int
+    model_response: str = ""
+    container_output: str = ""
+    duration_s: float = 0.0
+    error: Optional[str] = None
+
+
+@dataclass
+class SWESuiteResult:
+    """Results from running a full SWE suite against one model."""
+    model: str
+    results: list[SWETestResult] = field(default_factory=list)
+    total_passed: int = 0
+    total_tests: int = 0
+    total_duration_s: float = 0.0
+    categories: dict = field(default_factory=dict)  # category -> {passed, total}
+
+    def compute(self) -> None:
+        self.total_passed = sum(1 for r in self.results if r.passed)
+        self.total_tests = len(self.results)
+        self.total_duration_s = sum(r.duration_s for r in self.results)
+        cats: dict[str, dict] = {}
+        for r in self.results:
+            if r.category not in cats:
+                cats[r.category] = {"passed": 0, "total": 0}
+            cats[r.category]["total"] += 1
+            if r.passed:
+                cats[r.category]["passed"] += 1
+        self.categories = cats
+
+    @property
+    def pass_rate(self) -> float:
+        return self.total_passed / self.total_tests if self.total_tests > 0 else 0
+
+    def to_dict(self) -> dict:
+        return {
+            "model": self.model,
+            "total_passed": self.total_passed,
+            "total_tests": self.total_tests,
+            "pass_rate": round(self.pass_rate * 100, 1),
+            "total_duration_s": round(self.total_duration_s, 1),
+            "categories": self.categories,
+            "results": [
+                {
+                    "test_case": r.test_case,
+                    "category": r.category,
+                    "passed": r.passed,
+                    "tests_passed": r.tests_passed,
+                    "tests_total": r.tests_total,
+                    "duration_s": round(r.duration_s, 2),
+                    "error": r.error,
+                }
+                for r in self.results
+            ],
+        }
+
+
+SWE_SYSTEM_PROMPT = (
+    "You are a senior software engineer fixing bugs in a codebase. "
+    "You will be given the current code and a bug report. "
+    "Return ONLY the complete fixed file. No explanation, no markdown fencing, "
+    "just the corrected Python code."
+)
+
+
+async def run_swe_test(
+    model_spec: str,
+    test_case: TestCase,
+    use_docker: bool = False,
+) -> SWETestResult:
+    """Run one SWE test case against one model.
+
+    1. Send the buggy code + issue to the model
+    2. Get the model's fix
+    3. Run the fix in a Docker container against the test suite
+    4. Report pass/fail
+    """
+    start = time.perf_counter()
+
+    # Build prompt
+    prompt = (
+        f"## Bug Report\n{test_case.issue}\n\n"
+        f"## Current Code ({test_case.filename})\n"
+        f"```python\n{test_case.buggy_code}\n```\n\n"
+        f"Fix the bug. Return ONLY the complete corrected file."
+    )
+
+    # Get model's fix
+    config = resolve_model(model_spec)
+    provider, model_name = create_provider(config)
+
+    parts = []
+    try:
+        async for chunk in provider.stream_generate(
+            model=model_name, prompt=prompt, system=SWE_SYSTEM_PROMPT,
+        ):
+            parts.append(chunk.text)
+            if chunk.done:
+                break
+    except Exception as e:
+        return SWETestResult(
+            model=model_spec, test_case=test_case.name,
+            category=test_case.category, passed=False,
+            tests_passed=0, tests_total=test_case.expected_tests,
+            error=f"Model error: {e}",
+            duration_s=time.perf_counter() - start,
+        )
+
+    model_response = "".join(parts).strip()
+
+    # Clean markdown fencing if present
+    import re
+    if "```" in model_response:
+        match = re.search(r"```(?:python)?\s*(.*?)```", model_response, re.DOTALL)
+        if match:
+            model_response = match.group(1).strip()
+
+    # Run tests -- sandbox (default) or Docker
+    try:
+        if use_docker and check_docker():
+            result = run_in_container(
+                image=test_case.docker_image,
+                fix_code=model_response,
+                test_code=test_case.test_code,
+                test_command=test_case.test_command,
+                timeout=test_case.timeout,
+            )
+        else:
+            result = run_in_sandbox(
+                fix_code=model_response,
+                test_code=test_case.test_code,
+                timeout=test_case.timeout,
+            )
+
+        return SWETestResult(
+            model=model_spec, test_case=test_case.name,
+            category=test_case.category,
+            passed=result.all_passed,
+            tests_passed=result.tests_passed,
+            tests_total=result.tests_total or test_case.expected_tests,
+            model_response=model_response[:1000],
+            container_output=result.stdout[:500],
+            duration_s=time.perf_counter() - start,
+        )
+
+    except Exception as e:
+        return SWETestResult(
+            model=model_spec, test_case=test_case.name,
+            category=test_case.category, passed=False,
+            tests_passed=0, tests_total=test_case.expected_tests,
+            error=f"Execution error: {e}",
+            model_response=model_response[:500],
+            duration_s=time.perf_counter() - start,
+        )
+
+
+async def run_swe_suite(
+    model_spec: str,
+    test_pack: Optional[str] = None,
+    use_docker: bool = False,
+    on_progress=None,
+) -> SWESuiteResult:
+    """Run the full SWE suite against one model."""
+    packs = BUILT_IN_PACKS if test_pack is None else [
+        p for p in BUILT_IN_PACKS if p.name == test_pack
+    ]
+
+    all_cases = []
+    for pack in packs:
+        all_cases.extend(pack.cases)
+
+    result = SWESuiteResult(model=model_spec)
+
+    for i, case in enumerate(all_cases):
+        if on_progress:
+            on_progress(model_spec, i + 1, len(all_cases), case.name)
+
+        test_result = await run_swe_test(model_spec, case, use_docker=use_docker)
+        result.results.append(test_result)
+
+    result.compute()
+    return result
+
+
+async def run_swe_comparison(
+    model_specs: list[str],
+    test_pack: Optional[str] = None,
+    use_docker: bool = False,
+) -> list[SWESuiteResult]:
+    """Run SWE suite against multiple models (sequentially for fairness)."""
+    results = []
+    for spec in model_specs:
+        r = await run_swe_suite(spec, test_pack=test_pack, use_docker=use_docker)
+        results.append(r)
+    return results