Add ability to end agent runs as a result of a tool call #142
Conversation
dmontagu
changed the title
| @@ -45,6 +51,11 @@ class RunContext(Generic[AgentDeps]): | |||
| tool_name: str | None | |||
| """Name of the tool being called.""" | |||
|
|
|||
| def end_run(self, result: ResultData) -> NoReturn: | |||
| """End the call to `agent.run` as soon as possible, using the provided value as the result.""" | |||
| # NOTE: this means we ignore any other tools called concurrently | |||
There was a problem hiding this comment.
mention in docs what happens about cancelling other tools. Do the return values of functions that finish get added to messages?
There was a problem hiding this comment.
mention in docs what happens about cancelling other tools.
I added this comment to match the comments near the result schema things:
elif model_response.role == 'model-structured-response':
if self._result_schema is not None:
# if there's a result schema, and any of the calls match one of its tools, return the result
# NOTE: this means we ignore any other tools called here
if match := self._result_schema.find_tool(model_response):Is there a discussion of this in the docs? I didn't see one in a quick search, but if so I can adapt that; if not, I guess we could add a note about that too?
There was a problem hiding this comment.
I would propose that we wait to update the docs about this until after we merge Jeremiah's PR that adds the "eager" vs. "correct" mode, and support reflect those modes here too
Deploying pydantic-ai with
|
| Latest commit: |
d20034c
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://af7fb016.pydantic-ai.pages.dev |
| Branch Preview URL: | https://end-agent-run.pydantic-ai.pages.dev |
| messages.extend(task_results) | ||
| except _exceptions.StopAgentRun as e: | ||
| span.set_attribute('stop_agent_run', e.tool_name) | ||
| return _MarkFinalResult(data=e.result), [] |
There was a problem hiding this comment.
This is not correct, and basically should be a type error, but the problem is that e.result has type Any (it is the value passed to ctx.stop_run). However, I am not sure what's going on here so not sure how to tweak things to fix this. Maybe we can discuss tomorrow.
| ToolReturn(tool_name='final_result', content=('abcdef', 777), timestamp=IsNow(tz=timezone.utc)), | ||
| ] | ||
| ) | ||
| assert await result.get_data() == snapshot(('abcdef', 777)) |
There was a problem hiding this comment.
I'm hitting:
E pydantic_ai.exceptions.UnexpectedModelBehavior: Invalid message, unable to find tool. Called tool names: ['maybe_stop_run_tool']. Expected tool names: ['final_result']
And I can't for the life of me figure out what is happening or why. Could use some help understanding how streaming is supposed to work
With this change, you can have tool calls end the full run early by calling
ctx.stop_run(result)inside the tool call — this immediately ends tool execution (and type-checkers can even tell that subsequent code is unreachable).I've also added the ability to disable the standard result tools through a new keyword argument, so you can ensure that the only way for the agent to exit is to call a tool that exits by callingI think we should add this in a separate PR that exposes more control over the auto-generated result tools.ctx.stop_run.I still need to add documentation and examples.
NOTE: I no longer intend for this to resolve #127, as this is a fairly convoluted way to implement a "result tool", and I think we should just add an API for that directly. But @samuelcolvin has noted that there are other possible use cases for this feature, such as providing ways to explicitly interrupt agent execution during tool calls, etc., so I think it still makes sense to merge this.