File size: 2,554 Bytes
4397e12 | 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 | import random
from tiny_agent.chat import parse_assistant, render
from tiny_agent.tasks import KINDS, check, grounded, make_task, oracle_trajectory
from tiny_agent.tools import Workspace, truncate
def test_sandbox_isolation():
with Workspace({"a.txt": "hello\n"}) as ws:
assert "hello" in ws.tool_bash("cat a.txt")
assert "nonet" in ws.tool_bash("curl -s -m 2 https://example.com >/dev/null && echo net || echo nonet")
assert "No such file" in ws.tool_bash("ls /home")
assert "Read-only" in ws.tool_bash("touch /usr/x")
assert "timed out" in ws.tool_bash("sleep 5", timeout=1)
def test_read_edit_rules():
with Workspace({"f.py": "x = 1\ny = 1\n"}) as ws:
assert "read f.py before" in ws.tool_edit("f.py", "x = 1", "x = 2")
out = ws.tool_read("f.py")
assert " 1\tx = 1" in out
assert "appears 2 times" in ws.tool_edit("f.py", "= 1", "= 3")
assert "not found" in ws.tool_edit("f.py", "z = 1", "z = 2")
assert "Edited" in ws.tool_edit("f.py", "x = 1", "x = 2")
assert "outside /work" in ws.tool_read("../../etc/passwd")
assert "unknown tool" in ws.call("browse", {})
assert "bad arguments" in ws.call("read", {"file": "f.py"})
def test_truncation_marker():
t = truncate("x\n" * 10000, limit=600)
assert "truncated" in t and len(t) < 800
def test_parse_assistant():
m = parse_assistant('<think>\nhmm\n</think>\n<tool_call>\n{"name": "bash", "arguments": {"command": "ls"}}\n</tool_call>'
'<tool_call>{bad json}</tool_call><|im_end|>')
assert m["think"] == "hmm"
assert m["tool_calls"][0] == {"name": "bash", "arguments": {"command": "ls"}}
assert "error" in m["tool_calls"][1]
def test_render_roundtrip_assistant():
msg = {"role": "assistant", "think": "t", "content": "", "tool_calls": [{"name": "submit", "arguments": {"answer": "7"}}]}
text = render([msg]).split("<|im_start|>assistant\n", 1)[1]
back = parse_assistant(text)
assert back["tool_calls"] == msg["tool_calls"] and back["think"] == "t"
def test_oracle_solves_every_kind():
for i, kind in enumerate(KINDS * 3):
t = make_task(random.Random(5000 + i), kind)
msgs, ok = oracle_trajectory(t, random.Random(i))
assert ok, (kind, t.question)
assert grounded(t, msgs)
def test_wrong_answers_fail():
t = make_task(random.Random(1), "config_value")
assert not check(t, "definitely wrong")
assert not check(t, None)
assert check(t, f" {t.answer}. ")
|