File size: 1,228 Bytes
2abcc30 | 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 | from local_train.think import thought_text, wrap_completion
def test_wrap_converts_thought_prefix_and_keeps_one_bash_block():
gold = "THOUGHT: Look at retention.go\n\n```bash\ncat /testbed/retention.go\n```\n```bash\nls\n```"
out = wrap_completion(gold)
assert out is not None
assert out.startswith("<think>\nLook at retention.go")
assert "</think>" in out
assert out.count("```bash") == 1
assert "cat /testbed/retention.go" in out
assert "ls\n```" not in out
def test_wrap_override_rewrites_submit_command():
gold = "THOUGHT: done\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```"
out = wrap_completion(gold, "echo ALBEDO_TASK_DONE_SUBMIT_NOW")
assert out is not None
assert "echo ALBEDO_TASK_DONE_SUBMIT_NOW" in out
assert "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" not in out.split("</think>", 1)[-1]
def test_wrap_rejects_missing_bash():
assert wrap_completion("THOUGHT: just thinking") is None
def test_existing_think_block_is_reused():
gold = "<think>plan</think>\n\n```bash\nsed -i 's/a/b/' foo.py\n```"
assert thought_text(gold) == "plan"
out = wrap_completion(gold)
assert out is not None
assert "<think>\nplan\n</think>" in out
|