-
Notifications
You must be signed in to change notification settings - Fork 21
Expand file tree
/
Copy pathtest_quickstart_utils.py
More file actions
388 lines (317 loc) · 15.9 KB
/
Copy pathtest_quickstart_utils.py
File metadata and controls
388 lines (317 loc) · 15.9 KB
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
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
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
"""Tests for quickstart utility functions."""
import pytest
from eval_protocol.models import EvaluationRow, InputMetadata, Message
from eval_protocol.utils.evaluation_row_utils import (
multi_turn_assistant_to_ground_truth,
serialize_message,
assistant_to_ground_truth,
)
class TestSerializeMessage:
"""Tests for serialize_message function."""
def test_simple_message(self):
"""Test serialization of a simple message."""
message = Message(role="user", content="Hello, how are you?")
result = serialize_message(message)
assert result == "user: Hello, how are you?"
def test_assistant_message(self):
"""Test serialization of an assistant message."""
message = Message(role="assistant", content="I'm doing well, thank you!")
result = serialize_message(message)
assert result == "assistant: I'm doing well, thank you!"
def test_message_with_tool_calls(self):
"""Test serialization of a message with tool calls."""
tool_call = {
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "New York"}'},
}
message = Message(
role="assistant",
content="I'll check the weather for you.",
tool_calls=[tool_call], # pyright: ignore[reportArgumentType]
)
result = serialize_message(message)
expected = 'assistant: I\'ll check the weather for you.\n[Tool Call: get_weather({"location": "New York"})]'
assert result == expected
def test_message_with_multiple_tool_calls(self):
"""Test serialization of a message with multiple tool calls."""
tool_call1 = {
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "NYC"}'},
}
tool_call2 = {
"id": "call_456",
"type": "function",
"function": {"name": "get_time", "arguments": '{"timezone": "EST"}'},
}
message = Message(
role="assistant",
content="Let me get both for you.",
tool_calls=[tool_call1, tool_call2], # pyright: ignore[reportArgumentType]
)
result = serialize_message(message)
expected = (
"assistant: Let me get both for you.\n"
'[Tool Call: get_weather({"location": "NYC"})]\n'
'[Tool Call: get_time({"timezone": "EST"})]'
)
assert result == expected
def test_empty_content_message(self):
"""Test serialization of a message with empty content."""
message = Message(role="assistant", content="")
result = serialize_message(message)
assert result == "assistant: "
def test_none_content_message(self):
"""Test serialization of a message with None content."""
message = Message(role="assistant", content=None)
result = serialize_message(message)
assert result == "assistant: None"
class TestMultiTurnAssistantToGroundTruth:
"""Tests for multi_turn_assistant_to_ground_truth function."""
def test_single_turn_conversation(self):
"""Test that single-turn conversations are handled correctly."""
messages = [
Message(role="user", content="What's the weather like?"),
Message(role="assistant", content="It's sunny today!"),
]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 1
assert len(result[0].messages) == 1 # Only user message before assistant
assert result[0].messages[0].role == "user"
assert result[0].messages[0].content == "What's the weather like?"
assert result[0].ground_truth == "assistant: It's sunny today!"
def test_multi_turn_conversation(self):
"""Test that multi-turn conversations are split correctly."""
messages = [
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi there!"),
Message(role="user", content="How are you?"),
Message(role="assistant", content="I'm doing well, thanks!"),
]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 2
# First split: user -> assistant
assert len(result[0].messages) == 1
assert result[0].messages[0].content == "Hello"
assert result[0].ground_truth == "assistant: Hi there!"
# Second split: user -> assistant -> user -> assistant
assert len(result[1].messages) == 3
assert result[1].messages[0].content == "Hello"
assert result[1].messages[1].content == "Hi there!"
assert result[1].messages[2].content == "How are you?"
assert result[1].ground_truth == "assistant: I'm doing well, thanks!"
def test_conversation_with_system_message(self):
"""Test that system messages are preserved in splits."""
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi there!"),
Message(role="user", content="How are you?"),
Message(role="assistant", content="I'm doing well!"),
]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 2
# First split should include system message
assert len(result[0].messages) == 2
assert result[0].messages[0].role == "system"
assert result[0].messages[1].role == "user"
# Second split should include system message and previous conversation
assert len(result[1].messages) == 4
assert result[1].messages[0].role == "system"
def test_conversation_with_tool_calls(self):
"""Test that tool calls are preserved in ground truth."""
tool_call = {
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "NYC"}'},
}
messages = [
Message(role="user", content="What's the weather in NYC?"),
Message(
role="assistant",
content="I'll check that for you.",
tool_calls=[tool_call], # pyright: ignore[reportArgumentType]
),
]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 1
expected_ground_truth = 'assistant: I\'ll check that for you.\n[Tool Call: get_weather({"location": "NYC"})]'
assert result[0].ground_truth == expected_ground_truth
def test_multiple_rows_processing(self):
"""Test that multiple input rows are processed correctly."""
row1 = EvaluationRow(
messages=[Message(role="user", content="Hello"), Message(role="assistant", content="Hi!")]
)
row2 = EvaluationRow(
messages=[Message(role="user", content="Goodbye"), Message(role="assistant", content="Bye!")]
)
result = multi_turn_assistant_to_ground_truth([row1, row2])
assert len(result) == 2
assert result[0].messages[0].content == "Hello"
assert result[0].ground_truth == "assistant: Hi!"
assert result[1].messages[0].content == "Goodbye"
assert result[1].ground_truth == "assistant: Bye!"
def test_no_assistant_messages(self):
"""Test that rows with no assistant messages return empty list."""
messages = [Message(role="user", content="Hello"), Message(role="user", content="Anyone there?")]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 0
def test_only_assistant_messages(self):
"""Test handling of rows with only assistant messages."""
messages = [Message(role="assistant", content="Hello!"), Message(role="assistant", content="How can I help?")]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 2
# First assistant message (no context)
assert len(result[0].messages) == 0
assert result[0].ground_truth == "assistant: Hello!"
# Second assistant message (with first assistant as context)
assert len(result[1].messages) == 1
assert result[1].messages[0].content == "Hello!"
assert result[1].ground_truth == "assistant: How can I help?"
def test_duplicate_trace_filtering(self):
"""Test that duplicate traces are filtered out."""
# Create two rows with the same conversation leading to different assistant responses
messages1 = [
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi there!"),
Message(role="user", content="How are you?"),
Message(role="assistant", content="I'm good!"),
]
messages2 = [
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi there!"),
Message(role="user", content="How are you?"),
Message(role="assistant", content="I'm great!"), # Different response
]
row1 = EvaluationRow(messages=messages1)
row2 = EvaluationRow(messages=messages2)
result = multi_turn_assistant_to_ground_truth([row1, row2])
# Should only get 2 unique splits (not 4), because the context leading
# to the second assistant message is the same in both rows
assert len(result) == 2 # First "Hello" -> "Hi there!", then one unique context for second assistant
# Verify the unique traces
contexts = ["\n".join(serialize_message(m) for m in r.messages) for r in result]
assert len(set(contexts)) == len(contexts) # All contexts should be unique
def test_tools_and_metadata_preservation(self):
"""Test that tools and input_metadata are preserved in split rows."""
tools = [{"type": "function", "function": {"name": "test_tool"}}]
input_metadata = InputMetadata(
row_id="test_row", completion_params={"model": "gpt-4"}, session_data={"test": "data"}
)
messages = [Message(role="user", content="Hello"), Message(role="assistant", content="Hi!")]
row = EvaluationRow(messages=messages, tools=tools, input_metadata=input_metadata)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 1
assert result[0].tools == tools
assert result[0].input_metadata == input_metadata
def test_empty_input_list(self):
"""Test that empty input list returns empty result."""
result = multi_turn_assistant_to_ground_truth([])
assert len(result) == 0
def test_complex_multi_turn_with_tool_responses(self):
"""Test complex conversation with tool calls and responses."""
tool_call = {
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "NYC"}'},
}
messages = [
Message(role="user", content="What's the weather in NYC?"),
Message(
role="assistant",
content="I'll check that for you.",
tool_calls=[tool_call], # pyright: ignore[reportArgumentType]
),
Message(role="tool", tool_call_id="call_123", content="Sunny, 75°F"),
Message(role="assistant", content="It's sunny and 75°F in NYC!"),
Message(role="user", content="Thanks!"),
Message(role="assistant", content="You're welcome!"),
]
row = EvaluationRow(messages=messages)
result = multi_turn_assistant_to_ground_truth([row])
assert len(result) == 3 # Three assistant messages
# First assistant message with tool call
assert len(result[0].messages) == 1 # Just user message
assert "Tool Call: get_weather" in str(result[0].ground_truth or "")
# Second assistant message after tool response
assert len(result[1].messages) == 3 # user, assistant with tool call, tool response
assert result[1].ground_truth == "assistant: It's sunny and 75°F in NYC!"
# Third assistant message
assert len(result[2].messages) == 5 # All previous messages + "Thanks!"
assert result[2].ground_truth == "assistant: You're welcome!"
class TestAssistantToGroundTruth:
"""Tests for assistant_to_ground_truth function."""
def test_removes_last_assistant_message(self):
"""Test that the last assistant message is removed and set as ground truth."""
messages = [
Message(role="user", content="What's the weather like?"),
Message(role="assistant", content="It's sunny today!"),
]
row = EvaluationRow(messages=messages)
result = assistant_to_ground_truth([row])
assert len(result) == 1
assert len(result[0].messages) == 1 # Only user message remains
assert result[0].messages[0].role == "user"
assert result[0].messages[0].content == "What's the weather like?"
assert result[0].ground_truth == "assistant: It's sunny today!"
def test_multi_turn_with_last_assistant(self):
"""Test multi-turn conversation where last message is assistant."""
messages = [
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi there!"),
Message(role="user", content="How are you?"),
Message(role="assistant", content="I'm doing well!"),
]
row = EvaluationRow(messages=messages)
result = assistant_to_ground_truth([row])
assert len(result) == 1
assert len(result[0].messages) == 3 # All except last assistant
assert result[0].messages[-1].content == "How are you?"
assert result[0].ground_truth == "assistant: I'm doing well!"
def test_fails_when_last_message_not_assistant(self):
"""Test that function raises error when last message is not from assistant."""
messages = [
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi!"),
Message(role="user", content="Goodbye"),
]
row = EvaluationRow(messages=messages)
with pytest.raises(ValueError, match="Last message is not from assistant"):
assistant_to_ground_truth([row])
def test_preserves_metadata_and_tools(self):
"""Test that tools and metadata are preserved."""
messages = [
Message(role="user", content="Hello"),
Message(role="assistant", content="Hi there!"),
]
tools = [{"type": "function", "function": {"name": "test"}}]
input_metadata = InputMetadata(row_id="test_123", completion_params={})
row = EvaluationRow(messages=messages, tools=tools, input_metadata=input_metadata)
result = assistant_to_ground_truth([row])
assert len(result) == 1
assert result[0].tools == tools
assert result[0].input_metadata == input_metadata
assert result[0].ground_truth == "assistant: Hi there!"
def test_multiple_rows(self):
"""Test processing multiple rows."""
row1 = EvaluationRow(
messages=[Message(role="user", content="Hello"), Message(role="assistant", content="Hi!")]
)
row2 = EvaluationRow(
messages=[Message(role="user", content="Bye"), Message(role="assistant", content="Goodbye!")]
)
result = assistant_to_ground_truth([row1, row2])
assert len(result) == 2
assert result[0].ground_truth == "assistant: Hi!"
assert result[1].ground_truth == "assistant: Goodbye!"
def test_empty_input_list(self):
"""Test that empty input list returns empty result."""
result = assistant_to_ground_truth([])
assert len(result) == 0