mini-swe-agent 1.1.0__py3-none-any.whl

minisweagent/run/extra/swebench.py

@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+
+"""Run mini-SWE-agent on SWEBench instances.
+
+[not dim]
+More information about the usage: [bold green]https://mini-swe-agent.com/latest/usage/swebench/[/bold green]
+[/not dim]
+"""
+
+import concurrent.futures
+import json
+import random
+import re
+import threading
+import time
+import traceback
+from pathlib import Path
+
+import typer
+import yaml
+from datasets import load_dataset
+from rich.live import Live
+
+from minisweagent.agents.default import DefaultAgent
+from minisweagent.config import builtin_config_dir, get_config_path
+from minisweagent.environments.docker import DockerEnvironment
+from minisweagent.models import get_model
+from minisweagent.run.extra.utils.batch_progress import RunBatchProgressManager
+from minisweagent.run.utils.save import save_traj
+
+app = typer.Typer(rich_markup_mode="rich", add_completion=False)
+
+DATASET_MAPPING = {
+    "full": "princeton-nlp/SWE-Bench",
+    "verified": "princeton-nlp/SWE-Bench_Verified",
+    "lite": "princeton-nlp/SWE-Bench_Lite",
+    "multimodal": "princeton-nlp/SWE-Bench_Multimodal",
+    "multilingual": "swe-bench/SWE-Bench_Multilingual",
+    "smith": "SWE-bench/SWE-smith",
+    "_test": "klieret/swe-bench-dummy-test-dataset",
+}
+
+
+_OUTPUT_FILE_LOCK = threading.Lock()
+
+
+class ProgressTrackingAgent(DefaultAgent):
+    """Simple wrapper around DefaultAgent that provides progress updates."""
+
+    def __init__(self, *args, progress_manager: RunBatchProgressManager, instance_id: str = "", **kwargs):
+        super().__init__(*args, **kwargs)
+        self.progress_manager: RunBatchProgressManager = progress_manager
+        self.instance_id = instance_id
+
+    def step(self) -> dict:
+        """Override step to provide progress updates."""
+        self.progress_manager.update_instance_status(
+            self.instance_id, f"Step {self.model.n_calls + 1:3d} (${self.model.cost:.2f})"
+        )
+        return super().step()
+
+
+def get_swebench_docker_image_name(instance: dict) -> str:
+    """Get the image name for a SWEBench instance."""
+    image_name = instance.get("image_name", None)
+    if image_name is None:
+        # Docker doesn't allow double underscore, so we replace them with a magic token
+        iid = instance["instance_id"]
+        id_docker_compatible = iid.replace("__", "_1776_")
+        image_name = f"swebench/sweb.eval.x86_64.{id_docker_compatible}:latest".lower()
+    return image_name
+
+
+def update_preds_file(output_path: Path, instance_id: str, model_name: str, result: str):
+    """Update the output JSON file with results from a single instance."""
+    with _OUTPUT_FILE_LOCK:
+        output_data = {}
+        if output_path.exists():
+            output_data = json.loads(output_path.read_text())
+        output_data[instance_id] = {
+            "model_name_or_path": model_name,
+            "instance_id": instance_id,
+            "model_patch": result,
+        }
+        output_path.write_text(json.dumps(output_data, indent=2))
+
+
+def remove_from_preds_file(output_path: Path, instance_id: str):
+    """Remove an instance from the predictions file."""
+    if not output_path.exists():
+        return
+    with _OUTPUT_FILE_LOCK:
+        output_data = json.loads(output_path.read_text())
+        if instance_id in output_data:
+            del output_data[instance_id]
+            output_path.write_text(json.dumps(output_data, indent=2))
+
+
+def process_instance(
+    instance: dict,
+    output_dir: Path,
+    model_name: str | None,
+    config_path: str | Path,
+    progress_manager: RunBatchProgressManager,
+) -> None:
+    """Process a single SWEBench instance."""
+    instance_id = instance["instance_id"]
+    instance_dir = output_dir / instance_id
+    # avoid inconsistent state if something here fails and there's leftover previous files
+    remove_from_preds_file(output_dir / "preds.json", instance_id)
+    (instance_dir / f"{instance_id}.traj.json").unlink(missing_ok=True)
+
+    image_name = get_swebench_docker_image_name(instance)
+    config = yaml.safe_load(get_config_path(config_path).read_text())
+    model = get_model(model_name, config=config.get("model", {}))
+    task = instance["problem_statement"]
+
+    progress_manager.on_instance_start(instance_id)
+    progress_manager.update_instance_status(instance_id, "Pulling/starting docker")
+
+    agent = None
+    extra_info = None
+
+    try:
+        env = DockerEnvironment(**(config.get("environment", {}) | {"image": image_name}))
+        agent = ProgressTrackingAgent(
+            model,
+            env,
+            progress_manager=progress_manager,
+            instance_id=instance_id,
+            **config.get("agent", {}),
+        )
+        exit_status, result = agent.run(task)
+    except Exception as e:
+        print(f"Error processing instance {instance_id}: {e}\n{traceback.format_exc()}")
+        exit_status, result = type(e).__name__, str(e)
+        extra_info = {"traceback": traceback.format_exc()}
+    finally:
+        save_traj(
+            agent,
+            instance_dir / f"{instance_id}.traj.json",
+            exit_status=exit_status,
+            result=result,
+            extra_info=extra_info,
+            instance_id=instance_id,
+        )
+        update_preds_file(output_dir / "preds.json", instance_id, model.config.model_name, result)
+        progress_manager.on_instance_end(instance_id, exit_status)
+
+
+def filter_instances(
+    instances: list[dict], *, filter_spec: str, slice_spec: str = "", shuffle: bool = False
+) -> list[dict]:
+    """Filter and slice a list of SWEBench instances."""
+    if shuffle:
+        instances = sorted(instances.copy(), key=lambda x: x["instance_id"])
+        random.seed(42)
+        random.shuffle(instances)
+    before_filter = len(instances)
+    instances = [instance for instance in instances if re.match(filter_spec, instance["instance_id"])]
+    if (after_filter := len(instances)) != before_filter:
+        print(f"Instance filter: {before_filter} -> {after_filter} instances")
+    if slice_spec:
+        values = [int(x) if x else None for x in slice_spec.split(":")]
+        instances = instances[slice(*values)]
+        if (after_slice := len(instances)) != before_filter:
+            print(f"Instance slice: {before_filter} -> {after_slice} instances")
+    return instances
+
+
+@app.command()
+def main(
+    subset: str = typer.Option("lite", "--subset", help="SWEBench subset to use or path to a dataset"),
+    split: str = typer.Option("dev", "--split", help="Dataset split"),
+    slice_spec: str = typer.Option("", "--slice", help="Slice specification (e.g., '0:5' for first 5 instances)"),
+    filter_spec: str = typer.Option("", "--filter", help="Filter instance IDs by regex"),
+    shuffle: bool = typer.Option(False, "--shuffle", help="Shuffle instances"),
+    output: str = typer.Option("", "-o", "--output", help="Output directory"),
+    workers: int = typer.Option(1, "-w", "--workers", help="Number of worker threads for parallel processing"),
+    model: str | None = typer.Option(None, "-m", "--model", help="Model to use"),
+    redo_existing: bool = typer.Option(False, "--redo-existing", help="Redo existing instances"),
+    config: Path = typer.Option(
+        builtin_config_dir / "extra" / "swebench.yaml", "-c", "--config", help="Path to a config file"
+    ),
+) -> None:
+    """Run mini-SWE-agent on SWEBench instances"""
+    dataset_path = DATASET_MAPPING.get(subset, subset)
+    print(f"Loading dataset {dataset_path}, split {split}...")
+    instances = list(load_dataset(dataset_path, split=split))
+
+    instances = filter_instances(instances, filter_spec=filter_spec, slice_spec=slice_spec, shuffle=shuffle)
+    output_path = Path(output)
+    if not redo_existing and (output_path / "preds.json").exists():
+        existing_instances = list(json.loads((output_path / "preds.json").read_text()).keys())
+        print(f"Skipping {len(existing_instances)} existing instances")
+        instances = [instance for instance in instances if instance["instance_id"] not in existing_instances]
+
+    output_path.mkdir(parents=True, exist_ok=True)
+    print(f"Running on {len(instances)} instances...")
+    print(f"Results will be saved to {output_path}")
+
+    progress_manager = RunBatchProgressManager(len(instances), output_path / f"exit_statuses_{time.time()}.yaml")
+
+    def process_futures(futures: dict[concurrent.futures.Future, str]):
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                future.result()
+            except concurrent.futures.CancelledError:
+                pass
+            except Exception as e:
+                instance_id = futures[future]
+                print(f"Error in future for instance {instance_id}: {e}")
+                traceback.print_exc()
+                progress_manager.on_uncaught_exception(instance_id, e)
+
+    with Live(progress_manager.render_group, refresh_per_second=4):
+        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = {
+                executor.submit(process_instance, instance, output_path, model, config, progress_manager): instance[
+                    "instance_id"
+                ]
+                for instance in instances
+            }
+            try:
+                process_futures(futures)
+            except KeyboardInterrupt:
+                print("Cancelling all pending jobs. Press ^C again to exit immediately.")
+                for future in futures:
+                    if not future.running() and not future.done():
+                        future.cancel()
+                process_futures(futures)
+
+
+if __name__ == "__main__":
+    app()

