Coverage for pydantic_ai_slim/pydantic_ai/_agent_graph.py: 98.57%

1from __future__ import annotations as _annotations (empty)

2 

3import asyncio (empty)

4import dataclasses (empty)

5from abc import ABC (empty)

6from collections.abc import AsyncIterator, Iterator, Sequence (empty)

7from contextlib import asynccontextmanager, contextmanager (empty)

8from contextvars import ContextVar (empty)

9from dataclasses import field (empty)

10from typing import Any, Generic, Literal, Union, cast (empty)

11 

12import logfire_api (empty)

13from typing_extensions import TypeVar, assert_never (empty)

14 

15from pydantic_graph import BaseNode, Graph, GraphRunContext (empty)

16from pydantic_graph.nodes import End, NodeRunEndT (empty)

17 

18from . import ( (empty)

19 _result, 

20 _system_prompt, 

21 exceptions, 

22 messages as _messages, 

23 models, 

24 result, 

25 usage as _usage, 

26) 

27from .result import ResultDataT (empty)

28from .settings import ModelSettings, merge_model_settings (empty)

29from .tools import ( (empty)

30 RunContext, 

31 Tool, 

32 ToolDefinition, 

33) 

34 

35_logfire = logfire_api.Logfire(otel_scope='pydantic-ai') (empty)

36 

37# while waiting for https://github.com/pydantic/logfire/issues/745 

38try: (empty)

39 import logfire._internal.stack_info (empty)

40except ImportError: (empty)

41 pass (empty)

42else: 

43 from pathlib import Path (empty)

44 

45 logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),) (empty)

46 

47T = TypeVar('T') (empty)

48NoneType = type(None) (empty)

49EndStrategy = Literal['early', 'exhaustive'] (empty)

50"""The strategy for handling multiple tool calls when a final result is found. (empty)

51 

52- `'early'`: Stop processing other tool calls once a final result is found 

53- `'exhaustive'`: Process all tool calls even after finding a final result 

54""" 

55DepsT = TypeVar('DepsT') (empty)

56ResultT = TypeVar('ResultT') (empty)

57 

58 

59@dataclasses.dataclass (empty)

60class MarkFinalResult(Generic[ResultDataT]): (empty)

61 """Marker class to indicate that the result is the final result. 

62 

63 This allows us to use `isinstance`, which wouldn't be possible if we were returning `ResultDataT` directly. 

64 

65 It also avoids problems in the case where the result type is itself `None`, but is set. 

66 """ 

67 

68 data: ResultDataT (empty)

69 """The final result data.""" (empty)

70 tool_name: str | None (empty)

71 """Name of the final result tool, None if the result is a string.""" (empty)

72 

73 

74@dataclasses.dataclass (empty)

75class GraphAgentState: (empty)

76 """State kept across the execution of the agent graph.""" 

77 

78 message_history: list[_messages.ModelMessage] (empty)

79 usage: _usage.Usage (empty)

80 retries: int (empty)

81 run_step: int (empty)

82 

83 def increment_retries(self, max_result_retries: int) -> None: (empty)

84 self.retries += 1 (empty)

85 if self.retries > max_result_retries: (empty)

86 raise exceptions.UnexpectedModelBehavior( (empty)

87 f'Exceeded maximum retries ({max_result_retries}) for result validation' 

88 ) 

89 

90 

91@dataclasses.dataclass (empty)

92class GraphAgentDeps(Generic[DepsT, ResultDataT]): (empty)

93 """Dependencies/config passed to the agent graph.""" 

94 

95 user_deps: DepsT (empty)

96 

97 prompt: str (empty)

98 new_message_index: int (empty)

99 

100 model: models.Model (empty)

101 model_settings: ModelSettings | None (empty)

102 usage_limits: _usage.UsageLimits (empty)

103 max_result_retries: int (empty)

104 end_strategy: EndStrategy (empty)

105 

106 result_schema: _result.ResultSchema[ResultDataT] | None (empty)

107 result_tools: list[ToolDefinition] (empty)

108 result_validators: list[_result.ResultValidator[DepsT, ResultDataT]] (empty)

