Coverage for pydantic_ai_slim/pydantic_ai/_result.py: 95.65%

1from __future__ import annotations as _annotations (empty)

2 

3import inspect (empty)

4import sys (empty)

5import types (empty)

6from collections.abc import Awaitable, Iterable (empty)

7from dataclasses import dataclass, field (empty)

8from typing import Any, Callable, Generic, Literal, Union, cast, get_args, get_origin (empty)

9 

10from pydantic import TypeAdapter, ValidationError (empty)

11from typing_extensions import TypeAliasType, TypedDict, TypeVar (empty)

12 

13from . import _utils, messages as _messages (empty)

14from .exceptions import ModelRetry (empty)

15from .result import ResultDataT, ResultDataT_inv, ResultValidatorFunc (empty)

16from .tools import AgentDepsT, RunContext, ToolDefinition (empty)

17 

18T = TypeVar('T') (empty)

19"""An invariant TypeVar.""" (empty)

20 

21 

22@dataclass (empty)

23class ResultValidator(Generic[AgentDepsT, ResultDataT_inv]): (empty)

24 function: ResultValidatorFunc[AgentDepsT, ResultDataT_inv] (empty)

25 _takes_ctx: bool = field(init=False) (empty)

26 _is_async: bool = field(init=False) (empty)

27 

28 def __post_init__(self): (empty)

29 self._takes_ctx = len(inspect.signature(self.function).parameters) > 1 (empty)

30 self._is_async = inspect.iscoroutinefunction(self.function) (empty)

31 

32 async def validate( (empty)

33 self, 

34 result: T, 

35 tool_call: _messages.ToolCallPart | None, 

36 run_context: RunContext[AgentDepsT], 

37 ) -> T: 

38 """Validate a result but calling the function. 

39 

40 Args: 

41 result: The result data after Pydantic validation the message content. 

42 tool_call: The original tool call message, `None` if there was no tool call. 

43 run_context: The current run context. 

44 

45 Returns: 

46 Result of either the validated result data (ok) or a retry message (Err). 

47 """ 

48 if self._takes_ctx: (empty)

49 ctx = run_context.replace_with(tool_name=tool_call.tool_name if tool_call else None) (empty)

50 args = ctx, result (empty)

51 else: 

52 args = (result,) (empty)

53 

54 try: (empty)

55 if self._is_async: (empty)

56 function = cast(Callable[[Any], Awaitable[T]], self.function) (empty)

57 result_data = await function(*args) (empty)

58 else: 

59 function = cast(Callable[[Any], T], self.function) (empty)

60 result_data = await _utils.run_in_executor(function, *args) (empty)

61 except ModelRetry as r: (empty)

62 m = _messages.RetryPromptPart(content=r.message) (empty)

63 if tool_call is not None: (empty)

64 m.tool_name = tool_call.tool_name (empty)

65 m.tool_call_id = tool_call.tool_call_id (empty)

66 raise ToolRetryError(m) from r (empty)

67 else: 

68 return result_data (empty)

69 

70 

71class ToolRetryError(Exception): (empty)

72 """Internal exception used to signal a `ToolRetry` message should be returned to the LLM.""" 

73 

74 def __init__(self, tool_retry: _messages.RetryPromptPart): (empty)

75 self.tool_retry = tool_retry (empty)

76 super().__init__() (empty)

77 

78 

79@dataclass (empty)

80class ResultSchema(Generic[ResultDataT]): (empty)

81 """Model the final response from an agent run. 

82 

83 Similar to `Tool` but for the final result of running an agent. 

84 """ 

85 

86 tools: dict[str, ResultTool[ResultDataT]] (empty)

87 allow_text_result: bool (empty)

88 

89 @classmethod (empty)

90 def build( (empty)

91 cls: type[ResultSchema[T]], response_type: type[T], name: str, description: str | None 

92 ) -> ResultSchema[T] | None: 

93 """Build a ResultSchema dataclass from a response type.""" 

94 if response_type is str: (empty)

95 return None (empty)

96 

