diff --git a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py b/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py index cd3da225..8cc7d49b 100644 --- a/py/src/braintrust/integrations/pydantic_ai/test_pydantic_ai_integration.py +++ b/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 diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 1003fd35..3c7a1bc0 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/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: