fix(pydantic-ai): preserve unshaped message fields

py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py

@@ -1923,7 +1923,13 @@ def test_shape_message_and_response_pass_through_when_no_binary():
 
     from braintrust.integrations.pydantic_ai.tracing import _shape_message, _shape_model_response
 
-    message = SimpleNamespace(kind="request", state="complete", parts=["hello"])
+    message = SimpleNamespace(
+        kind="request",
+        state="complete",
+        instructions="Be concise.",
+        parts=[SimpleNamespace(part_kind="system-prompt", content="Use the supplied context."), "hello"],
+        new_pydantic_ai_field="preserve me",
+    )
     assert _shape_message(message) is message
 
     response = SimpleNamespace(
@@ -1936,6 +1942,56 @@ def test_shape_message_and_response_pass_through_when_no_binary():
     assert _shape_model_response(response) is response
 
 
+def test_conversation_messages_serialize_same_with_or_without_adapter_dump():
+    """Raw Pydantic AI messages log the same structure as adapter-dumped messages."""
+    from braintrust.bt_json import bt_safe_deep_copy
+    from braintrust.integrations.pydantic_ai.tracing import _shape_messages
+    from pydantic_ai.messages import (
+        ModelMessagesTypeAdapter,
+        ModelRequest,
+        ModelResponse,
+        SystemPromptPart,
+        TextPart,
+        ToolCallPart,
+        UserPromptPart,
+    )
+
+    messages = [
+        ModelRequest(
+            instructions="Be concise.",
+            parts=[
+                SystemPromptPart(content="Use the supplied context."),
+                UserPromptPart(content="Find the current weather."),
+            ],
+        ),
+        ModelResponse(
+            parts=[
+                TextPart(content="I will check."),
+                ToolCallPart(tool_name="weather", args={"city": "Toronto"}, tool_call_id="call_1"),
+            ]
+        ),
+    ]
+
+    raw_shape = bt_safe_deep_copy(_shape_messages(messages))
+    dumped_shape = bt_safe_deep_copy(_shape_messages(ModelMessagesTypeAdapter.dump_python(messages)))
+
+    assert raw_shape == dumped_shape
+
+    def assert_same_key_sets(raw, dumped):
+        if isinstance(raw, dict):
+            assert isinstance(dumped, dict)
+            assert raw.keys() == dumped.keys()
+            for key in raw:
+                assert_same_key_sets(raw[key], dumped[key])
+        elif isinstance(raw, list):
+            assert isinstance(dumped, list)
+            assert len(raw) == len(dumped)
+            for raw_item, dumped_item in zip(raw, dumped):
+                assert_same_key_sets(raw_item, dumped_item)
+
+    assert_same_key_sets(raw_shape, dumped_shape)
+
+
 def test_v2_model_provider_inference():
     from braintrust.integrations.pydantic_ai.tracing import _extract_model_info_from_model_instance
 

py/src/braintrust/integrations/pydantic_ai/tracing.py

@@ -1171,13 +1171,9 @@ def _shape_message(message: Any) -> Any:
     parts = _field_value(message, "parts")
     if not parts:
         return message
-    has_prompt_data = _field_value(message, "instructions") not in (_MISSING, None, "") or any(
-        _field_value(part, "part_kind") == "system-prompt" for part in parts
-    )
-    if not has_prompt_data and not any(_has_binary_leaf(part) for part in parts):
-        # Let Braintrust's serializer preserve every top-level field. It does not recurse
-        # into dataclasses nested in a list, so prompt-bearing parts have to be shaped
-        # here or they reach the span as repr strings; binary leaves need materializing.
+    if not any(_has_binary_leaf(part) for part in parts):
+        # Let Braintrust's serializer preserve every field unless binary content needs
+        # materializing. In particular, this keeps newer pydantic_ai fields intact.
         return message
     return _shape_object(
         message, fields=_MESSAGE_FIELDS, overrides={"parts": [_shape_content_part(part) for part in parts]}
@@ -1573,16 +1569,12 @@ def _build_agent_input_and_metadata(args: Any, kwargs: Any, instance: Any) -> tu
     conversation = kwargs.get("conversation")
     if conversation is not None:
         try:
-            from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-            # Nested message parts otherwise serialize as repr strings.
-            messages = ModelMessagesTypeAdapter.dump_python(_field_value(conversation, "messages"))
             # Use the existing normalizer fields after delegated kwargs have logged their None defaults.
             input_data.update(
                 _shape_object(
                     conversation,
                     fields=("conversation_id",),
-                    overrides={"message_history": _shape_messages(messages)},
+                    overrides={"message_history": _shape_messages(_field_value(conversation, "messages"))},
                 )
             )
         except Exception as e: