File size: 3,016 Bytes
21a9138 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Build the model's input state exactly as in training and evaluation.
python make_state.py diff CHANGE.diff --task-file TASK.txt [--baseline-file BASELINE.txt]
python make_state.py snippet CODE.py --lang py --task "Clean up the build cache"
Prints the state. Use it as `state=` in `model.system_one(...)` (see ../examples/usage.py).
"""
import argparse
import sys
from collections import Counter, defaultdict
from pathlib import Path
from behaviour import analyse_diff, analyse_snippet
BUDGET_CHARS = 8192 * 3 # ≈ 8k tokens; longer changes are cut with a marker, as in training
def summary(rep):
t = (rep or {}).get("text") or ""
return t.strip() or "(no security-relevant behaviour found)"
def baseline_text(counter_by_kind):
"""Repository baseline: the most common operands per behaviour kind across the repo's ordinary changes."""
if not counter_by_kind:
return "(no repository history available)"
lines = []
for kind in sorted(counter_by_kind):
top = counter_by_kind[kind].most_common(8)
lines.append(f"- {kind}: " + "; ".join(f"{op} ({n})" for op, n in top))
return "\n".join(lines)
def add_to_baseline(acc, rep):
"""Accumulate one ordinary change's report into acc (a defaultdict(Counter)) for baseline_text."""
for b in (rep or {}).get("behaviours", []):
acc[b["kind"]][str(b.get("operand") or b.get("path") or "?")[:80]] += 1
def new_baseline():
return defaultdict(Counter)
def state(task, baseline, beh, change):
parts = [f"Task:\n{(task or '(not provided)').strip()}",
f"Repository baseline (normal behaviour of this repo/org):\n{baseline}",
f"Behaviour summary (automated, not a verdict):\n{beh}"]
head = "\n\n".join(parts) + "\n\nChange:\n"
room = BUDGET_CHARS - len(head)
if len(change) > room:
change = change[:max(room - 80, 0)] + f"\n[... {len(change) - room + 80} more characters omitted]"
return head + change
def main(argv=None):
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
ap.add_argument("mode", choices=("diff", "snippet"))
ap.add_argument("input", help="file, or - for stdin")
ap.add_argument("--task"); ap.add_argument("--task-file")
ap.add_argument("--baseline-file", help="precomputed baseline text (default: no repository history)")
ap.add_argument("--lang", choices=("js", "ts", "py", "sh"), help="snippet language")
a = ap.parse_args(argv)
change = sys.stdin.read() if a.input == "-" else Path(a.input).read_text(encoding="utf-8", errors="replace")
task = a.task or (Path(a.task_file).read_text(errors="replace") if a.task_file else None)
rep = analyse_diff(change, task=task) if a.mode == "diff" else analyse_snippet(change, lang=a.lang, task=task)
base = Path(a.baseline_file).read_text().strip() if a.baseline_file else baseline_text(None)
print(state(task, base, summary(rep), change))
return 0
if __name__ == "__main__":
sys.exit(main())
|