forked from eval-protocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_langgraph_processor.py
More file actions
244 lines (206 loc) · 7.88 KB
/
Copy pathtest_langgraph_processor.py
File metadata and controls
244 lines (206 loc) · 7.88 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
from __future__ import annotations
import asyncio
from typing import Any, Dict, List
import pytest
from eval_protocol.models import EvaluationRow, Message
from eval_protocol.pytest.default_langchain_rollout_processor import LangGraphRolloutProcessor
from eval_protocol.pytest.types import RolloutProcessorConfig
class DummyLCMessage:
def __init__(self, message_type: str, content: str): # noqa: A002
self.type = message_type
self.content = content
class DummyGraph:
def __init__(self, out_messages: List[Any]):
self._out_messages = out_messages
async def ainvoke(self, payload: Dict[str, Any], **_: Any):
# Echo back the provided messages plus our out_messages
return {"messages": list(payload.get("messages") or []) + list(self._out_messages)}
def _make_processor_with_defaults(out_messages: List[Any]) -> LangGraphRolloutProcessor:
def graph_factory(_: RolloutProcessorConfig):
return DummyGraph(out_messages)
return LangGraphRolloutProcessor(graph_factory=graph_factory)
@pytest.mark.asyncio
async def test_apply_result_preserves_user_role_and_appends_assistant_from_lc():
# Arrange: EP user message in, LC assistant out
row = EvaluationRow(messages=[Message(role="user", content="hi")])
lc_assistant = DummyLCMessage(message_type="ai", content="hello")
processor = _make_processor_with_defaults([lc_assistant])
# Act
tasks = processor(
[row],
type(
"Cfg",
(),
{
"completion_params": {},
"semaphore": asyncio.Semaphore(10),
"mcp_config_path": "",
"logger": None,
"server_script_path": None,
"steps": 1,
"kwargs": {},
"exception_handler_config": None,
},
)(),
)
result_row = await asyncio.gather(*tasks)
out = result_row[0]
# Assert
assert out.messages[0].role == "user"
assert out.messages[-1].role == "assistant"
assert out.messages[-1].content == "hello"
@pytest.mark.asyncio
async def test_apply_result_handles_dict_messages_with_missing_role():
row = EvaluationRow(messages=[Message(role="user", content="Q")])
dict_msg = {"content": "A"} # no role provided
processor = _make_processor_with_defaults([dict_msg])
tasks = processor(
[row],
type(
"Cfg",
(),
{
"completion_params": {},
"semaphore": asyncio.Semaphore(10),
"mcp_config_path": "",
"logger": None,
"server_script_path": None,
"steps": 1,
"kwargs": {},
"exception_handler_config": None,
},
)(),
)
out = (await asyncio.gather(*tasks))[0]
assert out.messages[0].role == "user"
assert out.messages[-1].role == "assistant"
assert out.messages[-1].content == "A"
@pytest.mark.asyncio
async def test_to_input_converts_ep_messages_to_lc_via_adapter(monkeypatch):
# Arrange
ep_row = EvaluationRow(messages=[Message(role="user", content="Hello")])
called = {"ok": False}
def fake_to_lc(messages):
called["ok"] = True
return [DummyLCMessage(message_type="human", content=messages[0].content)]
# Patch the adapter function at its source module, since the processor imports it inside the function
import eval_protocol.adapters.langchain as lc_adapter
monkeypatch.setattr(lc_adapter, "serialize_ep_messages_to_lc", fake_to_lc, raising=True)
# Dummy graph that returns what it receives
class EchoGraph:
async def ainvoke(self, payload, **_):
# Ensure our adapter-produced messages flow through
return payload
processor = LangGraphRolloutProcessor(graph_factory=lambda _config: EchoGraph())
# Act
tasks = processor(
[ep_row],
type(
"Cfg",
(),
{
"completion_params": {},
"semaphore": asyncio.Semaphore(10),
"mcp_config_path": "",
"logger": None,
"server_script_path": None,
"steps": 1,
"kwargs": {},
"exception_handler_config": None,
},
)(),
)
await asyncio.gather(*tasks)
# Assert that adapter was used
assert called["ok"] is True
@pytest.mark.asyncio
async def test_apply_result_copies_tool_calls_from_lc_ai_and_toolmessage():
from langchain_core.messages import AIMessage, ToolMessage
# Arrange: EP user message in, LC assistant with tool_calls + LC tool message out
row = EvaluationRow(messages=[Message(role="user", content="count tracks")])
tool_call_id = "call_1"
# Use LangChain-native tool_call schema (name/args) so AIMessage validates
ai_with_tool = AIMessage(
content="I'll call the tool.",
tool_calls=[
{
"id": tool_call_id,
"name": "count_tracks",
"args": {},
}
],
)
tool_msg = ToolMessage(content="3503", name="count_tracks", tool_call_id=tool_call_id, status="success")
processor = _make_processor_with_defaults([ai_with_tool, tool_msg])
# Act
tasks = processor(
[row],
type(
"Cfg",
(),
{
"completion_params": {},
"semaphore": asyncio.Semaphore(10),
"mcp_config_path": "",
"logger": None,
"server_script_path": None,
"steps": 1,
"kwargs": {},
"exception_handler_config": None,
},
)(),
)
out = (await asyncio.gather(*tasks))[0]
# Assert: assistant message has tool_calls, and tool message is present
assistants = [m for m in out.messages if m.role == "assistant"]
tools = [m for m in out.messages if m.role == "tool"]
assert assistants, "No assistant messages found"
assert tools, "No tool messages found"
assert assistants[-1].tool_calls is not None and len(assistants[-1].tool_calls) == 1
assert assistants[-1].tool_calls[0].id, "tool_call id should be present"
assert tools[-1].content and "3503" in (tools[-1].content or "")
@pytest.mark.asyncio
async def test_apply_result_copies_tool_calls_from_additional_kwargs():
from langchain_core.messages import AIMessage, ToolMessage
# Arrange: tool_calls provided via additional_kwargs instead of attribute
row = EvaluationRow(messages=[Message(role="user", content="count tracks")])
tool_call_id = "call_2"
ai_with_tool = AIMessage(
content="I'll call the tool.",
additional_kwargs={
"tool_calls": [
{
"id": tool_call_id,
"name": "count_tracks",
"args": {},
}
]
},
)
tool_msg = ToolMessage(content="3503", name="count_tracks", tool_call_id=tool_call_id, status="success")
processor = _make_processor_with_defaults([ai_with_tool, tool_msg])
# Act
tasks = processor(
[row],
type(
"Cfg",
(),
{
"completion_params": {},
"semaphore": asyncio.Semaphore(10),
"mcp_config_path": "",
"logger": None,
"server_script_path": None,
"steps": 1,
"kwargs": {},
"exception_handler_config": None,
},
)(),
)
out = (await asyncio.gather(*tasks))[0]
# Assert
assistants = [m for m in out.messages if m.role == "assistant"]
tools = [m for m in out.messages if m.role == "tool"]
assert assistants and assistants[-1].tool_calls is not None
assert any(tc.id for tc in assistants[-1].tool_calls), "Expected tool_call with id"
assert tools and "3503" in (tools[-1].content or "")