minisweagent/run/extra/swebench_single.py

@@ -0,0 +1,53 @@
+"""Run on a single SWE-Bench instance."""
+
+from pathlib import Path
+
+import typer
+import yaml
+from datasets import load_dataset
+
+from minisweagent.agents.interactive import InteractiveAgent
+from minisweagent.config import builtin_config_dir, get_config_path
+from minisweagent.environments.docker import DockerEnvironment
+from minisweagent.models import get_model
+from minisweagent.run.extra.swebench import DATASET_MAPPING, get_swebench_docker_image_name
+
+app = typer.Typer(add_completion=False)
+
+
+@app.command()
+def main(
+    subset: str = typer.Option("lite", "--subset", help="SWEBench subset to use or path to a dataset"),
+    split: str = typer.Option("dev", "--split", help="Dataset split"),
+    instance_spec: str = typer.Option(None, "-i", "--instance", help="SWE-Bench instance ID"),
+    model_name: str | None = typer.Option(None, "-m", "--model", help="Model to use"),
+    config_path: Path = typer.Option(
+        builtin_config_dir / "extra" / "swebench.yaml", "-c", "--config", help="Path to a config file"
+    ),
+) -> None:
+    """Run on a single SWE-Bench instance."""
+    try:
+        dataset_path = DATASET_MAPPING[subset]
+    except KeyError:
+        dataset_path = subset
+    print(f"Loading dataset {dataset_path}, split {split}...")
+    instances = {
+        inst["instance_id"]: inst  # type: ignore
+        for inst in load_dataset(dataset_path, split=split)
+    }
+    if instance_spec.isnumeric():
+        instance_spec = sorted(instances.keys())[int(instance_spec)]
+    instance: dict = instances[instance_spec]  # type: ignore
+
+    _config = yaml.safe_load(get_config_path(config_path).read_text())
+    env = DockerEnvironment(**(_config.get("environment", {}) | {"image": get_swebench_docker_image_name(instance)}))
+    agent = InteractiveAgent(
+        get_model(model_name, _config.get("model", {})),
+        env,
+        **(_config.get("agent", {}) | {"mode": "yolo"}),
+    )
+    agent.run(instance["problem_statement"])
+
+
+if __name__ == "__main__":
+    app()

