Skip to content

test_trace

ipw.tests.core.test_trace

Tests for execution/trace.py — TurnTrace and QueryTrace.

TestTurnTrace

Test TurnTrace construction and serialization.

Source code in intelligence-per-watt/src/ipw/tests/core/test_trace.py 
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
class TestTurnTrace:
    """Test TurnTrace construction and serialization."""

    def test_creation_with_all_fields(self) -> None:
        turn = TurnTrace(
            turn_index=0,
            input_tokens=100,
            output_tokens=50,
            tool_result_tokens=20,
            tools_called=["calc", "search"],
            tool_latencies_s={"calc": 0.1, "search": 0.5},
            wall_clock_s=1.5,
            error=None,
            gpu_energy_joules=10.0,
            cpu_energy_joules=3.0,
            gpu_power_avg_watts=150.0,
            cpu_power_avg_watts=45.0,
            cost_usd=0.005,
        )
        assert turn.turn_index == 0
        assert turn.input_tokens == 100
        assert turn.tools_called == ["calc", "search"]
        assert turn.gpu_energy_joules == 10.0

    def test_to_dict_roundtrip(self) -> None:
        turn = TurnTrace(
            turn_index=2,
            input_tokens=50,
            output_tokens=25,
            tools_called=["think"],
            wall_clock_s=0.8,
            gpu_energy_joules=5.0,
            cost_usd=0.001,
        )
        d = turn.to_dict()
        restored = TurnTrace.from_dict(d)
        assert restored.turn_index == turn.turn_index
        assert restored.input_tokens == turn.input_tokens
        assert restored.output_tokens == turn.output_tokens
        assert restored.tools_called == turn.tools_called
        assert restored.wall_clock_s == turn.wall_clock_s
        assert restored.gpu_energy_joules == turn.gpu_energy_joules
        assert restored.cost_usd == turn.cost_usd

    def test_from_dict_defaults(self) -> None:
        d = {"turn_index": 0}
        turn = TurnTrace.from_dict(d)
        assert turn.input_tokens == 0
        assert turn.output_tokens == 0
        assert turn.tools_called == []
        assert turn.error is None

TestQueryTrace

Test QueryTrace construction, properties, and serialization.

Source code in intelligence-per-watt/src/ipw/tests/core/test_trace.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class TestQueryTrace:
    """Test QueryTrace construction, properties, and serialization."""

    def _make_trace(self) -> QueryTrace:
        turns = [
            TurnTrace(
                turn_index=0,
                input_tokens=100,
                output_tokens=50,
                tools_called=["calc"],
                wall_clock_s=1.0,
                gpu_energy_joules=5.0,
                cost_usd=0.002,
            ),
            TurnTrace(
                turn_index=1,
                input_tokens=80,
                output_tokens=30,
                tools_called=["search", "think"],
                wall_clock_s=2.0,
                gpu_energy_joules=8.0,
                cost_usd=0.003,
            ),
        ]
        return QueryTrace(
            query_id="q0001",
            workload_type="agentic",
            query_text="What is 2+2?",
            response_text="4",
            turns=turns,
            total_wall_clock_s=3.0,
            completed=True,
        )

    def test_num_turns(self) -> None:
        trace = self._make_trace()
        assert trace.num_turns == 2

    def test_total_input_tokens(self) -> None:
        trace = self._make_trace()
        assert trace.total_input_tokens == 180

    def test_total_output_tokens(self) -> None:
        trace = self._make_trace()
        assert trace.total_output_tokens == 80

    def test_total_tool_calls(self) -> None:
        trace = self._make_trace()
        assert trace.total_tool_calls == 3

    def test_total_gpu_energy_joules(self) -> None:
        trace = self._make_trace()
        assert trace.total_gpu_energy_joules == 13.0

    def test_total_gpu_energy_joules_none_when_no_values(self) -> None:
        trace = QueryTrace(
            query_id="q0",
            workload_type="test",
            turns=[TurnTrace(turn_index=0)],
        )
        assert trace.total_gpu_energy_joules is None

    def test_total_cost_usd(self) -> None:
        trace = self._make_trace()
        assert trace.total_cost_usd == pytest.approx(0.005)

    def test_total_cost_usd_none_when_no_values(self) -> None:
        trace = QueryTrace(
            query_id="q0",
            workload_type="test",
            turns=[TurnTrace(turn_index=0)],
        )
        assert trace.total_cost_usd is None

    def test_to_dict_from_dict_roundtrip(self) -> None:
        trace = self._make_trace()
        d = trace.to_dict()
        restored = QueryTrace.from_dict(d)
        assert restored.query_id == trace.query_id
        assert restored.workload_type == trace.workload_type
        assert restored.num_turns == trace.num_turns
        assert restored.total_input_tokens == trace.total_input_tokens
        assert restored.completed == trace.completed

    def test_jsonl_save_load_roundtrip(self, tmp_path: Path) -> None:
        trace = self._make_trace()
        jsonl_path = tmp_path / "traces.jsonl"

        trace.save_jsonl(jsonl_path)
        loaded = QueryTrace.load_jsonl(jsonl_path)

        assert len(loaded) == 1
        assert loaded[0].query_id == "q0001"
        assert loaded[0].num_turns == 2
        assert loaded[0].total_input_tokens == 180

    def test_jsonl_multiple_traces(self, tmp_path: Path) -> None:
        trace1 = self._make_trace()
        trace2 = QueryTrace(
            query_id="q0002",
            workload_type="agentic",
            query_text="Another question",
            response_text="Another answer",
            turns=[TurnTrace(turn_index=0, input_tokens=10, output_tokens=5)],
            total_wall_clock_s=0.5,
            completed=True,
        )

        jsonl_path = tmp_path / "traces.jsonl"
        trace1.save_jsonl(jsonl_path)
        trace2.save_jsonl(jsonl_path)

        loaded = QueryTrace.load_jsonl(jsonl_path)
        assert len(loaded) == 2
        assert loaded[0].query_id == "q0001"
        assert loaded[1].query_id == "q0002"

    def test_to_hf_dataset(self) -> None:
        try:
            from datasets import Dataset
        except ImportError:
            pytest.skip("datasets package not available")

        trace = self._make_trace()
        ds = QueryTrace.to_hf_dataset([trace])

        assert len(ds) == 1
        row = ds[0]
        assert row["query_id"] == "q0001"
        assert row["num_turns"] == 2
        assert row["total_input_tokens"] == 180
        assert row["total_output_tokens"] == 80
        assert row["total_tool_calls"] == 3
        assert row["completed"] is True
        # trace_json should be valid JSON
        parsed = json.loads(row["trace_json"])
        assert parsed["query_id"] == "q0001"