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 | @dataclass
class TestModel(Model):
"""A model specifically for testing purposes.
This will (by default) call all tools in the agent, then return a tool response if possible,
otherwise a plain response.
How useful this model is will vary significantly.
Apart from `__init__` derived by the `dataclass` decorator, all methods are private or match those
of the base class.
"""
# NOTE: Avoid test discovery by pytest.
__test__ = False
call_tools: list[str] | Literal['all'] = 'all'
"""List of tools to call. If `'all'`, all tools will be called."""
custom_result_text: str | None = None
"""If set, this text is returned as the final result."""
custom_result_args: Any | None = None
"""If set, these args will be passed to the result tool."""
seed: int = 0
"""Seed for generating random data."""
last_model_request_parameters: ModelRequestParameters | None = field(default=None, init=False)
"""The last ModelRequestParameters passed to the model in a request.
The ModelRequestParameters contains information about the function and result tools available during request handling.
This is set when a request is made, so will reflect the function tools from the last step of the last run.
"""
_model_name: str = field(default='test', repr=False)
_system: str = field(default='test', repr=False)
async def request(
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
) -> tuple[ModelResponse, Usage]:
self.last_model_request_parameters = model_request_parameters
model_response = self._request(messages, model_settings, model_request_parameters)
usage = _estimate_usage([*messages, model_response])
return model_response, usage
@asynccontextmanager
async def request_stream(
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
) -> AsyncIterator[StreamedResponse]:
self.last_model_request_parameters = model_request_parameters
model_response = self._request(messages, model_settings, model_request_parameters)
yield TestStreamedResponse(
_model_name=self._model_name, _structured_response=model_response, _messages=messages
)
@property
def model_name(self) -> str:
"""The model name."""
return self._model_name
@property
def system(self) -> str:
"""The system / model provider."""
return self._system
def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
def _get_tool_calls(self, model_request_parameters: ModelRequestParameters) -> list[tuple[str, ToolDefinition]]:
if self.call_tools == 'all':
return [(r.name, r) for r in model_request_parameters.function_tools]
else:
function_tools_lookup = {t.name: t for t in model_request_parameters.function_tools}
tools_to_call = (function_tools_lookup[name] for name in self.call_tools)
return [(r.name, r) for r in tools_to_call]
def _get_result(self, model_request_parameters: ModelRequestParameters) -> _TextResult | _FunctionToolResult:
if self.custom_result_text is not None:
assert model_request_parameters.allow_text_result, (
'Plain response not allowed, but `custom_result_text` is set.'
)
assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
return _TextResult(self.custom_result_text)
elif self.custom_result_args is not None:
assert model_request_parameters.result_tools is not None, (
'No result tools provided, but `custom_result_args` is set.'
)
result_tool = model_request_parameters.result_tools[0]
if k := result_tool.outer_typed_dict_key:
return _FunctionToolResult({k: self.custom_result_args})
else:
return _FunctionToolResult(self.custom_result_args)
elif model_request_parameters.allow_text_result:
return _TextResult(None)
elif model_request_parameters.result_tools:
return _FunctionToolResult(None)
else:
return _TextResult(None)
def _request(
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
) -> ModelResponse:
tool_calls = self._get_tool_calls(model_request_parameters)
result = self._get_result(model_request_parameters)
result_tools = model_request_parameters.result_tools
# if there are tools, the first thing we want to do is call all of them
if tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
return ModelResponse(
parts=[ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls],
model_name=self._model_name,
)
if messages:
last_message = messages[-1]
assert isinstance(last_message, ModelRequest), 'Expected last message to be a `ModelRequest`.'
# check if there are any retry prompts, if so retry them
new_retry_names = {p.tool_name for p in last_message.parts if isinstance(p, RetryPromptPart)}
if new_retry_names:
# Handle retries for both function tools and result tools
# Check function tools first
retry_parts: list[ModelResponsePart] = [
ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls if name in new_retry_names
]
# Check result tools
if result_tools:
retry_parts.extend(
[
ToolCallPart(
tool.name,
result.value
if isinstance(result, _FunctionToolResult) and result.value is not None
else self.gen_tool_args(tool),
)
for tool in result_tools
if tool.name in new_retry_names
]
)
return ModelResponse(parts=retry_parts, model_name=self._model_name)
if isinstance(result, _TextResult):
if (response_text := result.value) is None:
# build up details of tool responses
output: dict[str, Any] = {}
for message in messages:
if isinstance(message, ModelRequest):
for part in message.parts:
if isinstance(part, ToolReturnPart):
output[part.tool_name] = part.content
if output:
return ModelResponse(
parts=[TextPart(pydantic_core.to_json(output).decode())], model_name=self._model_name
)
else:
return ModelResponse(parts=[TextPart('success (no tool calls)')], model_name=self._model_name)
else:
return ModelResponse(parts=[TextPart(response_text)], model_name=self._model_name)
else:
assert result_tools, 'No result tools provided'
custom_result_args = result.value
result_tool = result_tools[self.seed % len(result_tools)]
if custom_result_args is not None:
return ModelResponse(
parts=[ToolCallPart(result_tool.name, custom_result_args)], model_name=self._model_name
)
else:
response_args = self.gen_tool_args(result_tool)
return ModelResponse(parts=[ToolCallPart(result_tool.name, response_args)], model_name=self._model_name)
|