81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from markovian_rsa_mlx.audit import (
|
|
AuditWriter,
|
|
RunStartEvent,
|
|
TraceCompleteEvent,
|
|
TailExtractedEvent,
|
|
AggregationPromptEvent,
|
|
RoundCompleteEvent,
|
|
FinalEvent,
|
|
RunEndEvent,
|
|
)
|
|
from markovian_rsa_mlx.config import RSAConfig
|
|
|
|
|
|
def test_writer_emits_jsonl(tmp_path: Path):
|
|
audit_path = tmp_path / "audit.jsonl"
|
|
writer = AuditWriter(audit_path)
|
|
writer.write(RunStartEvent(
|
|
run_id="run-1",
|
|
model_id="kyr0/zaya1-base-8b-MLX",
|
|
quantization="q4_g64",
|
|
config=RSAConfig(),
|
|
prompt="2+2",
|
|
created_at="2026-05-10T00:00:00Z",
|
|
))
|
|
writer.write(TraceCompleteEvent(
|
|
run_id="run-1",
|
|
round=0,
|
|
trace_id="t0",
|
|
text="The answer is 4.",
|
|
token_ids=[1, 2, 3],
|
|
generated_tokens=3,
|
|
finish_reason="eos",
|
|
elapsed_s=0.5,
|
|
))
|
|
writer.close()
|
|
|
|
lines = audit_path.read_text().strip().split("\n")
|
|
assert len(lines) == 2
|
|
e0 = json.loads(lines[0])
|
|
assert e0["event"] == "run_start"
|
|
assert e0["model_id"] == "kyr0/zaya1-base-8b-MLX"
|
|
assert e0["config"]["rounds"] == 2
|
|
e1 = json.loads(lines[1])
|
|
assert e1["event"] == "trace_complete"
|
|
assert e1["text"] == "The answer is 4."
|
|
assert e1["finish_reason"] == "eos"
|
|
|
|
|
|
def test_event_serialization_includes_event_field():
|
|
ev = TailExtractedEvent(
|
|
run_id="r", round=0, trace_id="t",
|
|
tail_token_ids=[1, 2, 3], tail_text="hi", tail_tokens=3,
|
|
)
|
|
payload = ev.to_dict()
|
|
assert payload["event"] == "tail_extracted"
|
|
|
|
|
|
def test_writer_appends_subsequent_events(tmp_path: Path):
|
|
audit_path = tmp_path / "audit.jsonl"
|
|
w = AuditWriter(audit_path)
|
|
w.write(RunEndEvent(run_id="r", elapsed_s=1.0,
|
|
total_generated_tokens=10, peak_memory_bytes=1000))
|
|
w.close()
|
|
w2 = AuditWriter(audit_path, mode="a")
|
|
w2.write(RunEndEvent(run_id="r2", elapsed_s=2.0,
|
|
total_generated_tokens=20, peak_memory_bytes=2000))
|
|
w2.close()
|
|
lines = audit_path.read_text().strip().split("\n")
|
|
assert len(lines) == 2
|
|
|
|
|
|
def test_writer_no_op_when_path_none():
|
|
w = AuditWriter(None)
|
|
w.write(RunEndEvent(run_id="x", elapsed_s=1.0,
|
|
total_generated_tokens=0, peak_memory_bytes=0))
|
|
w.close() # no exception
|