|
33 | 33 |
|
34 | 34 | sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "collection")) |
35 | 35 | import local_diffs as ld # noqa: E402 |
| 36 | +import llm_classify_fixes as llm # noqa: E402 (reuse the pluggable LLM engine) |
36 | 37 |
|
37 | 38 | CONSENSUS = {"lighthouse", "lodestar", "nimbus", "prysm", "teku", "grandine"} |
38 | 39 | PR_RE = re.compile(r"/pull/(\d+)") |
@@ -201,19 +202,81 @@ def _group(hunks): |
201 | 202 | return out |
202 | 203 |
|
203 | 204 |
|
| 205 | +# --- LLM fallback for rows the deterministic rules leave as "other" ---------- |
| 206 | +CONSENSUS_LABELS = [ |
| 207 | + "beacon-chain:justification-and-finality", "beacon-chain:rewards-and-penalties", |
| 208 | + "beacon-chain:registry-updates", "beacon-chain:effective-balance-updates", |
| 209 | + "beacon-chain:epoch-processing", "beacon-chain:block-processing", |
| 210 | + "beacon-chain:attestation", "beacon-chain:slashing", "beacon-chain:deposit", |
| 211 | + "beacon-chain:withdrawal", "beacon-chain:exit-consolidation", |
| 212 | + "beacon-chain:sync-committee", "beacon-chain:execution-payload", "fork-choice", |
| 213 | + "p2p-interface", "validator", "weak-subjectivity", "deposit-contract", "bls", |
| 214 | + "light-client", "fork-transition", "kzg-commitments", |
| 215 | + "data-availability-sampling", "builder"] |
| 216 | +EXECUTION_LABELS = [ |
| 217 | + "evm", "opcodes", "precompiles", "gas", "transactions", "txpool", |
| 218 | + "block-processing", "state-trie", "rlp", "p2p", "sync", "engine-api", |
| 219 | + "blobs", "eof", "rpc"] |
| 220 | +CROSS = ["crypto", "serialization", "database", "other"] |
| 221 | +RC_ENUM = [v for _, v in _RC] + ["improper_state_update", "other"] |
| 222 | +AP_ENUM = [v for _, v in _AP] + ["internal_only"] |
| 223 | + |
| 224 | + |
| 225 | +def llm_label(row, diff, lyr) -> dict: |
| 226 | + labels = (CONSENSUS_LABELS if lyr == "consensus" else EXECUTION_LABELS) + CROSS |
| 227 | + prompt = f"""Label this security fix in an Ethereum {lyr} client. |
| 228 | +
|
| 229 | +Pick the ONE best AREA label from this list (use "other" only if truly none fit): |
| 230 | +{', '.join(labels)} |
| 231 | +
|
| 232 | +Also pick root_cause from: {', '.join(sorted(set(RC_ENUM)))} |
| 233 | +and attack_path from: {', '.join(sorted(set(AP_ENUM)))} |
| 234 | +
|
| 235 | +Changed files: {row.get('files') or '(none)'} |
| 236 | +Title: {str(row.get('title') or '')[:200]} |
| 237 | +Description: {str(row.get('description') or '')[:400]} |
| 238 | +Code diff (truncated): |
| 239 | +{(diff or '')[:3000]} |
| 240 | +
|
| 241 | +Output ONLY one JSON object on the last line: |
| 242 | +{{"label": "...", "root_cause": "...", "attack_path": "..."}}""" |
| 243 | + try: |
| 244 | + out = llm._call_llm(prompt) |
| 245 | + m = re.search(r"\{[^{}]*\"label\"[^{}]*\}", out, re.S) |
| 246 | + obj = json.loads(m.group(0)) if m else {} |
| 247 | + except Exception: |
| 248 | + obj = {} |
| 249 | + valid = set(labels) |
| 250 | + lab = obj.get("label") if obj.get("label") in valid else None |
| 251 | + return {"label": lab, "root_cause": obj.get("root_cause"), |
| 252 | + "attack_path": obj.get("attack_path")} |
| 253 | + |
| 254 | + |
204 | 255 | def main() -> int: |
205 | 256 | ap = argparse.ArgumentParser() |
206 | 257 | ap.add_argument("--in", dest="inp", default=Path("data/ethereum_vulns.parquet"), type=Path) |
207 | 258 | ap.add_argument("--out", default=Path("data/labels.csv"), type=Path) |
208 | 259 | ap.add_argument("--pred-cache", default=Path("scratchpad_crawl/llm_pred_cache.json"), type=Path) |
209 | 260 | ap.add_argument("--diff-cache", default=Path("scratchpad_crawl/diff_cache.json"), type=Path) |
| 261 | + ap.add_argument("--llm", action="store_true", help="LLM fallback for 'other' rows") |
| 262 | + ap.add_argument("--llm-cache", default=Path("scratchpad_crawl/llm_label_cache.json"), type=Path) |
| 263 | + ap.add_argument("--engine", default="openai") |
| 264 | + ap.add_argument("--model", default="") |
| 265 | + ap.add_argument("--base-url", default="https://ollama.com/v1") |
| 266 | + ap.add_argument("--api-key-env", default="OLLAMA_API_KEY") |
| 267 | + ap.add_argument("--workers", type=int, default=6) |
210 | 268 | a = ap.parse_args() |
| 269 | + import os |
| 270 | + llm.ENGINE.update(engine=a.engine, |
| 271 | + model=a.model or ("gemma4:31b" if a.engine == "openai" else ""), |
| 272 | + base_url=a.base_url, |
| 273 | + api_key=os.environ.get(a.api_key_env, "") if a.api_key_env else "") |
211 | 274 |
|
212 | 275 | df = pd.read_parquet(a.inp) |
213 | 276 | preds = json.loads(a.pred_cache.read_text()) if a.pred_cache.exists() else {} |
214 | 277 | dcache = json.loads(a.diff_cache.read_text()) if a.diff_cache.exists() else {} |
215 | 278 |
|
216 | | - rows, n_diff, n_label = [], 0, 0 |
| 279 | + rows, metas, n_diff, n_label = [], [], 0, 0 |
217 | 280 | for i, r in enumerate(df.to_dict("records")): |
218 | 281 | client = r["source_platform"]; url = str(r["source_url"]); lyr = layer(client) |
219 | 282 | repo = ld.CLIENT_REPOS.get(client) |
@@ -251,11 +314,46 @@ def main() -> int: |
251 | 314 | "post_fix_code": json.dumps(post, ensure_ascii=False), |
252 | 315 | "fix_commit": fix_sha, "introduced_in_commit": introduced, |
253 | 316 | }) |
| 317 | + metas.append({"url": url, "layer": lyr, "files": ", ".join(files[:6]), |
| 318 | + "title": r.get("title"), "description": r.get("description")}) |
254 | 319 | if (i + 1) % 200 == 0: |
255 | 320 | a.diff_cache.write_text(json.dumps(dcache)) |
256 | 321 | print(f" [labels] {i+1}/{len(df)}", file=sys.stderr) |
257 | 322 | a.diff_cache.write_text(json.dumps(dcache)) |
258 | 323 |
|
| 324 | + # --- LLM fallback for rows still "other" ------------------------------- |
| 325 | + if a.llm: |
| 326 | + from concurrent.futures import ThreadPoolExecutor |
| 327 | + cache = json.loads(a.llm_cache.read_text()) if a.llm_cache.exists() else {} |
| 328 | + todo = [i for i, row in enumerate(rows) if row["label"] == "other"] |
| 329 | + print(f"[labels] LLM fallback on {len(todo)} 'other' rows " |
| 330 | + f"({sum(1 for i in todo if rows[i]['id'] in cache)} cached)", file=sys.stderr) |
| 331 | + |
| 332 | + def work(i): |
| 333 | + rid = rows[i]["id"] |
| 334 | + if rid in cache: |
| 335 | + return i, cache[rid] |
| 336 | + diff = dcache.get(metas[i]["url"]) or "" |
| 337 | + res = llm_label(metas[i], diff, metas[i]["layer"]) |
| 338 | + return i, res |
| 339 | + |
| 340 | + done = 0 |
| 341 | + with ThreadPoolExecutor(max_workers=a.workers) as ex: |
| 342 | + for i, res in ex.map(work, todo): |
| 343 | + cache[rows[i]["id"]] = res |
| 344 | + if res.get("label"): |
| 345 | + rows[i]["label"] = res["label"] |
| 346 | + if res.get("root_cause"): |
| 347 | + rows[i]["root_cause"] = res["root_cause"] |
| 348 | + if res.get("attack_path"): |
| 349 | + rows[i]["attack_path"] = res["attack_path"] |
| 350 | + done += 1 |
| 351 | + if done % 50 == 0: |
| 352 | + a.llm_cache.write_text(json.dumps(cache)) |
| 353 | + print(f" [labels-llm] {done}/{len(todo)}", file=sys.stderr) |
| 354 | + a.llm_cache.write_text(json.dumps(cache)) |
| 355 | + n_label = sum(1 for r in rows if r["label"] != "other") |
| 356 | + |
259 | 357 | a.out.parent.mkdir(parents=True, exist_ok=True) |
260 | 358 | with a.out.open("w", newline="", encoding="utf-8") as fh: |
261 | 359 | w = csv.DictWriter(fh, fieldnames=list(rows[0].keys())) |
|
0 commit comments