forked from eval-protocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_models_rl.py
More file actions
158 lines (141 loc) · 6.94 KB
/
Copy pathtest_models_rl.py
File metadata and controls
158 lines (141 loc) · 6.94 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
from typing import Any, Dict, List, Union
import pytest
from pydantic import ValidationError
from eval_protocol.agent.models import StepData
# Assuming these are the correct import paths based on our plan
from eval_protocol.models import EvaluateResult, Message as RewardKitMessage, StepOutput
# Minimal Message for StepData if direct import from eval_protocol.models is problematic in tests
# For now, assume RewardKitMessage from eval_protocol.models works.
class TestRLDataStructures:
def test_step_output_creation_valid(self):
"""Test valid creation of StepOutput."""
so = StepOutput(step_index=0, base_reward=0.5, reason="Good step", metrics={"accuracy": 0.9})
assert so.step_index == 0
assert so.base_reward == 0.5
assert so.reason == "Good step"
assert so.metrics == {"accuracy": 0.9}
so_str_index = StepOutput(step_index="turn_1", base_reward=-0.1)
assert so_str_index.step_index == "turn_1"
assert so_str_index.base_reward == -0.1
assert so_str_index.metrics == {}
assert so_str_index.reason is None
def test_step_output_invalid_types(self):
"""Test StepOutput validation errors for incorrect types."""
with pytest.raises(ValidationError):
StepOutput(step_index="0", base_reward="not_a_float") # base_reward should be float
with pytest.raises(ValidationError):
StepOutput(step_index=None, base_reward=0.5) # step_index is required
def test_evaluate_result_extended(self):
"""Test extended EvaluateResult with step_outputs."""
step_out1 = StepOutput(step_index=0, base_reward=0.1)
step_out2 = StepOutput(step_index="assistant_1", base_reward=0.2)
er_with_steps = EvaluateResult(
score=0.75,
reason="Overall good",
step_outputs=[step_out1, step_out2],
# metrics field is now Dict[str, MetricResult], not part of this basic test
# for simplicity, we'll test its default or assume it's handled elsewhere
)
assert er_with_steps.score == 0.75
assert er_with_steps.step_outputs is not None
assert len(er_with_steps.step_outputs) == 2
assert er_with_steps.step_outputs[0].base_reward == 0.1
assert er_with_steps.step_outputs[1].step_index == "assistant_1"
def test_evaluate_result_backward_compatibility(self):
"""Test EvaluateResult creation without new RL fields."""
# This test assumes MetricResult is defined and works as before.
# For simplicity, we might skip deep MetricResult testing here if it's complex to mock.
# Let's assume metrics can be an empty dict for this test if not focusing on MetricResult itself.
er_old_style = EvaluateResult(score=0.9, reason="Simple score")
assert er_old_style.score == 0.9
assert er_old_style.reason == "Simple score"
assert er_old_style.step_outputs is None
assert er_old_style.metrics == {} # Due to default_factory=dict
er_with_empty_steps = EvaluateResult(score=0.6, step_outputs=[])
assert er_with_empty_steps.score == 0.6
assert er_with_empty_steps.step_outputs == []
def test_evaluate_result_invalid_step_outputs(self):
"""Test EvaluateResult with invalid step_outputs type."""
with pytest.raises(ValidationError):
EvaluateResult(score=0.5, step_outputs="not_a_list")
with pytest.raises(ValidationError):
EvaluateResult(score=0.5, step_outputs=[{"step_index": 0, "base_reward": "wrong_type"}])
def test_step_data_creation_minimal(self):
"""Test minimal valid creation of StepData."""
msg_hist = [RewardKitMessage(role="user", content="Hello")]
action = {"type": "text", "content": "Hi"}
step = StepData(
system_step_index=0,
observation_data={"history": msg_hist}, # Example observation
action_taken=action,
resulting_messages_history=[
*msg_hist,
RewardKitMessage(role="assistant", content="Hi"),
],
)
assert step.system_step_index == 0
assert step.action_taken == action
assert step.base_reward is None
assert step.advantage is None
assert step.is_done is False
assert step.policy_value_estimate is None
def test_step_data_creation_full(self):
"""Test StepData creation with all optional fields."""
msg_hist1 = [RewardKitMessage(role="user", content="Hello")]
msg_hist2 = [*msg_hist1, RewardKitMessage(role="assistant", content="Hi there")]
step = StepData(
system_step_index=1,
observation_data=msg_hist1,
action_taken={"type": "text", "content": "Hi there"},
raw_policy_output="Hi there",
resulting_messages_history=msg_hist2,
policy_logprobs={"token_logprobs": [-0.1, -0.2]},
policy_value_estimate=0.95,
is_done=True,
step_info={"tool_used": "none", "latency_ms": 100},
base_reward=0.5,
advantage=0.1,
return_to_go=0.6,
)
assert step.policy_value_estimate == 0.95
assert step.is_done is True
assert step.base_reward == 0.5
assert step.advantage == 0.1
assert step.return_to_go == 0.6
assert step.step_info["latency_ms"] == 100
def test_step_data_field_validation(self):
"""Test StepData field type validations."""
with pytest.raises(ValidationError):
StepData( # Missing required fields
system_step_index=0,
observation_data=[],
)
with pytest.raises(ValidationError):
StepData(
system_step_index="not_an_int", # system_step_index should be int
observation_data=[],
action_taken={},
resulting_messages_history=[],
)
with pytest.raises(ValidationError):
StepData(
system_step_index=0,
observation_data=[],
action_taken={},
resulting_messages_history=[],
is_done="not_a_bool", # is_done should be bool
)
def test_step_data_message_import(self):
"""Test that Message can be used within StepData."""
# This test implicitly checks if the Message import in eval_protocol.agent.models
# is working or if the fallback is used. A more direct test might involve
# checking the type of resulting_messages_history items if possible.
m1 = RewardKitMessage(role="user", content="Test")
m2 = RewardKitMessage(role="assistant", content="Response")
step = StepData(
system_step_index=0,
observation_data=[m1],
action_taken={"type": "text", "content": "Response"},
resulting_messages_history=[m1, m2],
)
assert isinstance(step.resulting_messages_history[0], RewardKitMessage)