| """Plumbing for the one-GRPO-step walkthrough. |
| |
| Nothing in this file is a concept the notebook teaches -- it is the prompt |
| text, the sampling boilerplate and the fake environment, kept here so the |
| notebook itself can stay about the harness, the reward and the update. |
| |
| Every method of MockEnv is a lie. The comments say which kind. |
| """ |
|
|
| import difflib |
| import re |
|
|
| import pandas as pd |
| import torch |
| from IPython.display import display |
|
|
| TARGET = "astropy/modeling/separable.py" |
|
|
| |
| |
| MOCK_FILE = '''import numpy as np |
| |
| class Model: |
| """Mock model class.""" |
| pass |
| |
| def _compute_n_outputs(left, right): |
| return left.shape[0] if hasattr(left, "shape") else 1 |
| |
| def _coord_matrix(model, side, noutp): |
| return np.zeros((noutp, 1)) |
| |
| def _cstack(left, right): |
| """ |
| Function corresponding to '&' operation. |
| |
| Parameters |
| ---------- |
| left, right : `astropy.modeling.Model` or ndarray |
| If input is of an array, it is the output of `coord_matrix`. |
| |
| Returns |
| ------- |
| result : ndarray |
| Result from this operation. |
| |
| """ |
| noutp = _compute_n_outputs(left, right) |
| |
| if isinstance(left, Model): |
| cleft = _coord_matrix(left, 'left', noutp) |
| else: |
| cleft = np.zeros((noutp, left.shape[1])) |
| cleft[: left.shape[0], : left.shape[1]] = left |
| if isinstance(right, Model): |
| cright = _coord_matrix(right, 'right', noutp) |
| else: |
| cright = np.zeros((noutp, right.shape[1])) |
| cright[-right.shape[0]:, -right.shape[1]:] = 1 |
| |
| return np.hstack([cleft, cright])''' |
|
|
| RE_BASH = re.compile(r"```(?:bash|sh)?\s*\n(.*?)```", re.S) |
| RE_HEREDOC = re.compile(r"cat\s*>\s*(\S+)\s*<<\s*'?EOF'?\n(.*?)\nEOF", re.S) |
|
|
| SYSTEM = f"""You are a software engineering agent. You are in a Python repository. |
| Fix the bug described in the issue. |
| |
| Respond with exactly ONE bash command per message, in a fenced block: |
| |
| ```bash |
| your command here |
| ``` |
| |
| Useful commands: |
| ls |
| cat {TARGET} |
| grep -n "pattern" {TARGET} |
| python -m pytest |
| |
| To rewrite a file: |
| ```bash |
| cat > {TARGET} <<'EOF' |
| ...full new contents... |
| EOF |
| ``` |
| |
| One short sentence of reasoning, then exactly one command block.""" |
|
|
| NO_COMMAND = "No command found. Reply with exactly one ```bash block." |
|
|
|
|
| def first_bash_block(reply): |
| """The parse step: pull one action out of free text, or None.""" |
| m = RE_BASH.search(reply) |
| return m.group(1).strip() if m else None |
|
|
|
|
| def show_rollout(rollout, preview=160): |
| """Trajectory table, commands issued, and the candidate patch. |
| |
| `preview` caps the preview column; pass None to show messages in full. |
| """ |
| def author(i, m): |
| if m["role"] == "assistant": |
| return "model" |
| return "harness" if i < 2 else "environment" |
|
|
| def shorten(text): |
| one_line = text.replace("\n", " | ") |
| if preview is None or len(one_line) <= preview: |
| return one_line |
| return one_line[:preview] + " […]" |
|
|
| msgs = rollout["messages"] |
| n_turns = (len(msgs) - 2) // 2 |
| print(f"{len(msgs)} messages = 2 setup + 2 per turn x {n_turns} turns\n") |
|
|
| display(pd.DataFrame([ |
| {"i": i, "turn": "setup" if i < 2 else (i - 2) // 2, |
| "role": m["role"], "written by": author(i, m), |
| "chars": len(m["content"]), |
| "preview": shorten(m["content"])} |
| for i, m in enumerate(msgs) |
| ])) |
|
|
| print("\ncommands:") |
| for c in rollout["calls"]: |
| print(" $", c.splitlines()[0] + (" ..." if "\n" in c else "")) |
|
|
| print("\ncandidate patch:") |
| print(rollout["patch"] or "(no file changed)") |
|
|
|
|
| def show_group(group, reward): |
| """Summary table for a group of rollouts, then every trace verbatim.""" |
| display(pd.DataFrame([ |
| {"rollout": i, |
| "reward": reward[i], |
| "patch_lines": len(r["patch"].split("\n")) if r["patch"] else 0, |
| "n_commands": len(r["calls"])} |
| for i, r in enumerate(group) |
| ])) |
|
|
| for i, r in enumerate(group): |
| print(f"\n================ rollout {i} ================") |
| for cmd in r["calls"]: |
| print(" $", cmd) |
| print("\ncandidate patch:") |
| print(r["patch"] or "(no file changed)") |
|
|
|
|
| def scripted_sampler(replies): |
| """A stand-in for `generate`: replays canned replies, ignoring the prompt. |
| |
| Lets the same harness be driven by a competent 'model' so a demo can show |
| what a successful rollout looks like. |
| """ |
| it = iter(replies) |
|
|
| def sample(model, tok, prompt, **kwargs): |
| return next(it, "```bash\nls\n```") |
|
|
| return sample |
|
|
|
|
| def generate(model, tok, prompt, max_new_tokens=200, temperature=1.0, max_length=3072): |
| """The call step: sample a continuation, return only the new text.""" |
| device = next(model.parameters()).device |
| ip = tok(prompt, return_tensors="pt", truncation=True, max_length=max_length).to(device) |
| with torch.no_grad(): |
| out = model.generate(**ip, max_new_tokens=max_new_tokens, |
| do_sample=temperature > 0, |
| temperature=max(temperature, 1e-5), top_p=0.95, |
| pad_token_id=tok.pad_token_id or tok.eos_token_id) |
| return tok.decode(out[0][ip["input_ids"].shape[1]:], skip_special_tokens=True).strip() |
|
|
|
|
| class MockEnv: |
| """A shell that never runs a shell. |
| |
| Stands in for the container a real harness would exec into. Handles ls, |
| cat, grep, pytest and a heredoc write; anything else returns a stub. |
| """ |
|
|
| def __init__(self, failing_tests=()): |
| self.fs = {TARGET: MOCK_FILE} |
| self.orig = dict(self.fs) |
| self.failing_tests = list(failing_tests) |
| self.calls = [] |
|
|
| def run(self, cmd): |
| self.calls.append(cmd) |
|
|
| m = RE_HEREDOC.search(cmd) |
| if m: |
| path, body = m.group(1), m.group(2) |
| self.fs[path] = body |
| return f"[mock env] wrote {len(body.splitlines())} lines to {path}" |
|
|
| if cmd.strip().startswith("ls"): |
| return "\n".join(sorted(self.fs)) + "\nsetup.py\nREADME.rst\ntests/" |
|
|
| if cmd.strip().startswith("cat "): |
| path = cmd.split()[-1] |
| if path in self.fs: |
| return ("[mock env] only the region near the bug is available\n" |
| + self.fs[path]) |
| return f"cat: {path}: No such file or directory" |
|
|
| if cmd.strip().startswith("grep"): |
| pat = re.findall(r'"([^"]*)"|\'([^\']*)\'', cmd) |
| pat = next((a or b for a, b in pat), "") |
| src = self.fs.get(cmd.split()[-1], "") |
| hits = [f"{i + 1}:{line}" for i, line in enumerate(src.split("\n")) |
| if pat and pat in line] |
| return "\n".join(hits) if hits else "(no matches)" |
|
|
| if "pytest" in cmd: |
| body = "\n".join(f"FAILED {n}" for n in self.failing_tests) |
| return (f"{body}\n=== {len(self.failing_tests)} failed ===\n" |
| "[mock env] test names are real, but nothing was executed") |
|
|
| name = cmd.split()[0] if cmd.split() else cmd |
| return (f"bash: {name}: command not found\n" |
| "[mock env] implements only: ls, cat, grep, pytest, cat > FILE <<'EOF'") |
|
|
| def patch(self): |
| """The candidate patch: a real unified diff of the virtual FS. |
| |
| A file the agent created did not exist at the start, so it diffs |
| against empty -- the same as `git diff` for a new file. |
| """ |
| out = [] |
| for path, body in self.fs.items(): |
| before = self.orig.get(path, "") |
| if body != before: |
| out += list(difflib.unified_diff( |
| before.split("\n"), body.split("\n"), |
| fromfile=f"a/{path}", tofile=f"b/{path}", lineterm="")) |
| return "\n".join(out) |
|
|
