mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 13:26:48 +00:00
152 lines
4.7 KiB
Python
152 lines
4.7 KiB
Python
"""Tests for QueryTrace and TurnTrace data model."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from openjarvis.evals.core.trace import QueryTrace, TurnTrace
|
|
|
|
|
|
class TestTurnTrace:
|
|
def test_defaults(self):
|
|
t = TurnTrace(turn_index=0)
|
|
assert t.input_tokens == 0
|
|
assert t.output_tokens == 0
|
|
assert t.tools_called == []
|
|
assert t.gpu_energy_joules is None
|
|
assert t.cost_usd is None
|
|
|
|
def test_to_dict_roundtrip(self):
|
|
t = TurnTrace(
|
|
turn_index=1,
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
tools_called=["calculator"],
|
|
tool_latencies_s={"calculator": 0.5},
|
|
wall_clock_s=1.2,
|
|
gpu_energy_joules=10.0,
|
|
cost_usd=0.001,
|
|
)
|
|
d = t.to_dict()
|
|
t2 = TurnTrace.from_dict(d)
|
|
assert t2.turn_index == 1
|
|
assert t2.input_tokens == 100
|
|
assert t2.output_tokens == 50
|
|
assert t2.tools_called == ["calculator"]
|
|
assert t2.tool_latencies_s == {"calculator": 0.5}
|
|
assert t2.wall_clock_s == pytest.approx(1.2)
|
|
assert t2.gpu_energy_joules == pytest.approx(10.0)
|
|
assert t2.cost_usd == pytest.approx(0.001)
|
|
|
|
def test_from_dict_missing_keys(self):
|
|
d = {"turn_index": 0}
|
|
t = TurnTrace.from_dict(d)
|
|
assert t.input_tokens == 0
|
|
assert t.tools_called == []
|
|
assert t.error is None
|
|
|
|
|
|
class TestQueryTrace:
|
|
def _make_trace(self, **kwargs):
|
|
defaults = {
|
|
"query_id": "q0001",
|
|
"workload_type": "coding",
|
|
"query_text": "Hello",
|
|
"response_text": "World",
|
|
"turns": [
|
|
TurnTrace(
|
|
turn_index=0,
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
wall_clock_s=1.0,
|
|
gpu_energy_joules=5.0,
|
|
cost_usd=0.01,
|
|
),
|
|
TurnTrace(
|
|
turn_index=1,
|
|
input_tokens=150,
|
|
output_tokens=75,
|
|
wall_clock_s=1.5,
|
|
gpu_energy_joules=7.5,
|
|
cost_usd=0.015,
|
|
),
|
|
],
|
|
"total_wall_clock_s": 2.5,
|
|
"completed": True,
|
|
}
|
|
defaults.update(kwargs)
|
|
return QueryTrace(**defaults)
|
|
|
|
def test_num_turns(self):
|
|
t = self._make_trace()
|
|
assert t.num_turns == 2
|
|
|
|
def test_total_tokens(self):
|
|
t = self._make_trace()
|
|
assert t.total_input_tokens == 250
|
|
assert t.total_output_tokens == 125
|
|
assert t.total_tokens == 375
|
|
|
|
def test_total_gpu_energy(self):
|
|
t = self._make_trace()
|
|
assert t.total_gpu_energy_joules == pytest.approx(12.5)
|
|
|
|
def test_total_gpu_energy_fallback(self):
|
|
t = self._make_trace(
|
|
turns=[TurnTrace(turn_index=0)],
|
|
query_gpu_energy_joules=20.0,
|
|
)
|
|
assert t.total_gpu_energy_joules == pytest.approx(20.0)
|
|
|
|
def test_total_cost(self):
|
|
t = self._make_trace()
|
|
assert t.total_cost_usd == pytest.approx(0.025)
|
|
|
|
def test_total_cost_none(self):
|
|
t = self._make_trace(turns=[TurnTrace(turn_index=0)])
|
|
assert t.total_cost_usd is None
|
|
|
|
def test_throughput(self):
|
|
t = self._make_trace()
|
|
assert t.throughput_tokens_per_sec == pytest.approx(125 / 2.5)
|
|
|
|
def test_energy_per_token(self):
|
|
t = self._make_trace()
|
|
assert t.energy_per_token_joules == pytest.approx(12.5 / 125)
|
|
|
|
def test_avg_gpu_power(self):
|
|
t = self._make_trace()
|
|
# No per-turn power set, so falls back to query-level
|
|
assert t.avg_gpu_power_watts is None
|
|
|
|
def test_to_dict_roundtrip(self):
|
|
t = self._make_trace(is_resolved=True)
|
|
d = t.to_dict()
|
|
t2 = QueryTrace.from_dict(d)
|
|
assert t2.query_id == "q0001"
|
|
assert t2.workload_type == "coding"
|
|
assert t2.num_turns == 2
|
|
assert t2.completed is True
|
|
assert t2.is_resolved is True
|
|
assert t2.total_input_tokens == 250
|
|
|
|
def test_save_load_jsonl(self, tmp_path):
|
|
t1 = self._make_trace(query_id="q0001")
|
|
t2 = self._make_trace(query_id="q0002")
|
|
path = tmp_path / "traces.jsonl"
|
|
t1.save_jsonl(path)
|
|
t2.save_jsonl(path)
|
|
loaded = QueryTrace.load_jsonl(path)
|
|
assert len(loaded) == 2
|
|
assert loaded[0].query_id == "q0001"
|
|
assert loaded[1].query_id == "q0002"
|
|
|
|
def test_tool_call_count(self):
|
|
t = self._make_trace(
|
|
turns=[
|
|
TurnTrace(turn_index=0, tools_called=["calc", "search"]),
|
|
TurnTrace(turn_index=1, tools_called=["read"]),
|
|
]
|
|
)
|
|
assert t.total_tool_calls == 3
|