Make agent and graph runs serializable again (#1545)

Co-authored-by: Douwe Maan <hi@douwe.me>

Author: dmontagu

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -291,13 +291,4 @@ async def __anext__(self) -> T:
 
 
 def get_traceparent(x: AgentRun | AgentRunResult | GraphRun | GraphRunResult) -> str:
-    import logfire
-    import logfire_api
-    from logfire.experimental.annotations import get_traceparent
-
-    span: AbstractSpan | None = x._span(required=False)  # type: ignore[reportPrivateUsage]
-    if not span:  # pragma: no cover
-        return ''
-    if isinstance(span, logfire_api.LogfireSpan):  # pragma: no cover
-        assert isinstance(span, logfire.LogfireSpan)
-    return get_traceparent(span)
+    return x._traceparent(required=False) or ''  # type: ignore[reportPrivateUsage]

pydantic_ai_slim/pydantic_ai/agent.py

@@ -27,7 +27,6 @@
     result,
     usage as _usage,
 )
-from ._utils import AbstractSpan
 from .models.instrumented import InstrumentationSettings, InstrumentedModel
 from .result import FinalResult, OutputDataT, StreamedRunResult, ToolOutput
 from .settings import ModelSettings, merge_model_settings
@@ -1683,14 +1682,14 @@ async def main():
     ]
 
     @overload
-    def _span(self, *, required: Literal[False]) -> AbstractSpan | None: ...
+    def _traceparent(self, *, required: Literal[False]) -> str | None: ...
     @overload
-    def _span(self) -> AbstractSpan: ...
-    def _span(self, *, required: bool = True) -> AbstractSpan | None:
-        span = self._graph_run._span(required=False)  # type: ignore[reportPrivateUsage]
-        if span is None and required:  # pragma: no cover
-            raise AttributeError('Span is not available for this agent run')
-        return span
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        traceparent = self._graph_run._traceparent(required=False)  # type: ignore[reportPrivateUsage]
+        if traceparent is None and required:  # pragma: no cover
+            raise AttributeError('No span was created for this agent run')
+        return traceparent
 
     @property
     def ctx(self) -> GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any]]:
@@ -1729,7 +1728,7 @@ def result(self) -> AgentRunResult[OutputDataT] | None:
             graph_run_result.output.tool_name,
             graph_run_result.state,
             self._graph_run.deps.new_message_index,