109 

110 function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False) (empty)

111 

112 run_span: logfire_api.LogfireSpan (empty)

113 

114 

115@dataclasses.dataclass (empty)

116class BaseUserPromptNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], NodeRunEndT], ABC): (empty)

117 user_prompt: str (empty)

118 

119 system_prompts: tuple[str, ...] (empty)

120 system_prompt_functions: list[_system_prompt.SystemPromptRunner[DepsT]] (empty)

121 system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[DepsT]] (empty)

122 

123 async def _get_first_message( (empty)

124 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]] 

125 ) -> _messages.ModelRequest: 

126 run_context = _build_run_context(ctx) (empty)

127 history, next_message = await self._prepare_messages(self.user_prompt, ctx.state.message_history, run_context) (empty)

128 ctx.state.message_history = history (empty)

129 run_context.messages = history (empty)

130 

131 # TODO: We need to make it so that function_tools are not shared between runs 

132 # See comment on the current_retry field of `Tool` for more details. 

133 for tool in ctx.deps.function_tools.values(): (empty)

134 tool.current_retry = 0 (empty)

135 return next_message (empty)

136 

137 async def _prepare_messages( (empty)

138 self, user_prompt: str, message_history: list[_messages.ModelMessage] | None, run_context: RunContext[DepsT] 

139 ) -> tuple[list[_messages.ModelMessage], _messages.ModelRequest]: 

140 try: (empty)

141 ctx_messages = get_captured_run_messages() (empty)

142 except LookupError: (empty)

143 messages: list[_messages.ModelMessage] = [] (empty)

144 else: 

145 if ctx_messages.used: (empty)

146 messages = [] (empty)

147 else: 

148 messages = ctx_messages.messages (empty)

149 ctx_messages.used = True (empty)

150 

151 if message_history: (empty)

152 # Shallow copy messages 

153 messages.extend(message_history) (empty)

154 # Reevaluate any dynamic system prompt parts 

155 await self._reevaluate_dynamic_prompts(messages, run_context) (empty)

156 return messages, _messages.ModelRequest([_messages.UserPromptPart(user_prompt)]) (empty)

157 else: 

158 parts = await self._sys_parts(run_context) (empty)

159 parts.append(_messages.UserPromptPart(user_prompt)) (empty)

160 return messages, _messages.ModelRequest(parts) (empty)

161 

162 async def _reevaluate_dynamic_prompts( (empty)

163 self, messages: list[_messages.ModelMessage], run_context: RunContext[DepsT] 

164 ) -> None: 

165 """Reevaluate any `SystemPromptPart` with dynamic_ref in the provided messages by running the associated runner function.""" 

166 # Only proceed if there's at least one dynamic runner. 

167 if self.system_prompt_dynamic_functions: (empty)

168 for msg in messages: (empty)

169 if isinstance(msg, _messages.ModelRequest): (empty)

170 for i, part in enumerate(msg.parts): (empty)

171 if isinstance(part, _messages.SystemPromptPart) and part.dynamic_ref: (empty)

172 # Look up the runner by its ref 

173 if runner := self.system_prompt_dynamic_functions.get(part.dynamic_ref): 173 ↛ 170line 173 didn't jump to line 170 because the condition on line 173 was always true(empty)

174 updated_part_content = await runner.run(run_context) (empty)

175 msg.parts[i] = _messages.SystemPromptPart( (empty)

176 updated_part_content, dynamic_ref=part.dynamic_ref 

177 ) 

178 

179 async def _sys_parts(self, run_context: RunContext[DepsT]) -> list[_messages.ModelRequestPart]: (empty)

180 """Build the initial messages for the conversation.""" 

181 messages: list[_messages.ModelRequestPart] = [_messages.SystemPromptPart(p) for p in self.system_prompts] (empty)

182 for sys_prompt_runner in self.system_prompt_functions: (empty)

183 prompt = await sys_prompt_runner.run(run_context) (empty)

184 if sys_prompt_runner.dynamic: (empty)

185 messages.append(_messages.SystemPromptPart(prompt, dynamic_ref=sys_prompt_runner.function.__qualname__)) (empty)