97 if response_type_option := extract_str_from_union(response_type): (empty)

98 response_type = response_type_option.value (empty)

99 allow_text_result = True (empty)

100 else: 

101 allow_text_result = False (empty)

102 

103 def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[T]: (empty)

104 return cast(ResultTool[T], ResultTool(a, tool_name_, description, multiple)) (empty)

105 

106 tools: dict[str, ResultTool[T]] = {} (empty)

107 if args := get_union_args(response_type): (empty)

108 for i, arg in enumerate(args, start=1): (empty)

109 tool_name = union_tool_name(name, arg) (empty)

110 while tool_name in tools: (empty)

111 tool_name = f'{tool_name}_{i}' (empty)

112 tools[tool_name] = _build_tool(arg, tool_name, True) (empty)

113 else: 

114 tools[name] = _build_tool(response_type, name, False) (empty)

115 

116 return cls(tools=tools, allow_text_result=allow_text_result) (empty)

117 

118 def find_named_tool( (empty)

119 self, parts: Iterable[_messages.ModelResponsePart], tool_name: str 

120 ) -> tuple[_messages.ToolCallPart, ResultTool[ResultDataT]] | None: 

121 """Find a tool that matches one of the calls, with a specific name.""" 

122 for part in parts: 122 ↛ exitline 122 didn't return from function 'find_named_tool' because the loop on line 122 didn't complete(empty)

123 if isinstance(part, _messages.ToolCallPart): 123 ↛ 122line 123 didn't jump to line 122 because the condition on line 123 was always true(empty)

124 if part.tool_name == tool_name: (empty)

125 return part, self.tools[tool_name] (empty)

126 

127 def find_tool( (empty)

128 self, 

129 parts: Iterable[_messages.ModelResponsePart], 

130 ) -> tuple[_messages.ToolCallPart, ResultTool[ResultDataT]] | None: 

131 """Find a tool that matches one of the calls.""" 

132 for part in parts: (empty)

133 if isinstance(part, _messages.ToolCallPart): 133 ↛ 132line 133 didn't jump to line 132 because the condition on line 133 was always true(empty)

134 if result := self.tools.get(part.tool_name): (empty)

135 return part, result (empty)

136 

137 def tool_names(self) -> list[str]: (empty)

138 """Return the names of the tools.""" 

139 return list(self.tools.keys()) (empty)

140 

141 def tool_defs(self) -> list[ToolDefinition]: (empty)

142 """Get tool definitions to register with the model.""" 

143 return [t.tool_def for t in self.tools.values()] (empty)

144 

145 

146DEFAULT_DESCRIPTION = 'The final response which ends this conversation' (empty)

147 

148 

149@dataclass(init=False) (empty)

150class ResultTool(Generic[ResultDataT]): (empty)

151 tool_def: ToolDefinition (empty)

152 type_adapter: TypeAdapter[Any] (empty)

153 

154 def __init__(self, response_type: type[ResultDataT], name: str, description: str | None, multiple: bool): (empty)

155 """Build a ResultTool dataclass from a response type.""" 

156 assert response_type is not str, 'ResultTool does not support str as a response type' (empty)

157 

158 if _utils.is_model_like(response_type): (empty)

159 self.type_adapter = TypeAdapter(response_type) (empty)

160 outer_typed_dict_key: str | None = None (empty)

161 # noinspection PyArgumentList 

162 parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema()) (empty)

163 else: 

164 response_data_typed_dict = TypedDict('response_data_typed_dict', {'response': response_type}) # noqa (empty)

165 self.type_adapter = TypeAdapter(response_data_typed_dict) (empty)

166 outer_typed_dict_key = 'response' (empty)

167 # noinspection PyArgumentList 

168 parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema()) (empty)

169 # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM 

170 parameters_json_schema.pop('title') (empty)

171 

172 if json_schema_description := parameters_json_schema.pop('description', None): (empty)

173 if description is None: 173 ↛ 176line 173 didn't jump to line 176 because the condition on line 173 was always true(empty)

174 tool_description = json_schema_description (empty)

