fix(pydantic-ai): avoid duplicate request customization

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

@@ -45,6 +45,23 @@ def memory_logger():
         yield bgl
 
 
+def _model_with_customization_counter():
+    from pydantic_ai.models import infer_model
+
+    inferred_model = infer_model(MODEL)
+    customization_calls = []
+
+    class SideEffectModel(type(inferred_model)):
+        def __str__(self):
+            return MODEL
+
+        def customize_request_parameters(self, model_request_parameters):
+            customization_calls.append(model_request_parameters)
+            return super().customize_request_parameters(model_request_parameters)
+
+    return SideEffectModel(inferred_model.model_name), customization_calls
+
+
 def _assert_metrics_are_valid(metrics, start, end):
     """Assert that metrics contain expected fields and values."""
     assert "start" in metrics
@@ -648,18 +665,20 @@ def get_weather(city: str) -> str:
 @pytest.mark.vcr
 @pytest.mark.asyncio
 async def test_direct_model_request(memory_logger, direct):
-    """Test direct API model_request()."""
+    """Test direct API model_request() without tracing changing model preparation."""
     assert not memory_logger.pop()
 
+    model, customization_calls = _model_with_customization_counter()
     messages = [ModelRequest(parts=[UserPromptPart(content=TEST_PROMPT)])]
 
     start = time.time()
-    response = await direct.model_request(model=MODEL, messages=messages)
+    response = await direct.model_request(model=model, messages=messages)
     end = time.time()
 
     # Verify response
     assert response.parts
     assert "4" in str(response.parts[0].content)
+    assert len(customization_calls) == 1
 
     # Check spans
     spans = memory_logger.pop()
@@ -759,17 +778,19 @@ async def test_direct_model_request_stream(memory_logger, direct):
     """Test direct API model_request_stream() - verifies time_to_first_token is captured."""
     assert not memory_logger.pop()
 
+    model, customization_calls = _model_with_customization_counter()
     messages = [ModelRequest(parts=[UserPromptPart(content="Count from 1 to 3")])]
 
     start = time.time()
     chunk_count = 0
-    async with direct.model_request_stream(model=MODEL, messages=messages) as stream:
+    async with direct.model_request_stream(model=model, messages=messages) as stream:
         async for chunk in stream:
             chunk_count += 1
     end = time.time()
 
-    # Verify we got chunks
+    # Verify we got chunks and tracing did not rerun request customization.
     assert chunk_count > 0
+    assert len(customization_calls) == 1
 
     # Check spans
     spans = memory_logger.pop()

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

@@ -9,7 +9,7 @@
 from typing import Any
 
 from braintrust.integrations.utils import _materialize_attachment
-from braintrust.logger import _internal_get_global_state
+from braintrust.logger import _internal_get_global_state, current_span
 from braintrust.logger import start_span as _bt_start_span
 
 
@@ -429,7 +429,6 @@ def _build_model_class_input_and_metadata(instance: Any, args: Any, kwargs: Any)
 
     messages = args[0] if len(args) > 0 else kwargs.get("messages")
     model_settings = args[1] if len(args) > 1 else kwargs.get("model_settings")
-    model_request_parameters = args[2] if len(args) > 2 else kwargs.get("model_request_parameters")
 
     shaped_messages = _shape_messages(messages)
 
@@ -440,20 +439,31 @@ def _build_model_class_input_and_metadata(instance: Any, args: Any, kwargs: Any)
     metadata = _build_model_metadata(model_name, provider, model_settings=None)
     if model_settings is not None:
         metadata["invocation_params"] = model_settings
-    # Provider customization resolves inferred strictness and schema transformations used on the wire.
-    customize_request_parameters = getattr(instance, "customize_request_parameters", None)
-    if model_request_parameters is not None and callable(customize_request_parameters):
-        try:
-            model_request_parameters = customize_request_parameters(model_request_parameters)
-        except Exception as e:
-            logger.debug(f"Failed to customize model request parameters for tracing: {e}")
-    tools = _extract_model_request_tools(model_request_parameters)
-    if tools:
-        metadata["tools"] = tools
 
     return model_name, display_name, input_data, metadata
 
 
+def _model_prepare_request_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
+    prepared = wrapped(*args, **kwargs)
+    span = current_span()
+    if getattr(span, "_instrumentation", None) != _INSTRUMENTATION or not span.name.startswith("chat "):
+        return prepared
+
+    try:
+        model_request_parameters = prepared[1] if isinstance(prepared, tuple) and len(prepared) > 1 else None
+        tools = _extract_model_request_tools(model_request_parameters)
+    except Exception as e:
+        logger.debug(f"Failed to extract prepared model request parameters for tracing: {e}")
+        return prepared
+
+    if tools:
+        # prepare_request() runs inside the traced request/request_stream call, so
+        # this captures Pydantic AI's normal customization result without invoking
+        # a potentially stateful customization hook a second time.
+        span.log(metadata={"tools": tools})
+    return prepared
+
+
 def _wrap_concrete_model_class(model_class: Any):
     async def model_request_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
         model_name, display_name, input_data, metadata = _build_model_class_input_and_metadata(instance, args, kwargs)
@@ -486,6 +496,8 @@ def model_request_stream_wrapper(wrapped: Any, instance: Any, args: Any, kwargs:
 
     wrap_function_wrapper(model_class, "request", model_request_wrapper)
     wrap_function_wrapper(model_class, "request_stream", model_request_stream_wrapper)
+    if hasattr(model_class, "prepare_request"):
+        wrap_function_wrapper(model_class, "prepare_request", _model_prepare_request_wrapper)
     return model_class