186 else: 

187 messages.append(_messages.SystemPromptPart(prompt)) (empty)

188 return messages (empty)

189 

190 

191@dataclasses.dataclass (empty)

192class UserPromptNode(BaseUserPromptNode[DepsT, NodeRunEndT]): (empty)

193 async def run( (empty)

194 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]] 

195 ) -> ModelRequestNode[DepsT, NodeRunEndT]: 

196 return ModelRequestNode[DepsT, NodeRunEndT](request=await self._get_first_message(ctx)) (empty)

197 

198 

199@dataclasses.dataclass (empty)

200class StreamUserPromptNode(BaseUserPromptNode[DepsT, NodeRunEndT]): (empty)

201 async def run( (empty)

202 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]] 

203 ) -> StreamModelRequestNode[DepsT, NodeRunEndT]: 

204 return StreamModelRequestNode[DepsT, NodeRunEndT](request=await self._get_first_message(ctx)) (empty)

205 

206 

207async def _prepare_model( (empty)

208 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], 

209) -> models.AgentModel: 

210 """Build tools and create an agent model.""" 

211 function_tool_defs: list[ToolDefinition] = [] (empty)

212 

213 run_context = _build_run_context(ctx) (empty)

214 

215 async def add_tool(tool: Tool[DepsT]) -> None: (empty)

216 ctx = run_context.replace_with(retry=tool.current_retry, tool_name=tool.name) (empty)

217 if tool_def := await tool.prepare_tool_def(ctx): (empty)

218 function_tool_defs.append(tool_def) (empty)

219 

220 await asyncio.gather(*map(add_tool, ctx.deps.function_tools.values())) (empty)

221 

222 result_schema = ctx.deps.result_schema (empty)

223 return await run_context.model.agent_model( (empty)

224 function_tools=function_tool_defs, 

225 allow_text_result=_allow_text_result(result_schema), 

226 result_tools=result_schema.tool_defs() if result_schema is not None else [], 

227 ) 

228 

229 

230@dataclasses.dataclass (empty)

231class ModelRequestNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], NodeRunEndT]): (empty)

232 """Make a request to the model using the last message in state.message_history.""" 

233 

234 request: _messages.ModelRequest (empty)

235 

236 async def run( (empty)

237 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]] 

238 ) -> HandleResponseNode[DepsT, NodeRunEndT]: 

239 ctx.state.message_history.append(self.request) (empty)

240 

241 # Check usage 

242 if ctx.deps.usage_limits: 242 ↛ 246line 242 didn't jump to line 246 because the condition on line 242 was always true(empty)

243 ctx.deps.usage_limits.check_before_request(ctx.state.usage) (empty)

244 

245 # Increment run_step 

246 ctx.state.run_step += 1 (empty)

247 

248 with _logfire.span('preparing model and tools {run_step=}', run_step=ctx.state.run_step): (empty)

249 agent_model = await _prepare_model(ctx) (empty)

250 

251 # Actually make the model request 

252 model_settings = merge_model_settings(ctx.deps.model_settings, None) (empty)

253 with _logfire.span('model request') as span: (empty)

254 model_response, request_usage = await agent_model.request(ctx.state.message_history, model_settings) (empty)

255 span.set_attribute('response', model_response) (empty)

256 span.set_attribute('usage', request_usage) (empty)

257 

258 # Update usage 

259 ctx.state.usage.incr(request_usage, requests=1) (empty)

260 if ctx.deps.usage_limits: 260 ↛ 264line 260 didn't jump to line 264 because the condition on line 260 was always true(empty)

261 ctx.deps.usage_limits.check_tokens(ctx.state.usage) (empty)

262 

263 # Append the model response to state.message_history 

264 ctx.state.message_history.append(model_response) (empty)

265 return HandleResponseNode(model_response) (empty)

266 

267 

268@dataclasses.dataclass (empty)

269class HandleResponseNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], NodeRunEndT]): (empty)

270 """Process e response from a model, decide whether to end the run or make a new request.""" 

271 

272 model_response: _messages.ModelResponse (empty)

