-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Adding run_id to ModelRequest and ModelResponse #3366
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Adding run_id to ModelRequest and ModelResponse #3366
Conversation
…raphState to these ModelMessages
|
@DouweM A lot of tests will need to be modified to accommodate run_id in the snapshots. |
| ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]], | ||
| response: _messages.ModelResponse, | ||
| ) -> CallToolsNode[DepsT, NodeRunEndT]: | ||
| response.run_id = response.run_id or ctx.state.run_id |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We may also need to do this in StreamedRunResult._marked_completed to get this to work right when using agent.run_stream:
pydantic-ai/pydantic_ai_slim/pydantic_ai/result.py
Lines 555 to 556 in 4cc4f35
| if message is not None: | |
| self._all_messages.append(message) |
You'll want to test if it currently works, and if not add it.
We also definitely need to set the run_id here:
pydantic-ai/pydantic_ai_slim/pydantic_ai/agent/abstract.py
Lines 560 to 562 in 4cc4f35
# For backwards compatibility, append a new ModelRequest using the tool returns and retries if parts: messages.append(_messages.ModelRequest(parts)) pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py
Lines 742 to 744 in 4cc4f35
# For backwards compatibility, append a new ModelRequest using the tool returns and retries if tool_responses: messages.append(_messages.ModelRequest(parts=tool_responses))
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah I made it a part of RunContext and passed it, added for the others two as well
@adtyavrdhn Yep, we'll need |
…dding_run_id_to_modelmessage
|
This ended up being a much bigger diff than what I was expecting |
| def serialize_run_context(cls, ctx: RunContext[Any]) -> dict[str, Any]: | ||
| """Serialize the run context to a `dict[str, Any]`.""" | ||
| return { | ||
| 'run_id': ctx.run_id, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added it as a part of serialization for temporal and I added a small test for the round trip.
Closes #3171
Adds a UUID run_id at agent start, thread it through GraphAgentState, and assign it to every ModelRequest/ModelResponse as they’re created so all execution paths share one identifier.
Keep run_id optional on message classes for backward compatibility.
Cover the change with tests:
(1) assert captured messages and result.all_messages() all carry the same non-empty run_id and
(2) confirm ModelMessagesTypeAdapter preserves the field through serialization.