# Copyright (c) Microsoft. All rights reserved. """Instrumented driver for running Claude Code on SWE-bench with Agent-lightning. This script wires together the Lightning Store, LLM proxy, and Claude Code controller so that every SWE-bench instance is executed inside the official Claude container while capturing full Agent-lightning traces. It supports three backend modes: - `vllm`: wrap an OpenAI-compatible endpoint (e.g., vLLM) for hosted OSS models while collecting prompt/response token ids and logprobs. - `anthropic`: call the official Claude Code API via `ANTHROPIC_API_KEY` for prompt tuning. Backend model defaults to the provided frontend names. - `openai`: route through any OpenAI-compatible provider using `OPENAI_API_KEY`. Typical usage: hosted vLLM (requires model paths and --base-url) ```bash # Run vLLM in background vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \ --max-model-len 131072 \ --enable-auto-tool-choice \ --tool-call-parser qwen3_coder \ --port 45993 & python claude_code_agent.py vllm \ --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct \ --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct \ --base-url http://localhost:45993/v1 \ --dataset-path swebench_samples.jsonl \ ``` Official Claude Code via Anthropic: ```bash export ANTHROPIC_API_KEY=sk-... python claude_code_agent.py anthropic \ --dataset-path swebench_samples.jsonl \ --output-dir data_anthropic ``` Any OpenAI-compatible backend: ```bash export OPENAI_API_KEY=sk-... python claude_code_agent.py openai \ --backend-model-high gpt-5.1-codex-mini \ --backend-model-low gpt-4.1-mini \ --dataset-path swebench_samples.jsonl ``` Use `--debug` to enable debug loggings. """ import asyncio import json import logging import os import resource from argparse import ArgumentParser from typing import Any, Dict, List, Literal, Optional, Sequence, cast from claude_code_controller import ClaudeController from datasets import Dataset from extended_adapter import ExtendedLlmProxyTraceToTriplet from swebench.harness.constants import SWEbenchInstance from swebench.harness.utils import load_swebench_dataset # pyright: ignore[reportUnknownVariableType] from swebench_utils.evaluation import evaluate from swebench_utils.logging import log_for_evaluation from transformers import AutoTokenizer, PreTrainedTokenizerBase from agentlightning import ( InMemoryLightningStore, LightningStoreServer, LitAgentRunner, OtelTracer, setup_logging, setup_module_logging, ) from agentlightning.litagent import LitAgent from agentlightning.llm_proxy import LLMProxy, ModelConfig from agentlightning.store import LightningStore from agentlightning.types import AttemptedRollout, NamedResources, ProxyLLM, Rollout, RolloutRawResult, Span logger = logging.getLogger("claude_code_agent") def _load_dataset(path: str, epoch: int = 0, limit: Optional[int] = None) -> List[SWEbenchInstance]: instances: List[SWEbenchInstance] = [] with open(path) as f: for line in f: instance = json.loads(line) instance["epoch"] = epoch instances.append(instance) if limit is not None: instances = instances[:limit] return instances def _flatten_messages(messages: List[Any]) -> List[Dict[str, str]]: flattened: List[Dict[str, str]] = [] for msg in messages: if msg["role"] in ["system", "user"] and isinstance(msg["content"], list): msg_content: List[str] = [] for content in msg["content"]: msg_content.append(content["text"]) msg["content"] = "".join(msg_content) elif msg["role"] == "assistant" and "tool_calls" in msg: # NOTE: # Tool calls are list of dict, though in most case only one tool call is made per call # We serialize it as json string here to avoid nested structure msg["tool_calls"] = json.dumps(msg["tool_calls"]) for k in msg: assert isinstance(msg[k], str), f"\n>>> {msg}" flattened.append(msg) return flattened class ClaudeCodeAgent(LitAgent[SWEbenchInstance]): """Claude Code Agent implementation. This agent is a wrapper of the Claude Code controller, and it should be used to run the Claude Code agent on SWE-bench datasets. """ def __init__( self, namespace: Literal["swebench", "starryzhang"] = "swebench", max_turns: int = 5, run_method: Literal["python", "cli"] = "cli", open_file_limit: int = 4096, cache_level: str = "env", # ["none", "base", "env", "instance"] clean: bool = False, force_rebuild: bool = False, timeout: int = 1_800, # in sec instance_image_tag: str = "latest", rewrite_reports: bool = False, swebench_full_dataset: Optional[List[SWEbenchInstance]] = None, ) -> None: super().__init__() self.namespace = namespace self.max_turns = max_turns self.run_method = run_method self.cache_level = cache_level self.clean = clean self.force_rebuild = force_rebuild self.timeout = timeout self.instance_image_tag = instance_image_tag self.rewrite_reports = rewrite_reports self.swebench_full_dataset = ( {each["instance_id"]: each for each in swebench_full_dataset} if swebench_full_dataset is not None else {} ) # Set the maximum number of open files to the specified limit. resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit)) async def rollout_async( self, task: SWEbenchInstance, resources: NamedResources, rollout: Rollout ) -> RolloutRawResult: if not isinstance(rollout, AttemptedRollout): # Technically, rollout should be an AttemptedRollout here. # but the API is not stabilized yet. raise ValueError("Rollout is not an AttemptedRollout.") run_id = f"epoch_{task.get('epoch', 0)}" image = f"{self.namespace}/sweb.eval.x86_64.{task['instance_id'].lower()}".replace("__", "_1776_") llm = cast(ProxyLLM, resources["llm"]) try: # 1. init container controller = ClaudeController( image, task, run_id, llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id), llm.api_key or os.environ.get("ANTHROPIC_AUTH_TOKEN", "dummy"), ) # 2. execute task prediction = controller.run_instance( task, max_turns=self.max_turns, run_method=cast(Literal["python", "cli"], self.run_method) ) del controller except Exception as e: log_for_evaluation(run_id, task["instance_id"], f"Exception during rollout: {e}") return 0.0 # 3. obtain rewards (evaluation result) reward = 0.0 # empty patch if prediction["model_patch"] in ["", None]: return reward instance_id = prediction["instance_id"] result = evaluate( cast(Any, prediction), self.swebench_full_dataset[instance_id], self.cache_level, self.clean, self.force_rebuild, run_id, self.timeout, namespace=self.namespace, instance_image_tag=self.instance_image_tag, rewrite_reports=self.rewrite_reports, ) # error patch if result is None: return reward report = result[1] # resolved/unresolved patch if report[instance_id]["resolved"]: reward = 1.0 return reward def sanity_check_spans(spans: Sequence[Span]) -> None: assert len(spans) > 1, f"At least two spans are expected for a valid rollout. Found {len(spans)} spans." assert any(span.name == "raw_gen_ai_request" for span in spans), "raw_gen_ai_request span not found" assert any(span.name == "agentlightning.annotation" for span in spans), "agentlightning.annotation span not found" async def run_instance_async( instance: SWEbenchInstance, agent: ClaudeCodeAgent, runner: LitAgentRunner[SWEbenchInstance], store: LightningStore, output_dir: Optional[str], adapter: Optional[ExtendedLlmProxyTraceToTriplet], tokenizer: Optional[PreTrainedTokenizerBase], ) -> None: """Runs the agent on a specific SWE-bench instance. Running on specific SWE-bench instance and queries the traced spans. It then extracts the triplets and saves the dataset. """ instance_id = instance["instance_id"] logger.info(f"Starting to run instance: {instance_id}") # Run the agent and query the traced spans. with runner.run_context(agent=agent, store=store): rollout = await runner.step(instance) logger.info(f"Finished running instance: {instance_id}") spans = await store.query_spans(rollout.rollout_id) if output_dir is None: logger.info(f"Generated {len(spans)} spans for {instance_id}") return # 1. Dump raw spans (Common for both types) raw_path = os.path.join(output_dir, f"stream_{instance_id}.json") with open(raw_path, "w") as f: for span in spans: f.write(json.dumps(span.model_dump()) + "\n") logger.info(f"Dumped {len(spans)} spans to {raw_path}") # 2. Extract Triplets and Save Dataset (vLLM specific) if adapter is not None and tokenizer is not None: try: triplets = adapter.adapt(cast(List[Span], spans)) logger.info(f"Extracted {len(triplets)} triplets for {instance_id}") all_triplets: List[Dict[str, Any]] = [] recent_reward: Optional[float] = None # Process in reverse to propagate rewards if necessary/logic dictates for triplet in reversed(triplets): if triplet.reward is not None: recent_reward = triplet.reward prompt_text = tokenizer.decode(triplet.prompt["token_ids"]) # type: ignore all_triplets.append( { "repo": instance.get("repo", ""), "instance_id": instance_id, "turn": triplet.metadata["sequence_id"], "prompt_ids": triplet.prompt["token_ids"], "gold_completion_ids": triplet.response["token_ids"], "logprobs": triplet.response["logprobs"], "reward": recent_reward, "prompt": prompt_text, "messages": _flatten_messages(triplet.metadata["messages"]), } ) if all_triplets: ds = Dataset.from_list(all_triplets) # type: ignore save_path = os.path.join(output_dir, f"dataset-{instance_id}") ds.save_to_disk(save_path) # type: ignore logger.info(f"Saved HuggingFace dataset to {save_path}") except Exception as e: logger.error(f"Failed to extract triplets for {instance_id}: {e}") logger.info(f"Finished extracting spans and traces for instance: {instance_id}") # Quickly sanity check the spans sanity_check_spans(spans) logger.info(f"Sanity check passed for instance: {instance_id}") async def dry_run_claude_code( *, dataset_path: str, haiku_frontend_name: str, haiku_backend_name: str, sonnet_frontend_name: str, sonnet_backend_name: str, backend_type: Literal["vllm", "anthropic", "openai"], api_base_url: Optional[str], output_dir: Optional[str], max_turns: int, limit: Optional[int], cooldown_seconds: float, ) -> None: """Executes a dry run of the Claude Code agent on a dataset. This function handles both 'official' runs (interacting with Anthropic APIs) and 'hosted' runs (interacting with vLLM or compatible servers). It manages initialization of the Lightning Store, LLM Proxy, and the execution loop. If running in 'vllm' mode, it will also attempt to extract triplets using the provided backend name as the tokenizer path and save a HuggingFace Dataset. Args: dataset_path: Path to the JSONL dataset file. haiku_frontend_name: The model name used in the code to request the 'fast' model. haiku_backend_name: The actual model name/path on the backend. sonnet_frontend_name: The model name used in the code to request the 'strong' model. sonnet_backend_name: The actual model name/path on the backend. backend_type: The type of backend to configure ("vllm", "anthropic" or "openai"). api_base_url: Base URL for the API. Required for "vllm" or "openai". output_dir: Directory to save logs, spans, and datasets. max_turns: Maximum number of steps the agent can take per instance. limit: Optional limit on the number of instances to process. """ dataset = _load_dataset(dataset_path, limit=limit) # Initialize Infrastructure tracer = OtelTracer() runner = LitAgentRunner[SWEbenchInstance](tracer) store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654) await store.start() # Enable callbacks for training data extraction if using vLLM callbacks = ["return_token_ids", "opentelemetry", "logprobs"] if backend_type == "vllm" else ["opentelemetry"] llm_proxy = LLMProxy(port=12358, store=store, callbacks=callbacks) # Configure Models based on backend type model_configs: List[ModelConfig] = [] model_params: Dict[str, Any] = {} if backend_type == "vllm": model_namespace = "hosted_vllm" if api_base_url: model_params["api_base"] = api_base_url else: raise ValueError("api_base_url is required for vllm backend") elif backend_type == "anthropic": model_namespace = "anthropic" model_params["api_key"] = "os.environ/ANTHROPIC_API_KEY" if api_base_url: model_params["api_base"] = api_base_url elif backend_type == "openai": model_namespace = "openai" model_params["api_key"] = "os.environ/OPENAI_API_KEY" if api_base_url: # Users can still override this via environment variables, # even if they don't pass it in as an argument. model_params["api_base"] = api_base_url model_configs.extend( [ ModelConfig( model_name=sonnet_frontend_name, litellm_params={ "model": f"{model_namespace}/{sonnet_backend_name}", **model_params, }, ), ModelConfig( model_name=haiku_frontend_name, litellm_params={ "model": f"{model_namespace}/{haiku_backend_name}", **model_params, }, ), ] ) logger.info(f"Updating model list: {model_configs}") llm_proxy.update_model_list(model_configs) await llm_proxy.start() try: # Add the LLM proxy as a resource to the store await store.add_resources({"llm": llm_proxy.as_resource(model="local")}) # Prepare for triplet extraction if vllm adapter = ExtendedLlmProxyTraceToTriplet() if backend_type == "vllm" else None tokenizer = None if backend_type == "vllm": try: tokenizer = AutoTokenizer.from_pretrained(sonnet_backend_name) # type: ignore except Exception as e: logger.warning(f"Could not load tokenizer for {sonnet_backend_name}: {e}") # Load full swebench dataset. Mainly for evaluation purposes. swebench_full_dataset = load_swebench_dataset("princeton-nlp/SWE-bench", split="test") # Initialize Claude Code Agent claude_code_agent = ClaudeCodeAgent(swebench_full_dataset=swebench_full_dataset, max_turns=max_turns) # Execution Loop for instance in dataset: await run_instance_async( instance, claude_code_agent, runner, store, output_dir, adapter, cast(PreTrainedTokenizerBase, tokenizer), ) # Basic sleep to allow resource cleanup or rate limit cooling await asyncio.sleep(cooldown_seconds) finally: await llm_proxy.stop() await store.stop() if __name__ == "__main__": parser = ArgumentParser(description="Run Claude Code Agent experiments.") # Backend Selection parser.add_argument( "backend_type", type=str, choices=["vllm", "anthropic", "openai"], help="Backend type: 'vllm' for hosted models, 'anthropic' for official API, 'openai' for OpenAI API.", ) # Model Configuration parser.add_argument( "--backend-model-high", type=str, default=None, help="Backend model path/name for expensive model usages (used as vLLM model name / OpenAI model name).", ) parser.add_argument( "--backend-model-low", type=str, default=None, help="Backend model path/name for low-price model usages (used as vLLM model name / OpenAI model name).", ) parser.add_argument( "--base-url", type=str, default="http://localhost:8000/v1", help="LLM server address (required for vllm)." ) # Frontend/Agent Configuration parser.add_argument( "--frontend-model-high", type=str, default="claude-sonnet-4-5-20250929", help="The frontend high-price model name provided to Claude Code.", ) parser.add_argument( "--frontend-model-low", type=str, default="claude-haiku-4-5-20251001", help="The frontend low-price model name provided to Claude Code.", ) # Execution Configuration parser.add_argument("--dataset-path", type=str, default="swebench_samples.jsonl", help="Path to the dataset.") parser.add_argument("--max-turns", type=int, default=5, help="Maximum turns per instance.") parser.add_argument("--output-dir", type=str, default="data", help="Directory to save output logs.") parser.add_argument("--limit", type=int, default=None, help="Limit the number of instances to run (for debugging).") parser.add_argument("--cooldown-seconds", type=float, default=2.0, help="Cooldown seconds between instances.") parser.add_argument("--debug", action="store_true", help="Enable debug loggings.") args = parser.parse_args() if args.output_dir is not None: os.makedirs(args.output_dir, exist_ok=True) if args.debug: setup_logging() setup_module_logging("DEBUG", name="claude_code_agent") else: setup_logging(apply_to=[logger.name]) # Map backend_type to the appropriate args backend_mode = cast(Literal["vllm", "anthropic", "openai"], args.backend_type) # If using anthropic, the backend name usually matches the frontend or is specific API string. # Otherwise, the backend name is the model name/path (e.g., Qwen/...) and must be provided. if args.backend_model_high is None: if args.backend_type == "anthropic": backend_model_high = args.frontend_model_high else: raise ValueError("--backend-model-high is required for non-anthropic backends") else: backend_model_high = args.backend_model_high if args.backend_model_low is None: if args.backend_type == "anthropic": backend_model_low = args.frontend_model_low else: raise ValueError("--backend-model-low is required for non-anthropic backends") else: backend_model_low = args.backend_model_low asyncio.run( dry_run_claude_code( dataset_path=args.dataset_path, haiku_frontend_name=args.frontend_model_low, haiku_backend_name=backend_model_low, sonnet_frontend_name=args.frontend_model_high, sonnet_backend_name=backend_model_high, backend_type=backend_mode, api_base_url=args.base_url if backend_mode == "vllm" else None, output_dir=args.output_dir, max_turns=args.max_turns, limit=args.limit, cooldown_seconds=args.cooldown_seconds, ) )