Movatterモバイル変換


[0]ホーム

URL:


Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Sign up
Appearance settings

Commit21c2351

Browse files
Add spans as an attribute on agent/graph runs/runresults (pydantic#1415)
Co-authored-by: Alex Hall <alex.mojaki@gmail.com>
1 parentf7263ff commit21c2351

File tree

7 files changed

+193
-28
lines changed

7 files changed

+193
-28
lines changed

‎pydantic_ai_slim/pydantic_ai/_utils.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,14 @@
1515
frompydantic.json_schemaimportJsonSchemaValue
1616
fromtyping_extensionsimportParamSpec,TypeAlias,TypeGuard,is_typeddict
1717

18+
frompydantic_graph._utilsimportAbstractSpan
19+
20+
AbstractSpan=AbstractSpan
21+
1822
ifTYPE_CHECKING:
23+
frompydantic_ai.agentimportAgentRun,AgentRunResult
24+
frompydantic_graphimportGraphRun,GraphRunResult
25+
1926
from .importmessagesas_messages
2027
from .toolsimportObjectJsonSchema
2128

@@ -281,3 +288,16 @@ async def __anext__(self) -> T:
281288
exceptStopAsyncIteration:
282289
self._exhausted=True
283290
raise
291+
292+
293+
defget_traceparent(x:AgentRun|AgentRunResult|GraphRun|GraphRunResult)->str:
294+
importlogfire
295+
importlogfire_api
296+
fromlogfire.experimental.annotationsimportget_traceparent
297+
298+
span:AbstractSpan|None=x._span(required=False)# type: ignore[reportPrivateUsage]
299+
ifnotspan:# pragma: no cover
300+
return''
301+
ifisinstance(span,logfire_api.LogfireSpan):# pragma: no cover
302+
assertisinstance(span,logfire.LogfireSpan)
303+
returnget_traceparent(span)

‎pydantic_ai_slim/pydantic_ai/agent.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
fromopentelemetry.traceimportNoOpTracer,use_span
1212
frompydantic.json_schemaimportGenerateJsonSchema
13-
fromtyping_extensionsimportTypeGuard,TypeVar,deprecated
13+
fromtyping_extensionsimportLiteral,TypeGuard,TypeVar,deprecated
1414

1515
frompydantic_graphimportEnd,Graph,GraphRun,GraphRunContext
1616
frompydantic_graph._utilsimportget_event_loop
@@ -26,6 +26,7 @@
2626
result,
2727
usageas_usage,
2828
)
29+
from ._utilsimportAbstractSpan
2930
from .models.instrumentedimportInstrumentationSettings,InstrumentedModel
3031
from .resultimportFinalResult,ResultDataT,StreamedRunResult
3132
from .settingsimportModelSettings,merge_model_settings
@@ -52,6 +53,7 @@
5253
ifTYPE_CHECKING:
5354
frompydantic_ai.mcpimportMCPServer
5455

56+
5557
__all__= (
5658
'Agent',
5759
'AgentRun',
@@ -1402,6 +1404,16 @@ async def main():
14021404
_agent_graph.GraphAgentState,_agent_graph.GraphAgentDeps[AgentDepsT,Any],FinalResult[ResultDataT]
14031405
]
14041406

1407+
@overload
1408+
def_span(self,*,required:Literal[False])->AbstractSpan|None: ...
1409+
@overload
1410+
def_span(self)->AbstractSpan: ...
1411+
def_span(self,*,required:bool=True)->AbstractSpan|None:
1412+
span=self._graph_run._span(required=False)# type: ignore[reportPrivateUsage]
1413+
ifspanisNoneandrequired:# pragma: no cover
1414+
raiseAttributeError('Span is not available for this agent run')
1415+
returnspan
1416+
14051417
@property
14061418
defctx(self)->GraphRunContext[_agent_graph.GraphAgentState,_agent_graph.GraphAgentDeps[AgentDepsT,Any]]:
14071419
"""The current context of the agent run."""
@@ -1439,6 +1451,7 @@ def result(self) -> AgentRunResult[ResultDataT] | None:
14391451
graph_run_result.output.tool_name,
14401452
graph_run_result.state,
14411453
self._graph_run.deps.new_message_index,
1454+
self._graph_run._span(required=False),# type: ignore[reportPrivateUsage]
14421455
)
14431456

