Update logfire.instrument_pydantic_ai() parameters (#1480)

Author: alexmojaki

logfire/_internal/integrations/pydantic_ai.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 from typing import Any, Literal
 
 from pydantic_ai import Agent
@@ -13,27 +14,47 @@
 def instrument_pydantic_ai(
     logfire_instance: Logfire,
     obj: Agent | Model | None,
+    include_binary_content: bool | None,
+    include_content: bool | None,
+    version: Literal[1, 2, 3] | None,
     event_mode: Literal['attributes', 'logs'] | None,
     **kwargs: Any,
 ) -> None | InstrumentedModel:
-    if event_mode is None:  # pragma: no branch
-        event_mode = InstrumentationSettings.event_mode
-    kwargs = dict(
-        tracer_provider=logfire_instance.config.get_tracer_provider(),
-        event_logger_provider=logfire_instance.config.get_event_logger_provider(),
-        event_mode=event_mode,
-        **kwargs,
+    # Correctly handling all past and future versions is tricky.
+    # Since we provide these rather than the user, only include them if
+    # InstrumentationSettings has such parameters to prevent errors/warnings.
+    expected_kwarg_names = inspect.signature(InstrumentationSettings.__init__).parameters
+    final_kwargs: dict[str, Any] = {
+        k: v
+        for k, v in dict(
+            tracer_provider=logfire_instance.config.get_tracer_provider(),
+            meter_provider=logfire_instance.config.get_meter_provider(),  # not in old versions
+            event_logger_provider=logfire_instance.config.get_event_logger_provider(),  # may be removed in the future
+        ).items()
+        if k in expected_kwarg_names
+    }
+
+    # Now add known parameters that the user provided explicitly.
+    # None of these have `None` as a valid value, so we assume the user didn't pass that.
+    # Include even if not in expected_kwarg_names, to let Pydantic AI produce an error or warning.
+    final_kwargs.update(
+        {
+            k: v
+            for k, v in dict(
+                include_binary_content=include_binary_content,
+                include_content=include_content,
+                version=version,
+                event_mode=event_mode,
+            ).items()
+            if v is not None
+        }
     )
-    try:
-        settings = InstrumentationSettings(
-            meter_provider=logfire_instance.config.get_meter_provider(),
-            **kwargs,
-        )
-    except TypeError:  # pragma: no cover
-        # Handle older pydantic-ai versions that do not support meter_provider.
-        settings = InstrumentationSettings(
-            **kwargs,
-        )
+
+    # Finally, add any other kwargs the user provided. These are mainly for future compatibility.
+    # They also provide an escape hatch to override the tracer/meter/event_logger providers.
+    final_kwargs.update(kwargs)
+
+    settings = InstrumentationSettings(**final_kwargs)
     if isinstance(obj, Agent):
         obj.instrument = settings
     elif isinstance(obj, Model):

logfire/_internal/main.py

@@ -988,8 +988,10 @@ def instrument_pydantic_ai(
         obj: pydantic_ai.Agent | None = None,
         /,
         *,
-        event_mode: Literal['attributes', 'logs'] = 'attributes',
         include_binary_content: bool | None = None,
+        include_content: bool | None = None,
+        version: Literal[1, 2, 3] | None = None,
+        event_mode: Literal['attributes', 'logs'] | None = None,
         **kwargs: Any,
     ) -> None: ...
 
@@ -999,8 +1001,10 @@ def instrument_pydantic_ai(
         obj: pydantic_ai.models.Model,
         /,
         *,
-        event_mode: Literal['attributes', 'logs'] = 'attributes',
         include_binary_content: bool | None = None,
+        include_content: bool | None = None,
+        version: Literal[1, 2, 3] | None = None,
+        event_mode: Literal['attributes', 'logs'] | None = None,
         **kwargs: Any,
     ) -> pydantic_ai.models.Model: ...
 
@@ -1009,8 +1013,10 @@ def instrument_pydantic_ai(
         obj: pydantic_ai.Agent | pydantic_ai.models.Model | None = None,
         /,
         *,
-        event_mode: Literal['attributes', 'logs'] | None = None,
         include_binary_content: bool | None = None,
+        include_content: bool | None = None,
+        version: Literal[1, 2, 3] | None = None,
+        event_mode: Literal['attributes', 'logs'] | None = None,
         **kwargs: Any,
     ) -> pydantic_ai.models.Model | None:
         """Instrument Pydantic AI.
@@ -1020,10 +1026,24 @@ def instrument_pydantic_ai(
                 By default, all agents are instrumented.
                 You can also pass a specific model or agent.
                 If you pass a model, a new instrumented model will be returned.
-            event_mode: See the [Pydantic AI docs](https://ai.pydantic.dev/logfire/#data-format).
-                The default is whatever the default is in your version of Pydantic AI.
-            include_binary_content: Whether to include base64 encoded binary content (e.g. images) in the events.
+            include_binary_content: Whether to include base64 encoded binary content (e.g. images) in the telemetry.
                 On by default. Requires Pydantic AI 0.2.5 or newer.
+            include_content: Whether to include prompts, completions, and tool call arguments and responses
+                in the telemetry. On by default. Requires Pydantic AI 0.3.4 or newer.
+            version: Version of the data format. This is unrelated to the Pydantic AI package version.
+                Requires Pydantic AI 0.7.5 or newer.
+                Version 1 is based on the legacy event-based OpenTelemetry GenAI spec
+                    and will be removed in a future release.
+                    The parameter `event_mode` is only relevant for version 1.
+                Version 2 uses the newer OpenTelemetry GenAI spec and stores messages in the following attributes:
+                    - `gen_ai.system_instructions` for instructions passed to the agent.
+                    - `gen_ai.input.messages` and `gen_ai.output.messages` on model request spans.
+                    - `pydantic_ai.all_messages` on agent run spans.
+                Version 3 changes the names of some attributes and spans but not the shape of the data.
+                The default version depends on Pydantic AI.
+            event_mode: The mode for emitting events in version 1.
+                If `'attributes'`, events are attached to the span as attributes.
+                If `'logs'`, events are emitted as OpenTelemetry log-based events.
             kwargs: Additional keyword arguments to pass to
                 [`InstrumentationSettings`](https://ai.pydantic.dev/api/models/instrumented/#pydantic_ai.models.instrumented.InstrumentationSettings)
                 for future compatibility.
@@ -1032,13 +1052,13 @@ def instrument_pydantic_ai(
 
         self._warn_if_not_initialized_for_instrumentation()
 
-        if include_binary_content is not None:
-            kwargs['include_binary_content'] = include_binary_content
-
         return instrument_pydantic_ai(
             self,
             obj=obj,
             event_mode=event_mode,
+            version=version,
+            include_content=include_content,
+            include_binary_content=include_binary_content,
             **kwargs,
         )
 

tests/otel_integrations/test_pydantic_ai.py

@@ -78,6 +78,21 @@ def get_model(a: Agent):
     m2 = get_model(agent2)
     assert m2 is model
 
+    # Test all known parameters
+    logfire_inst.instrument_pydantic_ai(
+        include_binary_content=False,
+        include_content=False,
+        version=1,
+        event_mode='logs',
+    )
+    m = get_model(agent2)
+    assert isinstance(m, InstrumentedModel)
+    assert m.instrumentation_settings.version == 1
+    assert not m.instrumentation_settings.include_binary_content
+    assert not m.instrumentation_settings.include_content
+    assert m.instrumentation_settings.event_mode == 'logs'
+    Agent.instrument_all(False)
+
 
 def test_invalid_instrument_pydantic_ai():
     with pytest.raises(TypeError):