273 

274 async def run( (empty)

275 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]] 

276 ) -> Union[ModelRequestNode[DepsT, NodeRunEndT], FinalResultNode[DepsT, NodeRunEndT]]: # noqa UP007 

277 with _logfire.span('handle model response', run_step=ctx.state.run_step) as handle_span: (empty)

278 texts: list[str] = [] (empty)

279 tool_calls: list[_messages.ToolCallPart] = [] (empty)

280 for part in self.model_response.parts: (empty)

281 if isinstance(part, _messages.TextPart): (empty)

282 # ignore empty content for text parts, see #437 

283 if part.content: (empty)

284 texts.append(part.content) (empty)

285 elif isinstance(part, _messages.ToolCallPart): (empty)

286 tool_calls.append(part) (empty)

287 else: 

288 assert_never(part) 

289 

290 # At the moment, we prioritize at least executing tool calls if they are present. 

291 # In the future, we'd consider making this configurable at the agent or run level. 

292 # This accounts for cases like anthropic returns that might contain a text response 

293 # and a tool call response, where the text response just indicates the tool call will happen. 

294 if tool_calls: (empty)

295 return await self._handle_tool_calls_response(ctx, tool_calls, handle_span) (empty)

296 elif texts: (empty)

297 return await self._handle_text_response(ctx, texts, handle_span) (empty)

298 else: 

299 raise exceptions.UnexpectedModelBehavior('Received empty model response') (empty)

300 

301 async def _handle_tool_calls_response( (empty)

302 self, 

303 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], 

304 tool_calls: list[_messages.ToolCallPart], 

305 handle_span: logfire_api.LogfireSpan, 

306 ): 

307 result_schema = ctx.deps.result_schema (empty)

308 

309 # first look for the result tool call 

310 final_result: MarkFinalResult[NodeRunEndT] | None = None (empty)

311 parts: list[_messages.ModelRequestPart] = [] (empty)

312 if result_schema is not None: (empty)

313 if match := result_schema.find_tool(tool_calls): (empty)

314 call, result_tool = match (empty)

315 try: (empty)

316 result_data = result_tool.validate(call) (empty)

317 result_data = await _validate_result(result_data, ctx, call) (empty)

318 except _result.ToolRetryError as e: (empty)

319 # TODO: Should only increment retry stuff once per node execution, not for each tool call 

320 # Also, should increment the tool-specific retry count rather than the run retry count 

321 ctx.state.increment_retries(ctx.deps.max_result_retries) (empty)

322 parts.append(e.tool_retry) (empty)

323 else: 

324 final_result = MarkFinalResult(result_data, call.tool_name) (empty)

325 

326 # Then build the other request parts based on end strategy 

327 tool_responses = await _process_function_tools(tool_calls, final_result and final_result.tool_name, ctx) (empty)

328 

329 if final_result: (empty)

330 handle_span.set_attribute('result', final_result.data) (empty)

331 handle_span.message = 'handle model response -> final result' (empty)

332 return FinalResultNode[DepsT, NodeRunEndT](final_result, tool_responses) (empty)

333 else: 

334 if tool_responses: (empty)

335 handle_span.set_attribute('tool_responses', tool_responses) (empty)

336 tool_responses_str = ' '.join(r.part_kind for r in tool_responses) (empty)

337 handle_span.message = f'handle model response -> {tool_responses_str}' (empty)

338 parts.extend(tool_responses) (empty)

339 return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=parts)) (empty)

340 

341 async def _handle_text_response( (empty)

342 self, 

343 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], 

344 texts: list[str], 

345 handle_span: logfire_api.LogfireSpan, 

346 ): 

347 result_schema = ctx.deps.result_schema (empty)

348 

349 text = '\n\n'.join(texts) (empty)

350 if _allow_text_result(result_schema): (empty)

351 result_data_input = cast(NodeRunEndT, text) (empty)

352 try: (empty)

353 result_data = await _validate_result(result_data_input, ctx, None) (empty)

354 except _result.ToolRetryError as e: (empty)

355 ctx.state.increment_retries(ctx.deps.max_result_retries) (empty)

