Adding run_id to ModelRequest and ModelResponse, propagating it via GraphState to these ModelMessages

Author: adtyavrdhn

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -92,6 +92,7 @@ class GraphAgentState:
     usage: _usage.RunUsage = dataclasses.field(default_factory=_usage.RunUsage)
     retries: int = 0
     run_step: int = 0
+    run_id: str | None = None
 
     def increment_retries(
         self,
@@ -469,6 +470,7 @@ async def _make_request(
     async def _prepare_request(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> tuple[ModelSettings | None, models.ModelRequestParameters, list[_messages.ModelMessage], RunContext[DepsT]]:
+        self.request.run_id = self.request.run_id or ctx.state.run_id
         ctx.state.message_history.append(self.request)
 
         ctx.state.run_step += 1
@@ -510,6 +512,7 @@ def _finish_handling(
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
         response: _messages.ModelResponse,
     ) -> CallToolsNode[DepsT, NodeRunEndT]:
+        response.run_id = response.run_id or ctx.state.run_id
         # Update usage
         ctx.state.usage.incr(response.usage)
         if ctx.deps.usage_limits:  # pragma: no branch

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -3,6 +3,7 @@
 import dataclasses
 import inspect
 import json
+import uuid
 import warnings
 from asyncio import Lock
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
@@ -572,6 +573,7 @@ async def main():
             usage=usage,
             retries=0,
             run_step=0,
+            run_id=str(uuid.uuid4()),
         )
 
         # Merge model settings in order of precedence: run > agent > model

pydantic_ai_slim/pydantic_ai/messages.py

@@ -947,6 +947,9 @@ class ModelRequest:
     kind: Literal['request'] = 'request'
     """Message type identifier, this is available on all parts as a discriminator."""
 
+    run_id: str | None = None
+    """A unique identifier to identify the run."""
+
     @classmethod
     def user_text_prompt(cls, user_prompt: str, *, instructions: str | None = None) -> ModelRequest:
         """Create a `ModelRequest` with a single user prompt as text."""
@@ -1188,6 +1191,9 @@ class ModelResponse:
     finish_reason: FinishReason | None = None
     """Reason the model finished generating the response, normalized to OpenTelemetry values."""
 
+    run_id: str | None = None
+    """A unique identifier to identify the run."""
+
     @property
     def text(self) -> str | None:
         """Get the text in the response."""

tests/test_agent.py

@@ -2506,6 +2506,17 @@ def llm(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
     )
 
 
+def test_agent_message_history_includes_run_id() -> None:
+    agent = Agent(TestModel(custom_output_text='testing run_id'))
+
+    result = agent.run_sync('Hello')
+    history = result.all_messages()
+
+    run_ids = [message.run_id for message in history]
+    assert run_ids == snapshot([IsStr(), IsStr()])
+    assert len({*run_ids}) == snapshot(1)
+
+
 def test_unknown_tool():
     def empty(_: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
         return ModelResponse(parts=[ToolCallPart('foobar', '{}')])

tests/test_messages.py

@@ -469,6 +469,24 @@ def test_file_part_serialization_roundtrip():
     assert deserialized == messages
 
 
+def test_model_messages_type_adapter_preserves_run_id():
+    messages: list[ModelMessage] = [
+        ModelRequest(
+            parts=[UserPromptPart(content='Hi there', timestamp=datetime.now(tz=timezone.utc))],
+            run_id='run-123',
+        ),
+        ModelResponse(
+            parts=[TextPart(content='Hello!')],
+            run_id='run-123',
+        ),
+    ]
+
+    serialized = ModelMessagesTypeAdapter.dump_python(messages, mode='python')
+    deserialized = ModelMessagesTypeAdapter.validate_python(serialized)
+
+    assert [message.run_id for message in deserialized] == snapshot(['run-123', 'run-123'])
+
+
 def test_model_response_convenience_methods():
     response = ModelResponse(parts=[])
     assert response.text == snapshot(None)