minisweagent/run/extra/utils/batch_progress.py

@@ -0,0 +1,164 @@
+"""This module contains an auxiliary class for rendering progress of a batch run.
+It's identical to the one used in swe-agent.
+"""
+
+import collections
+from pathlib import Path
+from threading import Lock
+
+import yaml
+from rich.console import Group
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from rich.table import Table
+
+import minisweagent.models
+
+
+def _shorten_str(s: str, max_len: int, shorten_left=False) -> str:
+    if not shorten_left:
+        s = s[: max_len - 3] + "..." if len(s) > max_len else s
+    else:
+        s = "..." + s[-max_len + 3 :] if len(s) > max_len else s
+    return f"{s:<{max_len}}"
+
+
+class RunBatchProgressManager:
+    def __init__(
+        self,
+        num_instances: int,
+        yaml_report_path: Path | None = None,
+    ):
+        """This class manages a progress bar/UI for run-batch
+
+        Args:
+            num_instances: Number of task instances
+            yaml_report_path: Path to save a yaml report of the instances and their exit statuses
+        """
+
+        self._spinner_tasks: dict[str, TaskID] = {}
+        """We need to map instance ID to the task ID that is used by the rich progress bar."""
+
+        self._lock = Lock()
+
+        self._instances_by_exit_status = collections.defaultdict(list)
+        self._main_progress_bar = Progress(
+            SpinnerColumn(spinner_name="dots2"),
+            TextColumn("[progress.description]{task.description} (${task.fields[total_cost]})"),
+            BarColumn(),
+            MofNCompleteColumn(),
+            TaskProgressColumn(),
+            TimeElapsedColumn(),
+            TextColumn("[cyan]eta:[/cyan]"),
+            TimeRemainingColumn(),
+            # Wait 5 min before estimating speed
+            speed_estimate_period=60 * 5,
+        )
+        self._task_progress_bar = Progress(
+            SpinnerColumn(spinner_name="dots2"),
+            TextColumn("{task.fields[instance_id]}"),
+            TextColumn("{task.fields[status]}"),
+            TimeElapsedColumn(),
+        )
+        """Task progress bar for individual instances. There's only one progress bar
+        with one task for each instance.
+        """
+
+        self._main_task_id = self._main_progress_bar.add_task(
+            "[cyan]Overall Progress", total=num_instances, total_cost="0.00"
+        )
+
+        self.render_group = Group(Table(), self._task_progress_bar, self._main_progress_bar)
+        self._yaml_report_path = yaml_report_path
+
+    @property
+    def n_completed(self) -> int:
+        return sum(len(instances) for instances in self._instances_by_exit_status.values())
+
+    def update_exit_status_table(self):
+        # We cannot update the existing table, so we need to create a new one and
+        # assign it back to the render group.
+        t = Table()
+        t.add_column("Exit Status")
+        t.add_column("Count", justify="right", style="bold cyan")
+        t.add_column("Most recent instances")
+        t.show_header = False
+        with self._lock:
+            t.show_header = True
+            # Sort by number of instances in descending order
+            sorted_items = sorted(self._instances_by_exit_status.items(), key=lambda x: len(x[1]), reverse=True)
+            for status, instances in sorted_items:
+                instances_str = _shorten_str(", ".join(reversed(instances)), 55)
+                t.add_row(status, str(len(instances)), instances_str)
+        assert self.render_group is not None
+        self.render_group.renderables[0] = t
+
+    def _update_total_costs(self) -> None:
+        with self._lock:
+            self._main_progress_bar.update(
+                self._main_task_id, total_cost=f"{minisweagent.models.GLOBAL_MODEL_STATS.cost:.2f}"
+            )
+
+    def update_instance_status(self, instance_id: str, message: str):
+        assert self._task_progress_bar is not None
+        assert self._main_progress_bar is not None
+        with self._lock:
+            self._task_progress_bar.update(
+                self._spinner_tasks[instance_id],
+                status=_shorten_str(message, 30),
+                instance_id=_shorten_str(instance_id, 25, shorten_left=True),
+            )
+        self._update_total_costs()
+
+    def on_instance_start(self, instance_id: str):
+        with self._lock:
+            self._spinner_tasks[instance_id] = self._task_progress_bar.add_task(
+                description=f"Task {instance_id}",
+                status="Task initialized",
+                total=None,
+                instance_id=instance_id,
+            )
+
+    def on_instance_end(self, instance_id: str, exit_status: str | None) -> None:
+        self._instances_by_exit_status[exit_status].append(instance_id)
+        with self._lock:
+            try:
+                self._task_progress_bar.remove_task(self._spinner_tasks[instance_id])
+            except KeyError:
+                pass
+            self._main_progress_bar.update(TaskID(0), advance=1)
+        self.update_exit_status_table()
+        self._update_total_costs()
+        if self._yaml_report_path is not None:
+            self._save_overview_data_yaml(self._yaml_report_path)
+
+    def on_uncaught_exception(self, instance_id: str, exception: Exception) -> None:
+        self.on_instance_end(instance_id, f"Uncaught {type(exception).__name__}")
+
+    def print_report(self) -> None:
+        """Print complete list of instances and their exit statuses."""
+        for status, instances in self._instances_by_exit_status.items():
+            print(f"{status}: {len(instances)}")
+            for instance in instances:
+                print(f"  {instance}")
+
+    def _get_overview_data(self) -> dict:
+        """Get data like exit statuses, total costs, etc."""
+        return {
+            # convert defaultdict to dict because of serialization
+            "instances_by_exit_status": dict(self._instances_by_exit_status),
+        }
+
+    def _save_overview_data_yaml(self, path: Path) -> None:
+        """Save a yaml report of the instances and their exit statuses."""
+        with self._lock:
+            path.write_text(yaml.dump(self._get_overview_data(), indent=4))