356 return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry])) (empty)

357 else: 

358 handle_span.set_attribute('result', result_data) (empty)

359 handle_span.message = 'handle model response -> final result' (empty)

360 return FinalResultNode[DepsT, NodeRunEndT](MarkFinalResult(result_data, None)) (empty)

361 else: 

362 ctx.state.increment_retries(ctx.deps.max_result_retries) (empty)

363 return ModelRequestNode[DepsT, NodeRunEndT]( (empty)

364 _messages.ModelRequest( 

365 parts=[ 

366 _messages.RetryPromptPart( 

367 content='Plain text responses are not permitted, please call one of the functions instead.', 

368 ) 

369 ] 

370 ) 

371 ) 

372 

373 

374@dataclasses.dataclass (empty)

375class StreamModelRequestNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], NodeRunEndT]): (empty)

376 """Make a request to the model using the last message in state.message_history (or a specified request).""" 

377 

378 request: _messages.ModelRequest (empty)

379 _result: StreamModelRequestNode[DepsT, NodeRunEndT] | End[result.StreamedRunResult[DepsT, NodeRunEndT]] | None = ( (empty)

380 field(default=None, repr=False) 

381 ) 

382 

383 async def run( (empty)

384 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]] 

385 ) -> Union[StreamModelRequestNode[DepsT, NodeRunEndT], End[result.StreamedRunResult[DepsT, NodeRunEndT]]]: # noqa UP007 

386 if self._result is not None: 386 ↛ 389line 386 didn't jump to line 389 because the condition on line 386 was always true(empty)

387 return self._result (empty)

388 

389 async with self.run_to_result(ctx) as final_node: 

390 return final_node 

391 

392 @asynccontextmanager (empty)

393 async def run_to_result( (empty)

394 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]] 

395 ) -> AsyncIterator[StreamModelRequestNode[DepsT, NodeRunEndT] | End[result.StreamedRunResult[DepsT, NodeRunEndT]]]: 

396 result_schema = ctx.deps.result_schema (empty)

397 

398 ctx.state.message_history.append(self.request) (empty)

399 

400 # Check usage 

401 if ctx.deps.usage_limits: 401 ↛ 405line 401 didn't jump to line 405 because the condition on line 401 was always true(empty)

402 ctx.deps.usage_limits.check_before_request(ctx.state.usage) (empty)

403 

404 # Increment run_step 

405 ctx.state.run_step += 1 (empty)

406 

407 with _logfire.span('preparing model and tools {run_step=}', run_step=ctx.state.run_step): (empty)

408 agent_model = await _prepare_model(ctx) (empty)

409 

410 # Actually make the model request 

411 model_settings = merge_model_settings(ctx.deps.model_settings, None) (empty)

412 with _logfire.span('model request {run_step=}', run_step=ctx.state.run_step) as model_req_span: (empty)

413 async with agent_model.request_stream(ctx.state.message_history, model_settings) as streamed_response: (empty)

414 ctx.state.usage.requests += 1 (empty)

415 model_req_span.set_attribute('response_type', streamed_response.__class__.__name__) (empty)

416 # We want to end the "model request" span here, but we can't exit the context manager 

417 # in the traditional way 

418 model_req_span.__exit__(None, None, None) (empty)

419 

420 with _logfire.span('handle model response') as handle_span: (empty)

421 received_text = False (empty)

422 

423 async for maybe_part_event in streamed_response: (empty)

424 if isinstance(maybe_part_event, _messages.PartStartEvent): (empty)

425 new_part = maybe_part_event.part (empty)

426 if isinstance(new_part, _messages.TextPart): (empty)

427 received_text = True (empty)

428 if _allow_text_result(result_schema): (empty)

429 handle_span.message = 'handle model response -> final result' (empty)

430 streamed_run_result = _build_streamed_run_result(streamed_response, None, ctx) (empty)

431 self._result = End(streamed_run_result) (empty)

432 yield self._result (empty)

433 return (empty)

434 elif isinstance(new_part, _messages.ToolCallPart): (empty)

435 if result_schema is not None and (match := result_schema.find_tool([new_part])): (empty)

