test

Author: alexmojaki

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: ...
 

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):