23
23
from pydantic import BaseModel
24
24
import pytest
25
25
26
- from .. import utils
26
+ from .. import testing_utils
27
27
28
28
29
29
class MockBeforeModelCallback (BaseModel ):
@@ -35,7 +35,7 @@ def __call__(
35
35
llm_request : LlmRequest ,
36
36
) -> LlmResponse :
37
37
return LlmResponse (
38
- content = utils .ModelContent (
38
+ content = testing_utils .ModelContent (
39
39
[types .Part .from_text (text = self .mock_response )]
40
40
)
41
41
)
@@ -50,7 +50,7 @@ def __call__(
50
50
llm_response : LlmResponse ,
51
51
) -> LlmResponse :
52
52
return LlmResponse (
53
- content = utils .ModelContent (
53
+ content = testing_utils .ModelContent (
54
54
[types .Part .from_text (text = self .mock_response )]
55
55
)
56
56
)
@@ -65,7 +65,7 @@ async def __call__(
65
65
llm_request : LlmRequest ,
66
66
) -> LlmResponse :
67
67
return LlmResponse (
68
- content = utils .ModelContent (
68
+ content = testing_utils .ModelContent (
69
69
[types .Part .from_text (text = self .mock_response )]
70
70
)
71
71
)
@@ -80,7 +80,7 @@ async def __call__(
80
80
llm_response : LlmResponse ,
81
81
) -> LlmResponse :
82
82
return LlmResponse (
83
- content = utils .ModelContent (
83
+ content = testing_utils .ModelContent (
84
84
[types .Part .from_text (text = self .mock_response )]
85
85
)
86
86
)
@@ -97,7 +97,7 @@ async def async_noop_callback(**kwargs) -> Optional[LlmResponse]:
97
97
@pytest .mark .asyncio
98
98
async def test_before_model_callback ():
99
99
responses = ['model_response' ]
100
- mock_model = utils .MockModel .create (responses = responses )
100
+ mock_model = testing_utils .MockModel .create (responses = responses )
101
101
agent = Agent (
102
102
name = 'root_agent' ,
103
103
model = mock_model ,
@@ -106,8 +106,8 @@ async def test_before_model_callback():
106
106
),
107
107
)
108
108
109
- runner = utils .TestInMemoryRunner (agent )
110
- assert utils .simplify_events (
109
+ runner = testing_utils .TestInMemoryRunner (agent )
110
+ assert testing_utils .simplify_events (
111
111
await runner .run_async_with_new_session ('test' )
112
112
) == [
113
113
('root_agent' , 'before_model_callback' ),
@@ -117,15 +117,15 @@ async def test_before_model_callback():
117
117
@pytest .mark .asyncio
118
118
async def test_before_model_callback_noop ():
119
119
responses = ['model_response' ]
120
- mock_model = utils .MockModel .create (responses = responses )
120
+ mock_model = testing_utils .MockModel .create (responses = responses )
121
121
agent = Agent (
122
122
name = 'root_agent' ,
123
123
model = mock_model ,
124
124
before_model_callback = noop_callback ,
125
125
)
126
126
127
- runner = utils .TestInMemoryRunner (agent )
128
- assert utils .simplify_events (
127
+ runner = testing_utils .TestInMemoryRunner (agent )
128
+ assert testing_utils .simplify_events (
129
129
await runner .run_async_with_new_session ('test' )
130
130
) == [
131
131
('root_agent' , 'model_response' ),
@@ -135,7 +135,7 @@ async def test_before_model_callback_noop():
135
135
@pytest .mark .asyncio
136
136
async def test_after_model_callback ():
137
137
responses = ['model_response' ]
138
- mock_model = utils .MockModel .create (responses = responses )
138
+ mock_model = testing_utils .MockModel .create (responses = responses )
139
139
agent = Agent (
140
140
name = 'root_agent' ,
141
141
model = mock_model ,
@@ -144,8 +144,8 @@ async def test_after_model_callback():
144
144
),
145
145
)
146
146
147
- runner = utils .TestInMemoryRunner (agent )
148
- assert utils .simplify_events (
147
+ runner = testing_utils .TestInMemoryRunner (agent )
148
+ assert testing_utils .simplify_events (
149
149
await runner .run_async_with_new_session ('test' )
150
150
) == [
151
151
('root_agent' , 'after_model_callback' ),
@@ -155,7 +155,7 @@ async def test_after_model_callback():
155
155
@pytest .mark .asyncio
156
156
async def test_async_before_model_callback ():
157
157
responses = ['model_response' ]
158
- mock_model = utils .MockModel .create (responses = responses )
158
+ mock_model = testing_utils .MockModel .create (responses = responses )
159
159
agent = Agent (
160
160
name = 'root_agent' ,
161
161
model = mock_model ,
@@ -164,8 +164,8 @@ async def test_async_before_model_callback():
164
164
),
165
165
)
166
166
167
- runner = utils .TestInMemoryRunner (agent )
168
- assert utils .simplify_events (
167
+ runner = testing_utils .TestInMemoryRunner (agent )
168
+ assert testing_utils .simplify_events (
169
169
await runner .run_async_with_new_session ('test' )
170
170
) == [
171
171
('root_agent' , 'async_before_model_callback' ),
@@ -175,15 +175,15 @@ async def test_async_before_model_callback():
175
175
@pytest .mark .asyncio
176
176
async def test_async_before_model_callback_noop ():
177
177
responses = ['model_response' ]
178
- mock_model = utils .MockModel .create (responses = responses )
178
+ mock_model = testing_utils .MockModel .create (responses = responses )
179
179
agent = Agent (
180
180
name = 'root_agent' ,
181
181
model = mock_model ,
182
182
before_model_callback = async_noop_callback ,
183
183
)
184
184
185
- runner = utils .TestInMemoryRunner (agent )
186
- assert utils .simplify_events (
185
+ runner = testing_utils .TestInMemoryRunner (agent )
186
+ assert testing_utils .simplify_events (
187
187
await runner .run_async_with_new_session ('test' )
188
188
) == [
189
189
('root_agent' , 'model_response' ),
@@ -193,7 +193,7 @@ async def test_async_before_model_callback_noop():
193
193
@pytest .mark .asyncio
194
194
async def test_async_after_model_callback ():
195
195
responses = ['model_response' ]
196
- mock_model = utils .MockModel .create (responses = responses )
196
+ mock_model = testing_utils .MockModel .create (responses = responses )
197
197
agent = Agent (
198
198
name = 'root_agent' ,
199
199
model = mock_model ,
@@ -202,8 +202,8 @@ async def test_async_after_model_callback():
202
202
),
203
203
)
204
204
205
- runner = utils .TestInMemoryRunner (agent )
206
- assert utils .simplify_events (
205
+ runner = testing_utils .TestInMemoryRunner (agent )
206
+ assert testing_utils .simplify_events (
207
207
await runner .run_async_with_new_session ('test' )
208
208
) == [
209
209
('root_agent' , 'async_after_model_callback' ),
0 commit comments