436 call, _ = match (empty)

437 handle_span.message = 'handle model response -> final result' (empty)

438 streamed_run_result = _build_streamed_run_result( (empty)

439 streamed_response, call.tool_name, ctx 

440 ) 

441 self._result = End(streamed_run_result) (empty)

442 yield self._result (empty)

443 return (empty)

444 else: 

445 assert_never(new_part) (empty)

446 

447 tasks: list[asyncio.Task[_messages.ModelRequestPart]] = [] (empty)

448 parts: list[_messages.ModelRequestPart] = [] (empty)

449 model_response = streamed_response.get() (empty)

450 if not model_response.parts: (empty)

451 raise exceptions.UnexpectedModelBehavior('Received empty model response') (empty)

452 ctx.state.message_history.append(model_response) (empty)

453 

454 run_context = _build_run_context(ctx) (empty)

455 for p in model_response.parts: (empty)

456 if isinstance(p, _messages.ToolCallPart): (empty)

457 if tool := ctx.deps.function_tools.get(p.tool_name): (empty)

458 tasks.append(asyncio.create_task(tool.run(p, run_context), name=p.tool_name)) (empty)

459 else: 

460 parts.append(_unknown_tool(p.tool_name, ctx)) (empty)

461 

462 if received_text and not tasks and not parts: (empty)

463 # Can only get here if self._allow_text_result returns `False` for the provided result_schema 

464 ctx.state.increment_retries(ctx.deps.max_result_retries) (empty)

465 self._result = StreamModelRequestNode[DepsT, NodeRunEndT]( (empty)

466 _messages.ModelRequest( 

467 parts=[ 

468 _messages.RetryPromptPart( 

469 content='Plain text responses are not permitted, please call one of the functions instead.', 

470 ) 

471 ] 

472 ) 

473 ) 

474 yield self._result (empty)

475 return (empty)

476 

477 with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]): (empty)

478 task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks) (empty)

479 parts.extend(task_results) (empty)

480 

481 next_request = _messages.ModelRequest(parts=parts) (empty)

482 if any(isinstance(part, _messages.RetryPromptPart) for part in parts): (empty)

483 try: (empty)

484 ctx.state.increment_retries(ctx.deps.max_result_retries) (empty)

485 except: (empty)

486 # TODO: This is janky, so I think we should probably change it, but how? 

487 ctx.state.message_history.append(next_request) (empty)

488 raise (empty)

489 

490 handle_span.set_attribute('tool_responses', parts) (empty)

491 tool_responses_str = ' '.join(r.part_kind for r in parts) (empty)

492 handle_span.message = f'handle model response -> {tool_responses_str}' (empty)

493 # the model_response should have been fully streamed by now, we can add its usage 

494 streamed_response_usage = streamed_response.usage() (empty)

495 run_context.usage.incr(streamed_response_usage) (empty)

496 ctx.deps.usage_limits.check_tokens(run_context.usage) (empty)

497 self._result = StreamModelRequestNode[DepsT, NodeRunEndT](next_request) (empty)

498 yield self._result (empty)

499 return (empty)

500 

501 

502@dataclasses.dataclass (empty)

503class FinalResultNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], MarkFinalResult[NodeRunEndT]]): (empty)

504 """Produce the final result of the run.""" 

505 

506 data: MarkFinalResult[NodeRunEndT] (empty)

507 """The final result data.""" (empty)

508 extra_parts: list[_messages.ModelRequestPart] = dataclasses.field(default_factory=list) (empty)

509 

510 async def run( (empty)

511 self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]] 

512 ) -> End[MarkFinalResult[NodeRunEndT]]: 

513 run_span = ctx.deps.run_span (empty)

514 usage = ctx.state.usage (empty)

515 messages = ctx.state.message_history (empty)

516 

517 # TODO: For backwards compatibility, append a new ModelRequest using the tool returns and retries 

518 if self.extra_parts: (empty)

519 messages.append(_messages.ModelRequest(parts=self.extra_parts)) (empty)

520 

521 # TODO: Set this attribute somewhere 