14441457
def__aiter__(
@@ -1552,6 +1565,16 @@ class AgentRunResult(Generic[ResultDataT]):
15521565
_result_tool_name:str|None=dataclasses.field(repr=False)
15531566
_state:_agent_graph.GraphAgentState=dataclasses.field(repr=False)
15541567
_new_message_index:int=dataclasses.field(repr=False)
1568+
_span_value:AbstractSpan|None=dataclasses.field(repr=False)
1569+
1570+
@overload
1571+
def_span(self,*,required:Literal[False])->AbstractSpan|None: ...
1572+
@overload
1573+
def_span(self)->AbstractSpan: ...
1574+
def_span(self,*,required:bool=True)->AbstractSpan|None:
1575+
ifself._span_valueisNoneandrequired:# pragma: no cover
1576+
raiseAttributeError('Span is not available for this agent run')
1577+
returnself._span_value
15551578

15561579
def_set_result_tool_return(self,return_content:str)->list[_messages.ModelMessage]:
15571580
"""Set return content for the result tool.

‎pydantic_evals/pydantic_evals/otel/_context_in_memory_span_exporter.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
try:
1616
fromlogfire._internal.tracerimport (
17-
ProxyTracerProviderasLogfireProxyTracerProvider,# pyright: ignore[reportAssignmentType,reportPrivateImportUsage]
17+
ProxyTracerProviderasLogfireProxyTracerProvider,# pyright: ignore
1818
)
1919

2020
_LOGFIRE_IS_INSTALLED=True

‎pydantic_graph/pydantic_graph/_utils.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,14 @@
55
fromfunctoolsimportpartial
66
fromtypingimportAny,Callable,TypeVar
77

8-
fromtyping_extensionsimportParamSpec,TypeIs,get_args,get_origin
8+
fromlogfire_apiimportLogfireSpan
9+
fromopentelemetry.traceimportSpan
10+
fromtyping_extensionsimportParamSpec,TypeAlias,TypeIs,get_args,get_origin
911
fromtyping_inspectionimporttyping_objects
1012
fromtyping_inspection.introspectionimportis_union_origin
1113

14+
AbstractSpan:TypeAlias='LogfireSpan | Span'
15+
1216

1317
defget_event_loop():
1418
try:

‎pydantic_graph/pydantic_graph/graph.py

Lines changed: 52 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,16 @@
66
fromcontextlibimportAbstractContextManager,ExitStack,asynccontextmanager
77
fromdataclassesimportdataclass,field
88
fromfunctoolsimportcached_property
9-
fromtypingimportAny,Generic,cast
9+
fromtypingimportAny,Generic,cast,overload
1010

1111
importlogfire_api
1212
importtyping_extensions
13-
fromlogfire_apiimportLogfireSpan
13+
fromopentelemetry.traceimportSpan
1414
fromtyping_extensionsimportdeprecated
1515
fromtyping_inspectionimporttyping_objects
1616

1717
from .import_utils,exceptions,mermaid
18+
from ._utilsimportAbstractSpan
1819
from .nodesimportBaseNode,DepsT,End,GraphRunContext,NodeDef,RunEndT,StateT
1920
from .persistenceimportBaseStatePersistence
2021
from .persistence.in_memimportSimpleStatePersistence
@@ -125,7 +126,6 @@ async def run(
125126
deps:DepsT=None,
126127
persistence:BaseStatePersistence[StateT,RunEndT]|None=None,
127128
infer_name:bool=True,
128-
span:LogfireSpan|None=None,
129129
)->GraphRunResult[StateT,RunEndT]:
130130
"""Run the graph from a starting node until it ends.
131131
@@ -137,8 +137,6 @@ async def run(
137137
persistence: State persistence interface, defaults to
138138
[`SimpleStatePersistence`][pydantic_graph.SimpleStatePersistence] if `None`.
139139
infer_name: Whether to infer the graph name from the calling frame.
140-
span: The span to use for the graph run. If not provided, a span will be created depending on the value of
141-
the `auto_instrument` field.
142140
143141
Returns:
144142
A `GraphRunResult` containing information about the run, including its final result.
@@ -164,7 +162,7 @@ async def main():
164162
self._infer_name(inspect.currentframe())
165163

166164
asyncwithself.iter(
167-
start_node,state=state,deps=deps,persistence=persistence,span=span,infer_name=False
165+
start_node,state=state,deps=deps,persistence=persistence,infer_name=False
168166
)asgraph_run:
169167
asyncfor_nodeingraph_run:
170168
pass
@@ -214,7 +212,7 @@ async def iter(
214212
state:StateT=None,
215213
deps:DepsT=None,
216214
persistence:BaseStatePersistence[StateT,RunEndT]|None=None,
217-
span:AbstractContextManager[Any]|None=None,
215+
span:AbstractContextManager[Span]|None=None,
218216
infer_name:bool=True,
219217
)->AsyncIterator[GraphRun[StateT,DepsT,RunEndT]]:
220218
"""A contextmanager which can be used to iterate over the graph's nodes as they are executed.
@@ -252,14 +250,15 @@ async def iter(
252250
persistence=SimpleStatePersistence()
253251
persistence.set_graph_types(self)
254252

255-
ifself.auto_instrumentandspanisNone:
256-
span=logfire_api.span('run graph {graph.name}',graph=self)
257-
258253
withExitStack()asstack:
259-
ifspanisnotNone:
260-
stack.enter_context(span)
254+
entered_span:AbstractSpan|None=None
255+
ifspanisNone:
256+
ifself.auto_instrument:
257+
entered_span=stack.enter_context(logfire_api.span('run graph {graph.name}',graph=self))
258+
else:
259+
entered_span=stack.enter_context(span)
261260
yieldGraphRun[StateT,DepsT,RunEndT](
262-
graph=self,start_node=start_node,persistence=persistence,state=state,deps=deps
261+
graph=self,start_node=start_node,persistence=persistence,state=state,deps=deps,span=entered_span
263262
)
264263

265264
@asynccontextmanager
@@ -268,7 +267,7 @@ async def iter_from_persistence(
268267
persistence:BaseStatePersistence[StateT,RunEndT],
269268
*,
270269
deps:DepsT=None,
271-
span:AbstractContextManager[Any]|None=None,
270+
span:AbstractContextManager[AbstractSpan]|None=None,
272271
infer_name:bool=True,
273272
)->AsyncIterator[GraphRun[StateT,DepsT,RunEndT]]:
274273
"""A contextmanager to iterate over the graph's nodes as they are executed, created from a persistence object.
@@ -301,15 +300,15 @@ async def iter_from_persistence(
301300
span=logfire_api.span('run graph {graph.name}',graph=self)
302301

303302
withExitStack()asstack:
304-
ifspanisnotNone:
305-
stack.enter_context(span)
303+
entered_span=NoneifspanisNoneelsestack.enter_context(span)
306304
yieldGraphRun[StateT,DepsT,RunEndT](
307305
graph=self,
308306
start_node=snapshot.node,
309307
persistence=persistence,
310308
state=snapshot.state,
311309
deps=deps,
312310
snapshot_id=snapshot.id,
311+
span=entered_span,
313312
)
314313

315314
asyncdefinitialize(
@@ -370,6 +369,7 @@ async def next(
370369
persistence=persistence,
371370
state=state,
372371
deps=deps,
372+
span=None,
373373
)
374374
returnawaitrun.next(node)
375375

@@ -644,6 +644,7 @@ def __init__(
644644
persistence:BaseStatePersistence[StateT,RunEndT],
645645
state:StateT,
646646
deps:DepsT,
647+
span:AbstractSpan|None,
647648
snapshot_id:str|None=None,
648649
):
649650
"""Create a new run for a given graph, starting at the specified node.
@@ -658,6 +659,7 @@ def __init__(
658659
to all nodes via `ctx.state`.
659660
deps: Optional dependencies that each node can access via `ctx.deps`, e.g. database connections,
660661
configuration, or logging clients.
662+
span: The span used for the graph run.
661663
snapshot_id: The ID of the snapshot the node came from.
662664
"""
663665
self.graph=graph
@@ -666,9 +668,19 @@ def __init__(
666668
self.state=state
667669
self.deps=deps
668670

671+
self.__span=span
669672
self._next_node:BaseNode[StateT,DepsT,RunEndT]|End[RunEndT]=start_node
670673
self._is_started:bool=False
671674

675+
@overload
676+
def_span(self,*,required:typing_extensions.Literal[False])->AbstractSpan|None: ...
677+
@overload
678+
def_span(self)->AbstractSpan: ...
679+
def_span(self,*,required:bool=True)->AbstractSpan|None:
680+
ifself.__spanisNoneandrequired:# pragma: no cover
681+
raiseexceptions.GraphRuntimeError('No span available for this graph run.')
682+
returnself.__span
683+
672684
@property
673685
defnext_node(self)->BaseNode[StateT,DepsT,RunEndT]|End[RunEndT]:
674686
"""The next node that will be run in the graph.
@@ -682,10 +694,8 @@ def result(self) -> GraphRunResult[StateT, RunEndT] | None:
682694
"""The final result of the graph run if the run is completed, otherwise `None`."""
683695
ifnotisinstance(self._next_node,End):
684696
returnNone# The GraphRun has not finished running
685-
returnGraphRunResult(
686-
self._next_node.data,
687-
state=self.state,
688-
persistence=self.persistence,
697+
returnGraphRunResult[StateT,RunEndT](
698+
self._next_node.data,state=self.state,persistence=self.persistence,span=self._span(required=False)
689699
)
690700

691701
asyncdefnext(
@@ -793,10 +803,31 @@ def __repr__(self) -> str:
793803
returnf'<GraphRun graph={self.graph.nameor"[unnamed]"}>'
794804

795805

796-
@dataclass
806+
@dataclass(init=False)
797807
classGraphRunResult(Generic[StateT,RunEndT]):
798808
"""The final result of running a graph."""
799809

800810
output:RunEndT
801811
state:StateT
802812
persistence:BaseStatePersistence[StateT,RunEndT]=field(repr=False)
813+
814+
def__init__(
815+
self,
816+
output:RunEndT,
817+
state:StateT,
818+
persistence:BaseStatePersistence[StateT,RunEndT],
819+
span:AbstractSpan|None=None,
820+
):
821+
self.output=output
822+
self.state=state
823+
self.persistence=persistence
824+
self.__span=span
825+
826+
@overload
827+
def_span(self,*,required:typing_extensions.Literal[False])->AbstractSpan|None: ...
828+
@overload
829+
def_span(self)->AbstractSpan: ...
830+
def_span(self,*,required:bool=True)->AbstractSpan|None:# pragma: no cover
831+
ifself.__spanisNoneandrequired:
832+
raiseexceptions.GraphRuntimeError('No span available for this graph run.')
833+
returnself.__span

0 commit comments

Comments
 (0)

[8]ページ先頭

©2009-2025 Movatter.jp