| from __future__ import annotations |
|
|
| import re |
|
|
| from albedo_eval_service.shared.observation_format import first_bash_block |
|
|
| _THOUGHT_LINE_RE = re.compile(r"^\s*THOUGHT:\s*(.*)$", re.IGNORECASE) |
| _THINK_BLOCK_RE = re.compile(r"<think>(.*?)</think>", re.DOTALL | re.IGNORECASE) |
| _FENCE_RE = re.compile(r"```[\w-]*[ \t]*\n.*?```", re.DOTALL) |
|
|
|
|
| def thought_text(gold: str) -> str: |
| """Pull a thinking body out of rendered gold (`THOUGHT:` or `<think>`).""" |
| text = gold or "" |
| blocks = _THINK_BLOCK_RE.findall(text) |
| if blocks: |
| return blocks[0].strip() |
| lines: list[str] = [] |
| taking = False |
| for line in text.splitlines(): |
| match = _THOUGHT_LINE_RE.match(line) |
| if match: |
| taking = True |
| if match.group(1).strip(): |
| lines.append(match.group(1).strip()) |
| continue |
| if taking: |
| if line.strip().startswith("```") or line.strip().startswith("<"): |
| break |
| lines.append(line) |
| if lines: |
| return "\n".join(lines).strip() |
| before = _FENCE_RE.split(text, maxsplit=1)[0] |
| before = _THINK_BLOCK_RE.sub("", before) |
| cleaned = re.sub(r"^\s*THOUGHT:\s*", "", before, flags=re.I).strip() |
| return cleaned |
|
|
|
|
| def wrap_completion(gold: str, bash_override: str | None = None) -> str | None: |
| """One closed `<think>` block plus exactly one bash fence. Live thinking is on.""" |
| bash = (bash_override or first_bash_block(gold) or "").strip() |
| if not bash: |
| return None |
| thought = thought_text(gold) |
| if not thought: |
| thought = _fallback_thought(bash) |
| thought = thought.replace("</think>", "").strip() |
| return f"<think>\n{thought}\n</think>\n\n```bash\n{bash}\n```\n" |
|
|
|
|
| def _fallback_thought(bash: str) -> str: |
| head = bash.splitlines()[0][:160] |
| if bash.startswith("echo ") and "SUBMIT" in bash.upper(): |
| return f"Work is saved. Submit with the exact command from the instructions: {head}" |
| return f"Next step is to run this command and inspect the result: {head}" |
|
|