File size: 1,738 Bytes
55b60a8 | 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 64 65 66 67 68 69 70 | import json
from flashtrace.result import TokenScore, TraceResult
def make_result():
return TraceResult(
prompt_tokens=[" alpha", " beta", " gamma"],
generation_tokens=[" answer"],
scores=[0.2, 0.7, 0.1],
per_hop_scores=[[0.1, 0.4, 0.0], [0.1, 0.3, 0.1]],
thinking_ratios=[0.5, 0.2],
output_span=(0, 0),
reasoning_span=(0, 0),
method="flashtrace",
metadata={"model": "tiny"},
)
def test_topk_inputs_sorted():
result = make_result()
top = result.topk_inputs(2)
assert top == [
TokenScore(index=1, token=" beta", score=0.7),
TokenScore(index=0, token=" alpha", score=0.2),
]
def test_to_dict_is_json_serializable():
result = make_result()
payload = result.to_dict()
assert payload["method"] == "flashtrace"
assert payload["top_inputs"][0]["token"] == " beta"
json.dumps(payload)
def test_to_dict_sanitizes_tensor_metadata():
import torch
result = TraceResult(
prompt_tokens=[" alpha"],
generation_tokens=[" answer"],
scores=[1.0],
metadata={"tensor": torch.tensor([1.0, 2.0]), "object": object()},
)
payload = result.to_dict()
assert payload["metadata"]["tensor"] == [1.0, 2.0]
assert isinstance(payload["metadata"]["object"], str)
json.dumps(payload)
def test_json_and_html_export(tmp_path):
result = make_result()
json_path = tmp_path / "trace.json"
html_path = tmp_path / "trace.html"
result.to_json(json_path)
result.to_html(html_path)
assert json_path.read_text(encoding="utf-8").startswith("{")
html = html_path.read_text(encoding="utf-8")
assert "<html" in html
assert " beta" in html
|