175 else: 

176 tool_description = f'{description}. {json_schema_description}' 

177 else: 

178 tool_description = description or DEFAULT_DESCRIPTION (empty)

179 if multiple: (empty)

180 tool_description = f'{union_arg_name(response_type)}: {tool_description}' (empty)

181 

182 self.tool_def = ToolDefinition( (empty)

183 name=name, 

184 description=tool_description, 

185 parameters_json_schema=parameters_json_schema, 

186 outer_typed_dict_key=outer_typed_dict_key, 

187 ) 

188 

189 def validate( (empty)

190 self, tool_call: _messages.ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True 

191 ) -> ResultDataT: 

192 """Validate a result message. 

193 

194 Args: 

195 tool_call: The tool call from the LLM to validate. 

196 allow_partial: If true, allow partial validation. 

197 wrap_validation_errors: If true, wrap the validation errors in a retry message. 

198 

199 Returns: 

200 Either the validated result data (left) or a retry message (right). 

201 """ 

202 try: (empty)

203 pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off' (empty)

204 if isinstance(tool_call.args, str): (empty)

205 result = self.type_adapter.validate_json(tool_call.args, experimental_allow_partial=pyd_allow_partial) (empty)

206 else: 

207 result = self.type_adapter.validate_python(tool_call.args, experimental_allow_partial=pyd_allow_partial) (empty)

208 except ValidationError as e: (empty)

209 if wrap_validation_errors: 209 ↛ 217line 209 didn't jump to line 217 because the condition on line 209 was always true(empty)

210 m = _messages.RetryPromptPart( (empty)

211 tool_name=tool_call.tool_name, 

212 content=e.errors(include_url=False), 

213 tool_call_id=tool_call.tool_call_id, 

214 ) 

215 raise ToolRetryError(m) from e (empty)

216 else: 

217 raise (empty)

218 else: 

219 if k := self.tool_def.outer_typed_dict_key: (empty)

220 result = result[k] (empty)

221 return result (empty)

222 

223 

224def union_tool_name(base_name: str, union_arg: Any) -> str: (empty)

225 return f'{base_name}_{union_arg_name(union_arg)}' (empty)

226 

227 

228def union_arg_name(union_arg: Any) -> str: (empty)

229 return union_arg.__name__ (empty)

230 

231 

232def extract_str_from_union(response_type: Any) -> _utils.Option[Any]: (empty)

233 """Extract the string type from a Union, return the remaining union or remaining type.""" 

234 union_args = get_union_args(response_type) (empty)

235 if any(t is str for t in union_args): (empty)

236 remain_args: list[Any] = [] (empty)

237 includes_str = False (empty)

238 for arg in union_args: (empty)

239 if arg is str: (empty)

240 includes_str = True (empty)

241 else: 

242 remain_args.append(arg) (empty)

243 if includes_str: 243 ↛ exitline 243 didn't return from function 'extract_str_from_union' because the condition on line 243 was always true(empty)

244 if len(remain_args) == 1: 244 ↛ 247line 244 didn't jump to line 247 because the condition on line 244 was always true(empty)

245 return _utils.Some(remain_args[0]) (empty)

246 else: 

247 return _utils.Some(Union[tuple(remain_args)]) 

248 

249 

250def get_union_args(tp: Any) -> tuple[Any, ...]: (empty)

251 """Extract the arguments of a Union type if `response_type` is a union, otherwise return an empty union.""" 

252 if isinstance(tp, TypeAliasType): (empty)

253 tp = tp.__value__ (empty)

254 

255 origin = get_origin(tp) (empty)

256 if origin_is_union(origin): (empty)

257 return get_args(tp) (empty)

258 else: 

259 return () (empty)

260 

261 

262if sys.version_info < (3, 10): (empty)

263 

264 def origin_is_union(tp: type[Any] | None) -> bool: (empty)

265 return tp is Union (empty)

266 

267else: 

268 

269 def origin_is_union(tp: type[Any] | None) -> bool: (empty)

270 return tp is Union or tp is types.UnionType (empty)