From 98041ecade8288c1733960eedf777c8871271803 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 8 Sep 2026 14:27:47 -0400 Subject: [PATCH] fix(pydantic-ai): avoid duplicate request customization Capture tool metadata from the normal `prepare_request` result instead of calling `customize_request_parameters` during span construction. This preserves provider-adjusted tool schemas while avoiding side effects and divergent trace metadata. Add VCR-backed coverage for direct and streaming requests with a stateful customization hook. --- .../test_pydantic_ai_integration.py | 29 ++++++++++++--- .../integrations/pydantic_ai/tracing.py | 36 ++++++++++++------- 2 files changed, 49 insertions(+), 16 deletions(-) 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 d2fb0074..76430cd9 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 @@ -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() diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index e395ee61..c704a6f8 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/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