forked from strands-agents/harness-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_tool.py
More file actions
162 lines (131 loc) · 4.04 KB
/
Copy pathtest_tool.py
File metadata and controls
162 lines (131 loc) · 4.04 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
import json
from unittest.mock import ANY
import pytest
from strands import Agent, tool
from strands.hooks import BeforeToolCallEvent, HookProvider
from strands.interrupt import Interrupt
from strands.types.tools import ToolContext
@pytest.fixture
def interrupt_hook():
class Hook(HookProvider):
def register_hooks(self, registry):
registry.add_callback(BeforeToolCallEvent, self.interrupt)
def interrupt(self, event):
if event.tool_use["name"] != "time_tool":
return
response = event.interrupt("test_interrupt", reason="need approval")
if response != "APPROVE":
event.cancel_tool = "tool rejected"
return Hook()
@pytest.fixture
def time_tool():
@tool(name="time_tool", context=True)
def func(tool_context: ToolContext) -> str:
return tool_context.interrupt("test_interrupt", reason="need time")
return func
@pytest.fixture
def day_tool():
@tool(name="day_tool", context=True)
def func(tool_context: ToolContext) -> str:
return tool_context.interrupt("test_interrupt", reason="need day")
return func
@pytest.fixture
def weather_tool():
@tool(name="weather_tool")
def func() -> str:
return "sunny"
return func
@pytest.fixture
def agent(interrupt_hook, time_tool, day_tool, weather_tool):
return Agent(hooks=[interrupt_hook], tools=[time_tool, day_tool, weather_tool])
def test_interrupt(agent):
result = agent("What is the time, day, and weather?")
tru_stop_reason = result.stop_reason
exp_stop_reason = "interrupt"
assert tru_stop_reason == exp_stop_reason
tru_interrupts = sorted(result.interrupts, key=lambda interrupt: interrupt.reason)
exp_interrupts = [
Interrupt(
id=ANY,
name="test_interrupt",
reason="need approval",
),
Interrupt(
id=ANY,
name="test_interrupt",
reason="need day",
),
]
assert tru_interrupts == exp_interrupts
interrupt_approval, interrupt_day = result.interrupts
responses = [
{
"interruptResponse": {
"interruptId": interrupt_approval.id,
"response": "APPROVE",
},
},
{
"interruptResponse": {
"interruptId": interrupt_day.id,
"response": "monday",
},
},
]
result = agent(responses)
tru_stop_reason = result.stop_reason
exp_stop_reason = "interrupt"
assert tru_stop_reason == exp_stop_reason
tru_interrupts = result.interrupts
exp_interrupts = [
Interrupt(
id=ANY,
name="test_interrupt",
reason="need time",
),
]
assert tru_interrupts == exp_interrupts
interrupt_time = result.interrupts[0]
responses = [
{
"interruptResponse": {
"interruptId": interrupt_time.id,
"response": "12:01",
},
},
]
result = agent(responses)
result_message = json.dumps(result.message).lower()
assert all(string in result_message for string in ["12:01", "monday", "sunny"])
tru_tool_results = agent.messages[-2]["content"]
tru_tool_results.sort(key=lambda content: content["toolResult"]["content"][0]["text"])
exp_tool_results = [
{
"toolResult": {
"toolUseId": ANY,
"status": "success",
"content": [
{"text": "12:01"},
],
},
},
{
"toolResult": {
"toolUseId": ANY,
"status": "success",
"content": [
{"text": "monday"},
],
},
},
{
"toolResult": {
"toolUseId": ANY,
"status": "success",
"content": [
{"text": "sunny"},
],
},
},
]
assert tru_tool_results == exp_tool_results