Formatting

Author: WorldInnovationsDepartment

pydantic_ai_slim/pydantic_ai/agent.py

@@ -8,7 +8,7 @@
 from types import FrameType
 from typing import Any, Callable, Generic, cast, final, overload
 
-from opentelemetry.trace import NoOpTracer, use_span
+import logfire_api
 from typing_extensions import TypeGuard, TypeVar, deprecated
 
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
@@ -58,6 +58,17 @@
     'UserPromptNode',
 )
 
+_logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
+
+# while waiting for https://github.com/pydantic/logfire/issues/745
+try:
+    import logfire._internal.stack_info
+except ImportError:
+    pass
+else:
+    from pathlib import Path
+
+    logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
 T = TypeVar('T')
 S = TypeVar('S')
@@ -112,9 +123,6 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     The type of the result data, used to validate the result data, defaults to `str`.
     """
 
-    instrument: bool
-    """Automatically instrument with OpenTelemetry. Will use Logfire if it's configured."""
-
     _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _reflect_on_tool_call: bool = dataclasses.field(repr=False)
     _result_tool_name: str = dataclasses.field(repr=False)
@@ -149,7 +157,6 @@ def __init__(
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         reflect_on_tool_call: bool = True,
-        instrument: bool = False,
     ):
         """Create an agent.
 
@@ -182,7 +189,6 @@ def __init__(
             reflect_on_tool_call: Whether to generate a final response with an LLM call after tool calls.
                 When set to `False`, the agent will return tool call results directly without generating
                 a final reflection response. Defaults to `True`.
-            instrument: Automatically instrument with OpenTelemetry. Will use Logfire if it's configured.
         """
         if model is None or defer_model_check:
             self.model = model
@@ -193,7 +199,6 @@ def __init__(
         self.name = name
         self.model_settings = model_settings
         self.result_type = result_type
-        self.instrument = instrument
 
         self._deps_type = deps_type
 
@@ -397,7 +402,6 @@ async def main():
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
         model_used = self._get_model(model)
-        del model
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
@@ -427,18 +431,13 @@ async def main():
         model_settings = merge_model_settings(self.model_settings, model_settings)
         usage_limits = usage_limits or _usage.UsageLimits()
 
-        if isinstance(model_used, InstrumentedModel):
-            tracer = model_used.tracer
-        else:
-            tracer = NoOpTracer()
-        agent_name = self.name or 'agent'
-        run_span = tracer.start_span(
-            'agent run',
-            attributes={
-                'model_name': model_used.model_name if model_used else 'no-model',
-                'agent_name': agent_name,
-                'logfire.msg': f'{agent_name} run',
-            },
+        # Build the deps object for the graph
+        run_span = _logfire.span(
+            '{agent_name} run {prompt=}',
+            prompt=user_prompt,
+            agent=self,
+            model_name=model_used.model_name if model_used else 'no-model',
+            agent_name=self.name or 'agent',
         )
         graph_deps = self._build_graph_deps(
             deps,
@@ -463,7 +462,7 @@ async def main():
             state=state,
             deps=graph_deps,
             infer_name=False,
-            span=use_span(run_span, end_on_exit=True),
+            span=run_span,
         ) as graph_run:
             yield AgentRun(graph_run)
 
@@ -1119,7 +1118,7 @@ def _get_model(self, model: models.Model | models.KnownModelName | None) -> mode
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        if self.instrument and not isinstance(model_, InstrumentedModel):
+        if not isinstance(model_, InstrumentedModel):
             model_ = InstrumentedModel(model_)
 
         return model_

tests/test_reflect_on_tool_call.py

@@ -1,6 +1,5 @@
 """Tests for reflect_on_tool_call functionality."""
 
-
 import pytest
 from pydantic import BaseModel