minisweagent/run/github_issue.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+import os
+from pathlib import Path
+
+import requests
+import typer
+import yaml
+from rich.console import Console
+
+from minisweagent.agents.interactive import InteractiveAgent
+from minisweagent.config import builtin_config_dir, get_config_path
+from minisweagent.environments.docker import DockerEnvironment
+from minisweagent.models import get_model
+from minisweagent.run.extra.config import configure_if_first_time
+from minisweagent.run.utils.save import save_traj
+
+DEFAULT_CONFIG = Path(os.getenv("MSWEA_GITHUB_CONFIG_PATH", builtin_config_dir / "github_issue.yaml"))
+console = Console(highlight=False)
+app = typer.Typer(rich_markup_mode="rich", add_completion=False)
+
+
+def fetch_github_issue(issue_url: str) -> str:
+    """Fetch GitHub issue text from the URL."""
+    # Convert GitHub issue URL to API URL
+    api_url = issue_url.replace("github.com", "api.github.com/repos").replace("/issues/", "/issues/")
+
+    headers = {}
+    if github_token := os.getenv("GITHUB_TOKEN"):
+        headers["Authorization"] = f"token {github_token}"
+
+    response = requests.get(api_url, headers=headers)
+    issue_data = response.json()
+
+    title = issue_data["title"]
+    body = issue_data["body"] or ""
+
+    return f"GitHub Issue: {title}\n\n{body}"
+
+
+@app.command()
+def main(
+    issue_url: str = typer.Option(prompt="Enter GitHub issue URL", help="GitHub issue URL"),
+    config: Path = typer.Option(DEFAULT_CONFIG, "-c", "--config", help="Path to config file"),
+    model: str | None = typer.Option(None, "-m", "--model", help="Model to use"),
+    yolo: bool = typer.Option(False, "-y", "--yolo", help="Run without confirmation"),
+) -> InteractiveAgent:
+    """Run mini-SWE-agent on a GitHub issue"""
+    configure_if_first_time()
+
+    _config = yaml.safe_load(get_config_path(config).read_text())
+    _agent_config = _config.get("agent", {})
+    if yolo:
+        _agent_config["mode"] = "yolo"
+
+    task = fetch_github_issue(issue_url)
+
+    agent = InteractiveAgent(
+        get_model(model, _config.get("model", {})),
+        DockerEnvironment(**_config.get("environment", {})),
+        **_agent_config,
+    )
+
+    repo_url = issue_url.split("/issues/")[0]
+    if github_token := os.getenv("GITHUB_TOKEN"):
+        repo_url = repo_url.replace("https://github.com/", f"https://{github_token}@github.com/") + ".git"
+
+    agent.env.execute(f"git clone {repo_url} /testbed", cwd="/")
+
+    exit_status, result = None, None
+    try:
+        exit_status, result = agent.run(task)
+    except KeyboardInterrupt:
+        console.print("\n[bold red]KeyboardInterrupt -- goodbye[/bold red]")
+    finally:
+        save_traj(agent, Path("traj.json"), exit_status=exit_status, result=result)
+    return agent
+
+
+if __name__ == "__main__":
+    app()

minisweagent/run/hello_world.py

@@ -0,0 +1,36 @@
+import os
+from pathlib import Path
+
+import typer
+import yaml
+
+from minisweagent import package_dir
+from minisweagent.agents.default import DefaultAgent
+from minisweagent.environments.local import LocalEnvironment
+from minisweagent.models.litellm_model import LitellmModel
+
+app = typer.Typer()
+
+
+@app.command()
+def main(
+    task: str = typer.Option(..., "-t", "--task", help="Task/problem statement", show_default=False, prompt=True),
+    model_name: str = typer.Option(
+        os.getenv("MSWEA_MODEL_NAME"),
+        "-m",
+        "--model",
+        help="Model name (defaults to MSWEA_MODEL_NAME env var)",
+        prompt="What model do you want to use?",
+    ),
+) -> DefaultAgent:
+    agent = DefaultAgent(
+        LitellmModel(model_name=model_name),
+        LocalEnvironment(),
+        **yaml.safe_load(Path(package_dir / "config" / "default.yaml").read_text())["agent"],
+    )
+    agent.run(task)
+    return agent
+
+
+if __name__ == "__main__":
+    app()