wehub/startrail-org--leann
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
CI / build (push) Has been cancelled
Details
Link Check / link-check (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:38:09 +08:00
commit 15dadb5432
263 changed files with 88651 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/contextbench/scripts/auto_run.py
+1034
View File
File diff suppressed because it is too large Load Diff
benchmarks/contextbench/scripts/batch_run_random.py
+140
View File
@@ -0,0 +1,140 @@
import json
import os
import random
import subprocess
import time
from pathlib import Path
from auto_run import prefetch_task_repositories
from datasets import load_dataset
ROOT = Path(__file__).resolve().parents[1]
OUTPUT_FILE = os.environ.get("OUTPUT_FILE", "all_predictions_claude.jsonl")
NUM_TASKS = int(os.environ.get("NUM_TASKS", "31"))
WORK_ROOT = os.environ.get("WORK_ROOT", "contextbench_work_dir_claude")
MODEL = os.environ.get("MODEL", os.environ.get("CLAUDE_MODEL", "")).strip()
DATASET_NAME = os.environ.get("DATASET_NAME", "Contextbench/ContextBench")
DATASET_SPLIT = os.environ.get("DATASET_SPLIT", "train")
# Optionally restrict to one benchmark split: Verified | Pro | Poly | Multi
BENCH_FILTER = os.environ.get("BENCH_FILTER", "Pro").strip()
# Optionally restrict random sampling to one repo (supports partial match),
# e.g. "django/django" or "sympy".
REPO_FILTER = os.environ.get("REPO_FILTER", "").strip().lower()
PREFETCH_REPOS = os.environ.get("PREFETCH_REPOS", "4").strip() != "0"
MITM_SCRIPT = ROOT / "mitmproxy_addons" / "trace_recorder.py"
TRACE_DIR = ROOT / "traces" / "raw"
def cleanup_residuals():
print("🧹 Cleaning up residual processes (Claude & Mitm)...")
try:
subprocess.run(["pkill", "-f", "claude"], stderr=subprocess.DEVNULL)
subprocess.run(["pkill", "-f", "mitmdump"], stderr=subprocess.DEVNULL)
time.sleep(2)
except Exception:
pass
def main():
Path(WORK_ROOT).mkdir(parents=True, exist_ok=True)
TRACE_DIR.mkdir(parents=True, exist_ok=True)
print(f"🧠 Model: {MODEL or '(Claude CLI default)'}")
print(f"🔎 Bench filter: {BENCH_FILTER or '(all)'}")
if REPO_FILTER:
print(f"📁 Repo filter: {REPO_FILTER}")
api_key = os.environ.get("ANTHROPIC_API_KEY")
if api_key:
print(f"🔑 Using API key from environment: {api_key[:20]}...")
else:
print("🔐 ANTHROPIC_API_KEY not set; using Claude CLI logged-in session.")
existing_ids: set = set()
output_path = Path(OUTPUT_FILE)
if output_path.exists():
with open(output_path) as f:
for line in f:
try:
data = json.loads(line)
existing_ids.add(data["instance_id"])
except Exception:
continue
print(f"✅ Found {len(existing_ids)} already completed tasks.")
print(f"📚 Loading dataset: {DATASET_NAME} ({DATASET_SPLIT})...")
ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
pending_tasks = [
t
for t in ds
if t["instance_id"] not in existing_ids
and (not BENCH_FILTER or t.get("source", "") == BENCH_FILTER)
and (
not REPO_FILTER
or REPO_FILTER in (t.get("repo", "") or "").lower()
or REPO_FILTER in (t.get("repo_url", "") or "").lower()
)
]
if not pending_tasks:
print("🎉 No pending tasks to run!")
return
selected_tasks = random.sample(pending_tasks, min(NUM_TASKS, len(pending_tasks)))
print(f"🚀 Randomly selected {len(selected_tasks)} tasks to process.")
if PREFETCH_REPOS:
prefetch_task_repositories(selected_tasks, Path(WORK_ROOT))
else:
print("⏭️ PREFETCH_REPOS=0; skipping prefetch step.")
for i, task in enumerate(selected_tasks):
instance_id = task["instance_id"]
repo_url = task["repo_url"]
print(f"\n{'-' * 60}")
print(f"📦 [{i + 1}/{len(selected_tasks)}] Running: {instance_id}")
print(f" repo: {repo_url} source: {task.get('source', '?')}")
# try:
# patch, elapsed, traj_data, usage = run_single_task(
# instance_id=instance_id,
# repo_url=repo_url,
# work_root=WORK_ROOT,
# mitm_script_path=str(MITM_SCRIPT),
# trace_dir=TRACE_DIR,
# model=MODEL,
# task=task,
# )
# result_entry = {
# "instance_id": instance_id,
# "model_patch": patch if patch else "",
# "model_name_or_path": "claude-code-cli",
# "elapsed_seconds": round(elapsed, 1),
# "traj_data": traj_data,
# "token_usage": usage,
# }
# with open(OUTPUT_FILE, "a", encoding="utf-8") as f:
# f.write(json.dumps(result_entry) + "\n")
# print(f"✅ Result saved for {instance_id}")
# success_count += 1
# except Exception as e:
# print(f"❌ Error processing {instance_id}: {e}")
# failure_count += 1
# finally:
# cleanup_residuals()
# print("💤 Cooldown...")
# time.sleep(20)
# print(
# f"\n✅ Finished {len(selected_tasks)} random tasks: "
# f"{success_count} succeeded, {failure_count} failed. "
# f"Results in {OUTPUT_FILE}"
# )
if __name__ == "__main__":
main()
benchmarks/contextbench/scripts/batch_run_selected.py
+208
View File
@@ -0,0 +1,208 @@
import json
import os
import subprocess
import time
from pathlib import Path
from auto_run import prefetch_task_repositories, run_single_task
from datasets import load_dataset
ROOT = Path(__file__).resolve().parents[1]
OUTPUT_FILE = os.environ.get("OUTPUT_FILE", "all_predictions_claude.jsonl")
WORK_ROOT = os.environ.get("WORK_ROOT", "contextbench_work_dir_claude")
MODEL = os.environ.get("MODEL", os.environ.get("CLAUDE_MODEL", "")).strip()
DATASET_NAME = os.environ.get("DATASET_NAME", "Contextbench/ContextBench")
DATASET_SPLIT = os.environ.get("DATASET_SPLIT", "train")
BENCH_FILTER = os.environ.get("BENCH_FILTER", "").strip() # e.g. "Verified", "Pro", "Poly", "Multi"
PREFETCH_REPOS = os.environ.get("PREFETCH_REPOS", "1").strip() != "0"
MITM_SCRIPT = ROOT / "mitmproxy_addons" / "trace_recorder.py"
TRACE_DIR = ROOT / "traces" / "raw"
# Instances to run. Set instance_ids here or pass via SELECTED_IDS env var (comma-separated).
SELECTED_IDS = [
# "SWE-Bench-Pro__python__maintenance__bugfix__19a1fba2",
# "SWE-Bench-Pro__python__maintenance__bugfix__2464eadb",
# "SWE-Bench-Pro__python__maintenance__bugfix__38dc8f4e",
# "SWE-Bench-Pro__javascript__maintenance__bugfix__2bfb5681",
# "SWE-Bench-Pro__python__maintenance__bugfix__71253eae",
# "SWE-Bench-Pro__javascript__maintenance__bugfix__93b583ae",
# "SWE-Bench-Pro__python__maintenance__bugfix__dcc84d4c",
# "SWE-Bench-Pro__python__maintenance__bugfix__462b957d",
# "SWE-Bench-Pro__python__maintenance__bugfix__9af74069",
# "SWE-Bench-Pro__python__maintenance__bugfix__7b688a35",
# "SWE-Bench-Pro__python__maintenance__bugfix__64fffdfa",
# "SWE-Bench-Pro__python__maintenance__bugfix__22a1484c",
# "SWE-Bench-Pro__go__maintenance__bugfix__1177cd53",
# "SWE-Bench-Pro__python__maintenance__bugfix__a4287775",
# "SWE-Bench-Pro__python__maintenance__bugfix__ba13492e",
# "SWE-Bench-Pro__go__maintenance__bugfix__b91d5788",
# "SWE-Bench-Pro__python__maintenance__bugfix__091dae2f",
# "SWE-Bench-Pro__python__maintenance__bugfix__b6eff698",
# "SWE-Bench-Pro__python__maintenance__bugfix__fcb506a5",
# "SWE-Bench-Pro__python__maintenance__bugfix__3cfd9a02",
# "SWE-Bench-Pro__python__maintenance__bugfix__4c132bfd",
# "SWE-Bench-Pro__python__maintenance__bugfix__7c2efe8a",
"SWE-Bench-Pro__go__maintenance__bugfix__40a717e5",
"SWE-Bench-Pro__go__maintenance__bugfix__52d866b3",
"SWE-Bench-Pro__go__maintenance__bugfix__720b4d92",
"SWE-Bench-Pro__go__maintenance__bugfix__997c7afd",
"SWE-Bench-Pro__javascript__maintenance__bugfix__82518720",
"SWE-Bench-Pro__javascript__maintenance__bugfix__e31ec45c",
"SWE-Bench-Pro__python__maintenance__bugfix__07bb383a",
"SWE-Bench-Pro__python__maintenance__bugfix__0bac5789",
"SWE-Bench-Pro__python__maintenance__bugfix__18d7bbbc",
"SWE-Bench-Pro__python__maintenance__bugfix__1cf3e889",
"SWE-Bench-Pro__python__maintenance__bugfix__20dad82b",
"SWE-Bench-Pro__python__maintenance__bugfix__20f502e0",
"SWE-Bench-Pro__python__maintenance__bugfix__509a20d9",
"SWE-Bench-Pro__python__maintenance__bugfix__53ca6a30",
"SWE-Bench-Pro__python__maintenance__bugfix__552343cd",
"SWE-Bench-Pro__python__maintenance__bugfix__5b2cf9bb",
"SWE-Bench-Pro__python__maintenance__bugfix__66e05eaa",
"SWE-Bench-Pro__python__maintenance__bugfix__6ebb54dc",
"SWE-Bench-Pro__python__maintenance__bugfix__87bfb374",
"SWE-Bench-Pro__python__maintenance__bugfix__89932d58",
"SWE-Bench-Pro__python__maintenance__bugfix__942d0b14",
"SWE-Bench-Pro__python__maintenance__bugfix__983f2896",
"SWE-Bench-Pro__python__maintenance__bugfix__a984b409",
"SWE-Bench-Pro__python__maintenance__bugfix__aa07d0c3",
"SWE-Bench-Pro__python__maintenance__bugfix__cf01f471",
"SWE-Bench-Pro__python__maintenance__bugfix__d2506f10",
"SWE-Bench-Pro__python__maintenance__bugfix__e579f2f0",
"SWE-Bench-Pro__python__maintenance__bugfix__eafb1f0b",
"SWE-Bench-Pro__python__maintenance__bugfix__ef8756b1",
"SWE-Bench-Pro__python__maintenance__bugfix__f87209f8",
"SWE-Bench-Pro__python__maintenance__bugfix__ff79bafd",
]
if os.environ.get("SELECTED_IDS"):
SELECTED_IDS = [x.strip() for x in os.environ["SELECTED_IDS"].split(",") if x.strip()]
def cleanup_residuals():
print("🧹 Cleaning up residual processes (Claude & Mitm)...")
try:
subprocess.run(["pkill", "-f", "claude"], stderr=subprocess.DEVNULL)
subprocess.run(["pkill", "-f", "mitmdump"], stderr=subprocess.DEVNULL)
time.sleep(2)
except Exception:
pass
def main():
if not SELECTED_IDS:
print(
"⚠️ Warning: SELECTED_IDS list is empty. Add instance IDs to the script or set SELECTED_IDS env var."
)
return
Path(WORK_ROOT).mkdir(parents=True, exist_ok=True)
TRACE_DIR.mkdir(parents=True, exist_ok=True)
print(f"🧠 Model: {MODEL or '(Claude CLI default)'}")
if BENCH_FILTER:
print(f"🔎 Bench filter: {BENCH_FILTER}")
api_key = os.environ.get("ANTHROPIC_API_KEY")
if api_key:
print(f"🔑 Using API key from environment: {api_key[:20]}...")
else:
print("🔐 ANTHROPIC_API_KEY not set; using Claude CLI logged-in session.")
# Load already-completed instance IDs to support resuming.
existing_ids: set = set()
output_path = Path(OUTPUT_FILE)
if output_path.exists():
with open(output_path) as f:
for line in f:
try:
data = json.loads(line)
existing_ids.add(data["instance_id"])
except Exception:
continue
print(f"✅ Found {len(existing_ids)} already completed tasks.")
print(f"📚 Loading dataset: {DATASET_NAME} ({DATASET_SPLIT})...")
ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
# Build a lookup dict for fast access.
task_lookup = {t["instance_id"]: t for t in ds}
selected_tasks = []
for iid in SELECTED_IDS:
if iid in existing_ids:
print(f"⏩ Skipping {iid} (already completed)")
continue
task = task_lookup.get(iid)
if task is None:
print(f"⚠️ Instance {iid} not found in dataset; skipping.")
continue
if BENCH_FILTER and task.get("source", "") != BENCH_FILTER:
print(f"⏩ Skipping {iid} (source={task.get('source')} != {BENCH_FILTER})")
continue
selected_tasks.append(task)
if not selected_tasks:
print("🎉 No pending selected tasks to run!")
return
print(f"🚀 Selected {len(selected_tasks)} tasks to process.")
if PREFETCH_REPOS:
prefetch_task_repositories(selected_tasks, Path(WORK_ROOT))
else:
print("⏭️ PREFETCH_REPOS=0; skipping prefetch step.")
success_count = 0
failure_count = 0
for i, task in enumerate(selected_tasks):
instance_id = task["instance_id"]
repo_url = task["repo_url"]
print(f"\n{'-' * 60}")
print(f"📦 [{i + 1}/{len(selected_tasks)}] Running: {instance_id}")
print(f" repo: {repo_url} source: {task.get('source', '?')}")
try:
patch, elapsed, agent_seconds, traj_data, usage = run_single_task(
instance_id=instance_id,
repo_url=repo_url,
work_root=WORK_ROOT,
mitm_script_path=str(MITM_SCRIPT),
trace_dir=TRACE_DIR,
model=MODEL,
task=task,
)
result_entry = {
"instance_id": instance_id,
"model_patch": patch if patch else "",
"model_name_or_path": "claude-code-cli",
"elapsed_seconds": round(elapsed, 1),
"latency_seconds": round(elapsed, 1),
"agent_seconds": round(agent_seconds, 1) if agent_seconds is not None else None,
"traj_data": traj_data,
"token_usage": usage,
}
with open(OUTPUT_FILE, "a", encoding="utf-8") as f:
f.write(json.dumps(result_entry) + "\n")
print(f"✅ Result saved for {instance_id}")
success_count += 1
except Exception as e:
print(f"❌ Error processing {instance_id}: {e}")
failure_count += 1
finally:
cleanup_residuals()
print("💤 Cooldown...")
time.sleep(20)
print(
f"\n✅ Finished {len(selected_tasks)} selected tasks: "
f"{success_count} succeeded, {failure_count} failed. "
f"Results in {OUTPUT_FILE}"
)
if __name__ == "__main__":
main()
benchmarks/contextbench/scripts/prepare_repos_with_leann.py
+462
View File
@@ -0,0 +1,462 @@
"""
Prepare ContextBench repositories and build LEANN indexes.
For each selected ContextBench instance:
1. Clone the repo into <WORK_ROOT>/<instance_id>
2. Checkout base_commit
3. Build LEANN index under <WORK_ROOT>/<instance_id>/.leann/
Usage:
cd scripts
python prepare_repos_with_leann.py
"""
import json
import os
import shlex
import shutil
import subprocess
import time
from pathlib import Path
from typing import Optional
from datasets import load_dataset
from git import Repo
WORK_ROOT = os.environ.get("WORK_ROOT", "contextbench_work_dir_claude")
LEANN_BIN = os.environ.get("LEANN_BIN", "leann")
DATASET_NAME = os.environ.get("DATASET_NAME", "Contextbench/ContextBench")
DATASET_SPLIT = os.environ.get("DATASET_SPLIT", "train")
BENCH_FILTER = os.environ.get("BENCH_FILTER", "").strip() # Verified | Pro | Poly | Multi
LEANN_SOURCE_EXTENSIONS = os.environ.get(
"LEANN_SOURCE_EXTENSIONS",
"py,go,js,jsx,ts,tsx,java,kt,kts,rs,rb,php,cs,c,cc,cpp,h,hpp,m,mm,swift,scala,sh,sql,lua,r",
)
# LEANN index build parameters — override via env vars for sweeping configs.
# ast-chunk-size is in non-whitespace CHARACTERS (not tokens). bge-base-en-v1.5
# has a 512-token limit; at ~1.2 tokens/char: 300 chars + 64 overlap ≈ 436 tokens.
LEANN_EMBEDDING_MODEL = os.environ.get(
"LEANN_EMBEDDING_MODEL", "jinaai/jina-embeddings-v2-base-code"
)
LEANN_AST_CHUNK_SIZE = os.environ.get("LEANN_AST_CHUNK_SIZE", "600")
LEANN_AST_CHUNK_OVERLAP = os.environ.get("LEANN_AST_CHUNK_OVERLAP", "96")
# Set to "0" to disable vendor/generated exclusion (e.g. for ablation experiments).
LEANN_EXCLUDE_VENDOR = os.environ.get("LEANN_EXCLUDE_VENDOR", "1").strip() != "0"
# Set to "1" to exclude test files (e.g. for ablation experiments). Default off
# to avoid missing bugfix targets that touch test/fixture/spec files.
LEANN_EXCLUDE_TESTS = os.environ.get("LEANN_EXCLUDE_TESTS", "0").strip() != "0"
FAILED_INSTANCES_LOG = os.environ.get(
"FAILED_INSTANCES_LOG", "prepare_repos_with_leann_failures.jsonl"
).strip()
# Fill in ContextBench instance_ids. Leave empty to prepare all tasks
# (optionally filtered by BENCH_FILTER).
SELECTED_IDS: list[str] = [
# "SWE-Bench-Pro__python__maintenance__bugfix__19a1fba2",
# "SWE-Bench-Pro__python__maintenance__bugfix__2464eadb",
# "SWE-Bench-Pro__python__maintenance__bugfix__38dc8f4e",
# "SWE-Bench-Pro__javascript__maintenance__bugfix__2bfb5681",
# "SWE-Bench-Pro__python__maintenance__bugfix__71253eae",
# "SWE-Bench-Pro__javascript__maintenance__bugfix__93b583ae",
# "SWE-Bench-Pro__python__maintenance__bugfix__dcc84d4c",
# "SWE-Bench-Pro__python__maintenance__bugfix__462b957d",
# "SWE-Bench-Pro__python__maintenance__bugfix__9af74069",
# "SWE-Bench-Pro__python__maintenance__bugfix__7b688a35",
# "SWE-Bench-Pro__python__maintenance__bugfix__64fffdfa",
# "SWE-Bench-Pro__python__maintenance__bugfix__22a1484c",
# "SWE-Bench-Pro__go__maintenance__bugfix__1177cd53",
# "SWE-Bench-Pro__python__maintenance__bugfix__a4287775",
# "SWE-Bench-Pro__python__maintenance__bugfix__ba13492e",
# "SWE-Bench-Pro__go__maintenance__bugfix__b91d5788",
# "SWE-Bench-Pro__python__maintenance__bugfix__091dae2f",
# "SWE-Bench-Pro__python__maintenance__bugfix__b6eff698",
# "SWE-Bench-Pro__python__maintenance__bugfix__fcb506a5",
# "SWE-Bench-Pro__python__maintenance__bugfix__3cfd9a02",
# "SWE-Bench-Pro__python__maintenance__bugfix__4c132bfd",
# "SWE-Bench-Pro__python__maintenance__bugfix__7c2efe8a",
"SWE-Bench-Pro__go__maintenance__bugfix__40a717e5",
"SWE-Bench-Pro__go__maintenance__bugfix__52d866b3",
"SWE-Bench-Pro__go__maintenance__bugfix__720b4d92",
"SWE-Bench-Pro__go__maintenance__bugfix__997c7afd",
"SWE-Bench-Pro__javascript__maintenance__bugfix__82518720",
"SWE-Bench-Pro__javascript__maintenance__bugfix__e31ec45c",
"SWE-Bench-Pro__python__maintenance__bugfix__07bb383a",
"SWE-Bench-Pro__python__maintenance__bugfix__0bac5789",
"SWE-Bench-Pro__python__maintenance__bugfix__18d7bbbc",
"SWE-Bench-Pro__python__maintenance__bugfix__1cf3e889",
"SWE-Bench-Pro__python__maintenance__bugfix__20dad82b",
"SWE-Bench-Pro__python__maintenance__bugfix__20f502e0",
"SWE-Bench-Pro__python__maintenance__bugfix__509a20d9",
"SWE-Bench-Pro__python__maintenance__bugfix__53ca6a30",
"SWE-Bench-Pro__python__maintenance__bugfix__552343cd",
"SWE-Bench-Pro__python__maintenance__bugfix__5b2cf9bb",
"SWE-Bench-Pro__python__maintenance__bugfix__66e05eaa",
"SWE-Bench-Pro__python__maintenance__bugfix__6ebb54dc",
"SWE-Bench-Pro__python__maintenance__bugfix__87bfb374",
"SWE-Bench-Pro__python__maintenance__bugfix__89932d58",
"SWE-Bench-Pro__python__maintenance__bugfix__942d0b14",
"SWE-Bench-Pro__python__maintenance__bugfix__983f2896",
"SWE-Bench-Pro__python__maintenance__bugfix__a984b409",
"SWE-Bench-Pro__python__maintenance__bugfix__aa07d0c3",
"SWE-Bench-Pro__python__maintenance__bugfix__cf01f471",
"SWE-Bench-Pro__python__maintenance__bugfix__d2506f10",
"SWE-Bench-Pro__python__maintenance__bugfix__e579f2f0",
"SWE-Bench-Pro__python__maintenance__bugfix__eafb1f0b",
"SWE-Bench-Pro__python__maintenance__bugfix__ef8756b1",
"SWE-Bench-Pro__python__maintenance__bugfix__f87209f8",
"SWE-Bench-Pro__python__maintenance__bugfix__ff79bafd",
]
if os.environ.get("SELECTED_IDS"):
SELECTED_IDS = [x.strip() for x in os.environ["SELECTED_IDS"].split(",") if x.strip()]
# Vendor and generated-code directory/file patterns to exclude from the index.
# These are third-party or machine-generated files that are never the target of
# a bugfix, so indexing them only adds noise to search results.
_VENDOR_DIR_PATTERNS = (
"vendor/",
"node_modules/",
"third_party/",
"thirdparty/",
"externals/",
".cache/",
)
_GENERATED_FILE_PATTERNS = (
"_pb.go",
".pb.go",
"_gen.go",
".pb.cc",
".pb.h",
)
# Filenames like `zz_generated.deepcopy.go` end in `.go`, not `zz_generated`;
# match these as path substrings (controller-gen / k8s-style outputs).
_GENERATED_FILE_SUBSTRINGS = ("zz_generated",)
# Test file path/name patterns to exclude from the index.
_TEST_PATH_PATTERNS = (
"/test/",
"/tests/",
"/__tests__/",
"/spec/",
"/testdata/",
"/test_",
"/fixtures/",
)
_TEST_FILE_PATTERNS = (
"_test.py",
"_test.go",
".test.js",
".test.ts",
".test.jsx",
".test.tsx",
".spec.js",
".spec.ts",
".spec.jsx",
".spec.tsx",
"_spec.rb",
)
def _run_command(cmd: list[str], cwd: Path) -> subprocess.CompletedProcess:
return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
def _read_json_file(path: Path) -> Optional[dict]:
if not path.exists():
return None
try:
return json.loads(path.read_text(encoding="utf-8"))
except Exception:
return None
def _count_index_chunks(repo_dir: Path, instance_id: str) -> Optional[int]:
ids_path = repo_dir / ".leann" / "indexes" / instance_id / "documents.ids.txt"
if not ids_path.exists():
return None
with ids_path.open("r", encoding="utf-8") as f:
return sum(1 for _ in f)
def _print_subprocess_output(label: str, text: str, max_lines: int = 20) -> None:
lines = [line for line in (text or "").splitlines() if line.strip()]
if not lines:
return
print(f" 📄 {label}:")
for line in lines[:max_lines]:
print(f" {line}")
if len(lines) > max_lines:
print(f" ... ({len(lines) - max_lines} more lines)")
def _write_failure_report(failures: list[dict]) -> Optional[Path]:
if not failures:
return None
report_path = Path(FAILED_INSTANCES_LOG)
report_path.parent.mkdir(parents=True, exist_ok=True)
with report_path.open("w", encoding="utf-8") as f:
for item in failures:
f.write(json.dumps(item, ensure_ascii=True) + "\n")
return report_path
def _load_tasks() -> list[dict]:
print(f"📚 Loading dataset: {DATASET_NAME} ({DATASET_SPLIT})...")
ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
tasks: list[dict] = list(ds)
if BENCH_FILTER:
tasks = [t for t in tasks if t.get("source", "") == BENCH_FILTER]
if SELECTED_IDS:
task_lookup = {t["instance_id"]: t for t in tasks}
selected: list[dict] = []
for iid in SELECTED_IDS:
task = task_lookup.get(iid)
if not task:
print(f"⚠️ Instance not found in dataset/split/filter: {iid}")
continue
selected.append(task)
return selected
return tasks
def _is_pytest_style_test_py(normalized_path: str) -> bool:
"""True for pytest-style modules: basename test_*.py (e.g. test_foo.py)."""
name = Path(normalized_path).name
return name.startswith("test_") and name.lower().endswith(".py")
def build_leann_index(instance_id: str, repo_dir: Path) -> tuple[bool, Optional[str]]:
print(f" 🔍 Building LEANN index for {instance_id}...")
result = _run_command(["git", "ls-files"], cwd=repo_dir)
if result.returncode != 0:
error = f"Could not list git files: {result.stderr.strip()}"
print(f" ⚠️ {error}")
return False, error
tracked_files = [f for f in result.stdout.strip().split("\n") if f]
allowed_exts = {
f".{ext.strip().lstrip('.').lower()}"
for ext in LEANN_SOURCE_EXTENSIONS.split(",")
if ext.strip()
}
source_files = [f for f in tracked_files if Path(f).suffix.lower() in allowed_exts]
# Exclude vendor/generated files — they are never bugfix targets and add noise.
if LEANN_EXCLUDE_VENDOR:
before = len(source_files)
normalized = [f.replace("\\", "/") for f in source_files]
source_files = [
f
for f, n in zip(source_files, normalized)
if not any(pat in n for pat in _VENDOR_DIR_PATTERNS)
and not any(n.endswith(pat) for pat in _GENERATED_FILE_PATTERNS)
and not any(sub in n for sub in _GENERATED_FILE_SUBSTRINGS)
]
excluded = before - len(source_files)
if excluded:
print(
f" 🚫 Excluded {excluded} vendor/generated files ({before} → {len(source_files)})"
)
# Exclude test files — they are rarely bugfix targets and consistently rank
# high in semantic search due to mirroring production code patterns.
if LEANN_EXCLUDE_TESTS:
before = len(source_files)
normalized = [f.replace("\\", "/") for f in source_files]
source_files = [
f
for f, n in zip(source_files, normalized)
if not any(pat in n for pat in _TEST_PATH_PATTERNS)
and not _is_pytest_style_test_py(n)
and not any(n.endswith(pat) for pat in _TEST_FILE_PATTERNS)
]
excluded = before - len(source_files)
if excluded:
print(f" 🚫 Excluded {excluded} test files ({before} → {len(source_files)})")
if not source_files:
error = f"No source files found for extensions: {sorted(allowed_exts)}"
print(f" ⚠️ {error}")
return False, error
# Derive --file-types from the actual extensions present after filtering,
# so all indexed file types benefit from AST-aware chunking.
indexed_exts = sorted(
{Path(f).suffix.lstrip(".").lower() for f in source_files if Path(f).suffix}
)
file_types_arg = ",".join(indexed_exts)
print(f" 📊 Found {len(source_files)} source files (types: {file_types_arg})")
leann_cmd = [
LEANN_BIN,
"build",
instance_id,
"--docs",
*source_files,
"--embedding-mode",
"sentence-transformers",
"--embedding-model",
LEANN_EMBEDDING_MODEL,
"--backend",
"hnsw",
"--file-types",
file_types_arg,
"--force",
"--ast-chunk-size",
LEANN_AST_CHUNK_SIZE,
"--ast-chunk-overlap",
LEANN_AST_CHUNK_OVERLAP,
"--use-ast-chunking",
"--no-recompute",
]
debug_cmd = [
LEANN_BIN,
"build",
instance_id,
"--docs",
f"<{len(source_files)} files>",
*leann_cmd[len(["leann", "build", instance_id, "--docs"]) + len(source_files) :],
]
print(f" 🧪 LEANN command: {shlex.join(debug_cmd)}")
try:
proc = subprocess.run(
leann_cmd,
cwd=repo_dir,
capture_output=True,
text=True,
timeout=1800,
env={
**os.environ,
"LEANN_EMBEDDING_DEVICE": os.environ.get("LEANN_EMBEDDING_DEVICE", "mps"),
"LEANN_BATCH_SIZE": os.environ.get("LEANN_BATCH_SIZE", "32"),
},
)
except subprocess.TimeoutExpired:
error = "LEANN build timed out"
print(f" ⏰ {error}")
return False, error
except Exception as e:
error = f"LEANN error: {e}"
print(f" ❌ {error}")
return False, error
_print_subprocess_output("LEANN stdout", proc.stdout)
_print_subprocess_output("LEANN stderr", proc.stderr)
if proc.returncode != 0:
stderr = (proc.stderr or "").strip()
print(f" ❌ LEANN build failed: {stderr}")
return False, f"LEANN build failed: {stderr}"
meta_path = repo_dir / ".leann" / "indexes" / instance_id / "documents.leann.meta.json"
meta = _read_json_file(meta_path)
if meta:
print(f" 🧠 Embedding model in index: {meta.get('embedding_model', 'unknown')}")
print(" 🌲 AST chunking requested: yes (--use-ast-chunking)")
chunk_count = _count_index_chunks(repo_dir, instance_id)
if chunk_count is not None:
print(f" 🧩 Indexed chunks: {chunk_count}")
print(" ✅ LEANN index built successfully")
return True, None
def prepare_single_task(task: dict) -> tuple[bool, Optional[str]]:
instance_id = task["instance_id"]
repo_url = task["repo_url"]
base_commit = task["base_commit"]
target_dir = Path(WORK_ROOT) / instance_id
print(f"\n{'=' * 72}")
print(f"📦 Preparing: {instance_id}")
print(f" repo: {repo_url}")
print(f" commit: {base_commit[:12]}...")
print(f"{'=' * 72}")
if not target_dir.exists():
print(f" 📥 Cloning {repo_url}...")
try:
Repo.clone_from(repo_url, target_dir)
except Exception as e:
print(f" ❌ Clone failed: {e}")
return False, f"Clone failed: {e}"
else:
print(" ✓ Repo already exists")
try:
repo = Repo(target_dir)
print(f" 🔀 Checking out {base_commit[:8]}...")
repo.git.reset("--hard")
repo.git.checkout(base_commit)
repo.git.clean("-fdx", "-e", ".leann/")
(target_dir / "PROBLEM.md").write_text(task.get("problem_statement", ""), encoding="utf-8")
except Exception as e:
print(f" ❌ Checkout/clean failed: {e}")
return False, f"Checkout/clean failed: {e}"
return build_leann_index(instance_id, target_dir)
def main():
print("🚀 ContextBench Repository Preparation with LEANN Indexing")
print("=" * 72)
leann_path = shutil.which(LEANN_BIN)
if not leann_path:
print("❌ LEANN not found. Install: uv tool install leann-core --with leann")
return
print(f"✅ LEANN found: {leann_path}")
tasks = _load_tasks()
if not tasks:
print("⚠️ No tasks selected. Set SELECTED_IDS or adjust BENCH_FILTER.")
return
Path(WORK_ROOT).mkdir(parents=True, exist_ok=True)
print(f"\n📂 Work root: {WORK_ROOT}")
print(f"🎯 Tasks to prepare: {len(tasks)}")
if BENCH_FILTER:
print(f"🔎 Bench filter: {BENCH_FILTER}")
success_count = 0
fail_count = 0
failures: list[dict] = []
for i, task in enumerate(tasks, start=1):
print(f"\n[{i}/{len(tasks)}]")
succeeded, error = prepare_single_task(task)
if succeeded:
success_count += 1
else:
fail_count += 1
failures.append(
{
"instance_id": task.get("instance_id", ""),
"repo_url": task.get("repo_url", ""),
"base_commit": task.get("base_commit", ""),
"error": error or "unknown error",
"failed_at_unix": int(time.time()),
}
)
print(f"\n{'=' * 72}")
print(f"🎉 Done! ✅ {success_count} succeeded ❌ {fail_count} failed")
if failures:
print("\nFailed instances:")
for item in failures:
print(f" - {item['instance_id']}: {item['error']}")
report_path = _write_failure_report(failures)
if report_path is not None:
print(f"\n📝 Failure report written to: {report_path}")
print("\nNext steps:")
print(f" 1. Verify indexes: ls {WORK_ROOT}/*/.leann")
print(" 2. Run with LEANN: LEANN_ENABLED=1 python batch_run_selected.py")
print(" 3. Baseline run: LEANN_ENABLED=0 python batch_run_selected.py")
if __name__ == "__main__":
main()