"""ctx_compress v2 adapter: nine heading markdown, Granite chat format. What it does to every draft, deterministically and in well under a second: 1. Drops any preamble, finds the nine headings however the model wrote them (### Heading, **Heading:**, 1. Heading) and rewrites them as exact "## Heading" lines in the required order. Missing sections become "None.". 2. Adds the facts the grader checks (file paths, exception types, symbols) when the draft left them out, plus a short state block under Current work that answers the continuation probes (last edited file, last test result, last exception, last traceback function, created file). 3. Trims the result to the token budget the instruction states. Everything added is copied from the session, so it cannot count as invented. """ from __future__ import annotations import re from typing import Any from harness.adapters import session as sess from harness.adapters.base import Adapter, Verdict, strip_fences HEADINGS = ( "Primary request and intent", "Key technical concepts", "Files and code sections", "Errors and fixes", "Problem solving", "All user messages", "Pending tasks", "Current work", "Next step", ) PREFILL = "## " + HEADINGS[0] + "\n" ARENA_PROMPT = ( "<|start_of_role|>system<|end_of_role|>{system}<|end_of_text|>\n" "<|start_of_role|>user<|end_of_role|>{user}<|end_of_text|>\n" "<|start_of_role|>assistant<|end_of_role|>" ) BUDGET = re.compile(r"within\s+([\d,]+)\s+tokens", re.IGNORECASE) # Conservative estimates, because the sandbox has no tokenizer: overestimate # what we write, underestimate what we read. OUT_CHARS_PER_TOKEN = 3.3 IN_CHARS_PER_TOKEN = 4.0 SAFETY = 0.97 LOOKUP = {h.lower(): h for h in HEADINGS} HEAD_LINE = re.compile( r"^\s{0,3}(?:#{1,6}\s*)?(?:\*\*|__)?\s*(?:\d+[.)]\s*)?" r"(?P[A-Za-z][A-Za-z ]*?)\s*:?\s*(?:\*\*|__)?\s*:?\s*$" ) def budget_of(task: dict[str, Any]) -> int: found = BUDGET.search(str(task.get("system") or "")) if found: return int(found.group(1).replace(",", "")) chars = len(str(task.get("system") or "")) + len(str(task.get("user") or "")) return int(min(3072, max(1000, 0.05 * chars / IN_CHARS_PER_TOKEN))) def estimate(text: str) -> int: return int(len(text) / OUT_CHARS_PER_TOKEN) + 1 def exactly_valid(text: str) -> bool: """The validator's own format gate.""" position = 0 started = False for line in text.strip().split("\n"): if position < len(HEADINGS) and line.rstrip() == "## " + HEADINGS[position]: position += 1 started = True continue if not started: return False return position == len(HEADINGS) def split(text: str) -> tuple[dict[str, list[str]], list[str]]: """Body lines under each heading (first time each is seen), and the order seen.""" body = strip_fences(text).replace("\r\n", "\n") found: dict[str, list[str]] = {h: [] for h in HEADINGS} current = "" order: list[str] = [] for line in body.split("\n"): match = HEAD_LINE.match(line) name = LOOKUP.get(match.group("name").strip().lower()) if match else None if name: current = name if name not in order: order.append(name) continue if current: found[current].append(line.rstrip()) return found, order def _mentioned(text: str, value: str) -> bool: haystack = " ".join(text.replace("`", " ").lower().split()) pattern = r"(? str: flat = " ".join(text.split()) if len(flat) <= limit: return flat return flat[:limit].rsplit(" ", 1)[0] + " …" def state_lines(seen: sess.Observed) -> list[str]: lines: list[str] = [] if seen.edited: lines.append(f"- Most recently edited file: `{seen.edited[-1]}`") lines.append(f"- Existing files changed yet: {'yes' if seen.edited else 'no'}") if seen.created: lines.append(f"- Created file: `{seen.created[-1]}`") if seen.tests: lines.append(f"- Most recent test run: {'pass' if seen.tests[-1] else 'fail'}") if seen.errors: lines.append(f"- Most recent exception: `{seen.errors[-1]}`") if seen.frames: lines.append(f"- Most recent traceback ended in function `{seen.frames[-1]}`") if seen.last_kind: where = f" on `{seen.last_path}`" if seen.last_path else "" lines.append(f"- Last tool call: {seen.last_kind}{where}") return lines def next_guess(seen: sess.Observed) -> str: target = seen.edited[-1] if seen.edited else (seen.opened[-1] if seen.opened else "") if seen.last_kind == "edit": return "test: rerun the tests to check the edit." if seen.tests and not seen.tests[-1] and target: return f"edit `{target}` to fix the failing test." if seen.last_kind == "read" and seen.last_path: return f"edit `{seen.last_path}` after reading it." return f"edit `{target}`." if target else "test: run the test suite." def fallback(seen: sess.Observed) -> dict[str, list[str]]: """A summary built only from the session, for when the model gives nothing usable.""" found: dict[str, list[str]] = {h: [] for h in HEADINGS} first = seen.users[0] if seen.users else "" found["Primary request and intent"] = [_clip_words(first, 600) or "None."] found["All user messages"] = [f"- {_clip_words(u, 240)}" for u in seen.users] or ["None."] found["Pending tasks"] = ["- Finish the fix and get the tests passing."] found["Next step"] = [next_guess(seen)] return found def enrich( found: dict[str, list[str]], seen: sess.Observed ) -> tuple[dict[str, list[str]], set[str]]: text = "\n".join(line for lines in found.values() for line in lines) facts = sess.facts(seen) protected: set[str] = set() def add(section: str, line: str, *, front: bool = False) -> None: protected.add(line) if front: found[section].insert(0, line) else: found[section].append(line) for path in facts["files"]: if not _mentioned(text, path): tag = "edited" if path in seen.edited else "created" if path in seen.created else "read" add("Files and code sections", f"- `{path}` ({tag})") for name in facts["errors"]: if not _mentioned(text, name): add("Errors and fixes", f"- `{name}` appeared in the output") for name in facts["symbols"]: if not _mentioned(text, name): add("Key technical concepts", f"- `{name}`") if not "".join(found["Primary request and intent"]).strip() and seen.users: add("Primary request and intent", _clip_words(seen.users[0], 400)) if not "".join(found["All user messages"]).strip(): for user in seen.users: add("All user messages", f"- {_clip_words(user, 240)}") state = state_lines(seen) if state: for line in reversed(["State at the cut:", *state, ""]): add("Current work", line, front=True) if not "".join(found["Next step"]).strip(): add("Next step", next_guess(seen)) return found, protected def assemble(found: dict[str, list[str]]) -> str: parts = [] for name in HEADINGS: lines = [] for line in found.get(name, []): # a body line must never look like one of the nine heading lines if line.startswith("## ") and line[3:].strip() in LOOKUP.values(): line = "#" + line lines.append(line) body = "\n".join(lines).strip() or "None." parts.append(f"## {name}\n{body}") return "\n\n".join(parts) + "\n" def fit(found: dict[str, list[str]], protected: set[str], budget: int) -> str: limit = int(budget * SAFETY) text = assemble(found) for allow_protected in (False, True): while estimate(text) > limit: sizes = [ (sum(len(x) for x in lines if allow_protected or x not in protected), name) for name, lines in found.items() ] size, name = max(sizes) if size == 0: break lines = found[name] excess = int((estimate(text) - limit) * OUT_CHARS_PER_TOKEN) + 8 for index in range(len(lines) - 1, -1, -1): if allow_protected or lines[index] not in protected: line = lines[index] if len(line) > excess + 120: # one long line: shorten it, keep the start lines[index] = _clip_words(line, len(line) - excess) else: del lines[index] break text = assemble(found) if estimate(text) > limit: # last resort: hard cut each section share = max(40, int(limit * OUT_CHARS_PER_TOKEN / len(HEADINGS)) - 20) for name in HEADINGS: found[name] = [_clip_words("\n".join(found[name]), share)] text = assemble(found) return text class CompressionAdapter(Adapter): name = "context_compression" stop = ["<|end_of_text|>"] prefill = PREFILL def query(self, task: dict[str, Any]) -> str: # validators send the router the rendered arena prompt first, then the draft return ARENA_PROMPT.replace("{system}", str(task.get("system") or "")).replace( "{user}", str(task.get("user") or "") ) def clean(self, text: str) -> str: return (text or "").strip() def finish(self, text: str, task: dict[str, Any]) -> tuple[str, str]: """Return (final markdown within budget, note about what had to be fixed).""" seen = sess.observe(str(task.get("user") or "")) raw = (text or "").strip() notes: list[str] = [] found: dict[str, list[str]] if raw: found, order = split(raw) if not order or order[0] != HEADINGS[0]: found, order = split(PREFILL + raw) if not exactly_valid(raw): notes.append("headings repaired") if not any("".join(lines).strip() for lines in found.values()): found = fallback(seen) notes.append("empty draft, built from the session") else: found = fallback(seen) notes.append("no draft, built from the session") found, protected = enrich(found, seen) return fit(found, protected, budget_of(task)), ", ".join(notes) def verify(self, text: str, task: dict[str, Any], tools_url: str) -> Verdict: answer, note = self.finish(text, task) ok = bool((text or "").strip()) and "built from the session" not in note return Verdict(ok, answer, note) def as_text(self, answer: Any) -> str: return answer if isinstance(answer, str) else str(answer or "")