AgentnessBench / tests /runtime /test_trace_accounting.py
irregular6612's picture
refactor(scenario): delete predator_evade; template is the canonical scenario
93cd78f
Raw
History Blame Contribute Delete
2.8 kB
from proteus.game.agents import VanillaAgent
from proteus.providers import FakeProvider
from proteus.game.runtime import SessionRunner, SessionTrace
def _run(response):
agent = VanillaAgent(FakeProvider(responses=[response], model_name="fake-1"))
return SessionRunner(
"template", agent, seed=42, play_turns=3, use_probe=False,
).run()
def test_act_tokens_and_raw_text_persisted_per_turn():
# FakeProvider sets output_tokens = word count; inline <think> gives thinking_tokens.
trace = _run("<think>predator is two cells east</think>ACTION: up")
t0 = trace.turns[0]
assert t0.raw_text == "<think>predator is two cells east</think>ACTION: up"
assert t0.output_tokens > 0 # propagated from CompletionResult
assert t0.input_tokens == 0 # FakeProvider reports 0; output_tokens below is the live propagation check
assert t0.thinking_tokens == 5 # whitespace-split of the think block
# use_probe=False: probe fields stay at their defaults.
assert t0.probe_raw_text == ""
assert t0.probe_thinking_tokens == 0
def test_accounting_survives_jsonl_roundtrip():
trace = _run("<think>a b c</think>ACTION: up")
reloaded = SessionTrace.model_validate_json(trace.model_dump_json())
assert reloaded.model_dump() == trace.model_dump()
assert reloaded.turns[0].thinking_tokens == 3
assert reloaded.turns[0].raw_text == "<think>a b c</think>ACTION: up"
def test_probe_accounting_persisted_when_enabled():
# Distinct probe vs act responses (different <think> word counts) so a
# probe/act field cross-wiring bug would be visible: turn 1's probe consumes
# responses[0] (3-word think) and turn 1's act consumes responses[1] (5-word think).
agent = VanillaAgent(FakeProvider(
responses=[
"<think>it is east</think>predator is east", # probe: 3 think-words
"<think>move up and away now</think>I should go up\nACTION: up", # act: 5 think-words
],
model_name="fake-1",
))
trace = SessionRunner(
"template", agent, seed=42, play_turns=3, use_probe=True,
).run()
t0 = trace.turns[0]
assert t0.probe_a # answer recorded
assert "it is east" in t0.probe_reasoning # probe reasoning preserved
assert t0.probe_raw_text # raw probe output preserved
assert t0.probe_thinking_tokens == 3 # from the PROBE response
assert t0.thinking_tokens == 5 # from the ACT response — NOT cross-wired
assert t0.probe_output_tokens > 0
# survives JSONL round-trip
reloaded = SessionTrace.model_validate_json(trace.model_dump_json())
assert reloaded.model_dump() == trace.model_dump()