522 # handle_span = self.handle_model_response_span 

523 # handle_span.set_attribute('final_data', self.data) 

524 run_span.set_attribute('usage', usage) (empty)

525 run_span.set_attribute('all_messages', messages) (empty)

526 

527 # End the run with self.data 

528 return End(self.data) (empty)

529 

530 

531def _build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]]) -> RunContext[DepsT]: (empty)

532 return RunContext[DepsT]( (empty)

533 deps=ctx.deps.user_deps, 

534 model=ctx.deps.model, 

535 usage=ctx.state.usage, 

536 prompt=ctx.deps.prompt, 

537 messages=ctx.state.message_history, 

538 run_step=ctx.state.run_step, 

539 ) 

540 

541 

542def _build_streamed_run_result( (empty)

543 result_stream: models.StreamedResponse, 

544 result_tool_name: str | None, 

545 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], 

546) -> result.StreamedRunResult[DepsT, NodeRunEndT]: 

547 new_message_index = ctx.deps.new_message_index (empty)

548 result_schema = ctx.deps.result_schema (empty)

549 run_span = ctx.deps.run_span (empty)

550 usage_limits = ctx.deps.usage_limits (empty)

551 messages = ctx.state.message_history (empty)

552 run_context = _build_run_context(ctx) (empty)

553 

554 async def on_complete(): (empty)

555 """Called when the stream has completed. 

556 

557 The model response will have been added to messages by now 

558 by `StreamedRunResult._marked_completed`. 

559 """ 

560 last_message = messages[-1] (empty)

561 assert isinstance(last_message, _messages.ModelResponse) (empty)

562 tool_calls = [part for part in last_message.parts if isinstance(part, _messages.ToolCallPart)] (empty)

563 parts = await _process_function_tools( (empty)

564 tool_calls, 

565 result_tool_name, 

566 ctx, 

567 ) 

568 # TODO: Should we do something here related to the retry count? 

569 # Maybe we should move the incrementing of the retry count to where we actually make a request? 

570 # if any(isinstance(part, _messages.RetryPromptPart) for part in parts): 

571 # ctx.state.increment_retries(ctx.deps.max_result_retries) 

572 if parts: (empty)

573 messages.append(_messages.ModelRequest(parts)) (empty)

574 run_span.set_attribute('all_messages', messages) (empty)

575 

576 return result.StreamedRunResult[DepsT, NodeRunEndT]( (empty)

577 messages, 

578 new_message_index, 

579 usage_limits, 

580 result_stream, 

581 result_schema, 

582 run_context, 

583 ctx.deps.result_validators, 

584 result_tool_name, 

585 on_complete, 

586 ) 

587 

588 

589async def _process_function_tools( (empty)

590 tool_calls: list[_messages.ToolCallPart], 

591 result_tool_name: str | None, 

592 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], 

593) -> list[_messages.ModelRequestPart]: 

594 """Process function (non-result) tool calls in parallel. 

595 

596 Also add stub return parts for any other tools that need it. 

597 """ 

598 parts: list[_messages.ModelRequestPart] = [] (empty)

599 tasks: list[asyncio.Task[_messages.ToolReturnPart | _messages.RetryPromptPart]] = [] (empty)

600 

601 stub_function_tools = bool(result_tool_name) and ctx.deps.end_strategy == 'early' (empty)

602 result_schema = ctx.deps.result_schema (empty)

603 

604 # we rely on the fact that if we found a result, it's the first result tool in the last 

605 found_used_result_tool = False (empty)

606 run_context = _build_run_context(ctx) (empty)

607 

608 for call in tool_calls: (empty)

609 if call.tool_name == result_tool_name and not found_used_result_tool: (empty)

610 found_used_result_tool = True (empty)

611 parts.append( (empty)

612 _messages.ToolReturnPart( 

613 tool_name=call.tool_name, 

614 content='Final result processed.', 

615 tool_call_id=call.tool_call_id, 

616 ) 

617 ) 

618 elif tool := ctx.deps.function_tools.get(call.tool_name): (empty)

619 if stub_function_tools: (empty)