-            self._graph_run._span(required=False),  # type: ignore[reportPrivateUsage]
+            self._traceparent(required=False),
         )
 
     def __aiter__(
@@ -1847,16 +1846,16 @@ class AgentRunResult(Generic[OutputDataT]):
     _output_tool_name: str | None = dataclasses.field(repr=False)
     _state: _agent_graph.GraphAgentState = dataclasses.field(repr=False)
     _new_message_index: int = dataclasses.field(repr=False)
-    _span_value: AbstractSpan | None = dataclasses.field(repr=False)
+    _traceparent_value: str | None = dataclasses.field(repr=False)
 
     @overload
-    def _span(self, *, required: Literal[False]) -> AbstractSpan | None: ...
+    def _traceparent(self, *, required: Literal[False]) -> str | None: ...
     @overload
-    def _span(self) -> AbstractSpan: ...
-    def _span(self, *, required: bool = True) -> AbstractSpan | None:
-        if self._span_value is None and required:  # pragma: no cover
-            raise AttributeError('Span is not available for this agent run')
-        return self._span_value
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        if self._traceparent_value is None and required:  # pragma: no cover
+            raise AttributeError('No span was created for this agent run')
+        return self._traceparent_value
 
     @property
     @deprecated('`result.data` is deprecated, use `result.output` instead.')

pydantic_graph/pydantic_graph/_utils.py

@@ -3,16 +3,47 @@
 import asyncio
 import types
 from functools import partial
-from typing import Any, Callable, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
 from logfire_api import LogfireSpan
-from opentelemetry.trace import Span
 from typing_extensions import ParamSpec, TypeAlias, TypeIs, get_args, get_origin
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+
 AbstractSpan: TypeAlias = 'LogfireSpan | Span'
 
+try:
+    from opentelemetry.trace import Span, set_span_in_context
+    from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+    TRACEPARENT_PROPAGATOR = TraceContextTextMapPropagator()
+    TRACEPARENT_NAME = 'traceparent'
+    assert TRACEPARENT_NAME in TRACEPARENT_PROPAGATOR.fields
+
+    # Logic taken from logfire.experimental.annotations
+    def get_traceparent(span: AbstractSpan) -> str | None:
+        """Get a string representing the span context to use for annotating spans."""
+        real_span: Span
+        if isinstance(span, Span):
+            real_span = span
+        else:
+            real_span = span._span
+            assert real_span
+        context = set_span_in_context(real_span)
+        carrier: dict[str, Any] = {}
+        TRACEPARENT_PROPAGATOR.inject(carrier, context)
+        return carrier.get(TRACEPARENT_NAME, '')
+
+except ImportError:  # pragma: no cover
+
+    def get_traceparent(span: AbstractSpan) -> str | None:
+        # Opentelemetry wasn't installed, so we can't get the traceparent
+        return None
+
 
 def get_event_loop():
     try:

pydantic_graph/pydantic_graph/graph.py

@@ -15,7 +15,7 @@
 from typing_inspection import typing_objects
 
 from . import _utils, exceptions, mermaid
-from ._utils import AbstractSpan
+from ._utils import AbstractSpan, get_traceparent
 from .nodes import BaseNode, DepsT, End, GraphRunContext, NodeDef, RunEndT, StateT
 from .persistence import BaseStatePersistence
 from .persistence.in_mem import SimpleStatePersistence
@@ -257,8 +257,14 @@ async def iter(
                     entered_span = stack.enter_context(logfire_api.span('run graph {graph.name}', graph=self))
             else:
                 entered_span = stack.enter_context(span)
+            traceparent = None if entered_span is None else get_traceparent(entered_span)
             yield GraphRun[StateT, DepsT, RunEndT](
-                graph=self, start_node=start_node, persistence=persistence, state=state, deps=deps, span=entered_span
+                graph=self,
+                start_node=start_node,
+                persistence=persistence,
+                state=state,
+                deps=deps,
+                traceparent=traceparent,
             )
 
     @asynccontextmanager
@@ -301,14 +307,15 @@ async def iter_from_persistence(
 
         with ExitStack() as stack:
             entered_span = None if span is None else stack.enter_context(span)
+            traceparent = None if entered_span is None else get_traceparent(entered_span)
             yield GraphRun[StateT, DepsT, RunEndT](
                 graph=self,
                 start_node=snapshot.node,
                 persistence=persistence,
                 state=snapshot.state,
                 deps=deps,
                 snapshot_id=snapshot.id,
-                span=entered_span,
+                traceparent=traceparent,
             )
 
     async def initialize(
@@ -369,7 +376,7 @@ async def next(
             persistence=persistence,
             state=state,
             deps=deps,
-            span=None,
+            traceparent=None,
         )
         return await run.next(node)
 
@@ -644,7 +651,7 @@ def __init__(
         persistence: BaseStatePersistence[StateT, RunEndT],
         state: StateT,
         deps: DepsT,
-        span: AbstractSpan | None,
+        traceparent: str | None,
         snapshot_id: str | None = None,
     ):
         """Create a new run for a given graph, starting at the specified node.
@@ -659,7 +666,7 @@ def __init__(
                 to all nodes via `ctx.state`.
             deps: Optional dependencies that each node can access via `ctx.deps`, e.g. database connections,
                 configuration, or logging clients.
-            span: The span used for the graph run.
+            traceparent: The traceparent for the span used for the graph run.
             snapshot_id: The ID of the snapshot the node came from.
         """
         self.graph = graph
@@ -668,18 +675,18 @@ def __init__(
         self.state = state
         self.deps = deps
 
-        self.__span = span
+        self.__traceparent = traceparent
         self._next_node: BaseNode[StateT, DepsT, RunEndT] | End[RunEndT] = start_node
         self._is_started: bool = False
 
     @overload
-    def _span(self, *, required: typing_extensions.Literal[False]) -> AbstractSpan | None: ...
+    def _traceparent(self, *, required: typing_extensions.Literal[False]) -> str | None: ...
     @overload
-    def _span(self) -> AbstractSpan: ...
-    def _span(self, *, required: bool = True) -> AbstractSpan | None:
-        if self.__span is None and required:  # pragma: no cover
-            raise exceptions.GraphRuntimeError('No span available for this graph run.')
-        return self.__span
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        if self.__traceparent is None and required:  # pragma: no cover
+            raise exceptions.GraphRuntimeError('No span was created for this graph run')
+        return self.__traceparent
 
     @property
     def next_node(self) -> BaseNode[StateT, DepsT, RunEndT] | End[RunEndT]:
@@ -695,7 +702,10 @@ def result(self) -> GraphRunResult[StateT, RunEndT] | None:
         if not isinstance(self._next_node, End):
             return None  # The GraphRun has not finished running
         return GraphRunResult[StateT, RunEndT](
-            self._next_node.data, state=self.state, persistence=self.persistence, span=self._span(required=False)
+            self._next_node.data,
+            state=self.state,
+            persistence=self.persistence,
+            traceparent=self._traceparent(required=False),
         )
 
     async def next(
@@ -816,18 +826,18 @@ def __init__(
         output: RunEndT,
         state: StateT,
         persistence: BaseStatePersistence[StateT, RunEndT],
-        span: AbstractSpan | None = None,
+        traceparent: str | None = None,
     ):
         self.output = output
         self.state = state
         self.persistence = persistence
-        self.__span = span
+        self.__traceparent = traceparent
 
     @overload
-    def _span(self, *, required: typing_extensions.Literal[False]) -> AbstractSpan | None: ...
+    def _traceparent(self, *, required: typing_extensions.Literal[False]) -> str | None: ...
     @overload
-    def _span(self) -> AbstractSpan: ...
-    def _span(self, *, required: bool = True) -> AbstractSpan | None:  # pragma: no cover
-        if self.__span is None and required:
-            raise exceptions.GraphRuntimeError('No span available for this graph run.')
-        return self.__span
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:  # pragma: no cover
+        if self.__traceparent is None and required:
+            raise exceptions.GraphRuntimeError('No span was created for this graph run.')
+        return self.__traceparent

tests/test_agent.py

@@ -8,10 +8,11 @@
 import pytest
 from dirty_equals import IsJson
 from inline_snapshot import snapshot
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, TypeAdapter, field_validator
 from pydantic_core import to_json
 
 from pydantic_ai import Agent, ModelRetry, RunContext, UnexpectedModelBehavior, UserError, capture_run_messages
+from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import (
     BinaryContent,
     ModelMessage,
@@ -1869,3 +1870,16 @@ def my_tool(x: int) -> int:
             ModelResponse(parts=[], model_name='function:llm:', timestamp=IsNow(tz=timezone.utc)),
         ]
     )
+
+
+def test_agent_run_result_serialization() -> None:
+    agent = Agent('test', output_type=Foo)
+    result = agent.run_sync('Hello')
+
+    # Check that dump_json doesn't raise an error
+    adapter = TypeAdapter(AgentRunResult[Foo])
+    serialized_data = adapter.dump_json(result)
+
+    # Check that we can load the data back
+    deserialized_result = adapter.validate_json(serialized_data)
+    assert deserialized_result == result