620 parts.append( (empty)

621 _messages.ToolReturnPart( 

622 tool_name=call.tool_name, 

623 content='Tool not executed - a final result was already processed.', 

624 tool_call_id=call.tool_call_id, 

625 ) 

626 ) 

627 else: 

628 tasks.append(asyncio.create_task(tool.run(call, run_context), name=call.tool_name)) (empty)

629 elif result_schema is not None and call.tool_name in result_schema.tools: (empty)

630 # if tool_name is in _result_schema, it means we found a result tool but an error occurred in 

631 # validation, we don't add another part here 

632 if result_tool_name is not None: (empty)

633 parts.append( (empty)

634 _messages.ToolReturnPart( 

635 tool_name=call.tool_name, 

636 content='Result tool not used - a final result was already processed.', 

637 tool_call_id=call.tool_call_id, 

638 ) 

639 ) 

640 else: 

641 parts.append(_unknown_tool(call.tool_name, ctx)) (empty)

642 

643 # Run all tool tasks in parallel 

644 if tasks: (empty)

645 with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]): (empty)

646 task_results: Sequence[_messages.ToolReturnPart | _messages.RetryPromptPart] = await asyncio.gather(*tasks) (empty)

647 for result in task_results: (empty)

648 if isinstance(result, _messages.ToolReturnPart): (empty)

649 parts.append(result) (empty)

650 elif isinstance(result, _messages.RetryPromptPart): (empty)

651 parts.append(result) (empty)

652 else: 

653 assert_never(result) (empty)

654 return parts (empty)

655 

656 

657def _unknown_tool( (empty)

658 tool_name: str, 

659 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], 

660) -> _messages.RetryPromptPart: 

661 ctx.state.increment_retries(ctx.deps.max_result_retries) (empty)

662 tool_names = list(ctx.deps.function_tools.keys()) (empty)

663 if result_schema := ctx.deps.result_schema: (empty)

664 tool_names.extend(result_schema.tool_names()) (empty)

665 

666 if tool_names: (empty)

667 msg = f'Available tools: {", ".join(tool_names)}' (empty)

668 else: 

669 msg = 'No tools available.' (empty)

670 

671 return _messages.RetryPromptPart(content=f'Unknown tool name: {tool_name!r}. {msg}') (empty)

672 

673 

674async def _validate_result( (empty)

675 result_data: T, 

676 ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, T]], 

677 tool_call: _messages.ToolCallPart | None, 

678) -> T: 

679 for validator in ctx.deps.result_validators: (empty)

680 run_context = _build_run_context(ctx) (empty)

681 result_data = await validator.validate(result_data, tool_call, run_context) (empty)

682 return result_data (empty)

683 

684 

685def _allow_text_result(result_schema: _result.ResultSchema[Any] | None) -> bool: (empty)

686 return result_schema is None or result_schema.allow_text_result (empty)

687 

688 

689@dataclasses.dataclass (empty)

690class _RunMessages: (empty)

691 messages: list[_messages.ModelMessage] (empty)

692 used: bool = False (empty)

693 

694 

695_messages_ctx_var: ContextVar[_RunMessages] = ContextVar('var') (empty)

696 

697 

698@contextmanager (empty)

699def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]: (empty)

700 """Context manager to access the messages used in a [`run`][pydantic_ai.Agent.run], [`run_sync`][pydantic_ai.Agent.run_sync], or [`run_stream`][pydantic_ai.Agent.run_stream] call. 

701 

702 Useful when a run may raise an exception, see [model errors](../agents.md#model-errors) for more information. 

703 

704 Examples: 

705 ```python 

706 from pydantic_ai import Agent, capture_run_messages 

707 

708 agent = Agent('test') 

709 

710 with capture_run_messages() as messages: 

711 try: 

712 result = agent.run_sync('foobar') 

713 except Exception: 

714 print(messages) 

715 raise 

716 ``` 

717 

718 !!! note 

719 If you call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context, 

720 `messages` will represent the messages exchanged during the first call only. 

721 """ 

722 try: (empty)

723 yield _messages_ctx_var.get().messages (empty)

724 except LookupError: (empty)