fix(pydantic_ai): capture tools in llm metadata

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

@@ -906,6 +906,18 @@ def is_descendant(child_span, ancestor_id):
     assert is_descendant(chat_span, agent_span["span_id"]), "chat span should be nested under agent_run"
     assert chat_span["metadata"]["model"] == "gpt-4o-mini"
     assert chat_span["metadata"]["provider"] == "openai"
+    output_tool = next(
+        (
+            tool
+            for tool in chat_span["metadata"]["tools"]
+            if tool.get("type") == "function" and tool.get("function", {}).get("name") == "final_result"
+        ),
+        None,
+    )
+    assert output_tool is not None
+    assert set(output_tool["function"]) == {"name", "description", "parameters", "strict"}
+    assert output_tool["function"]["parameters"]["properties"]["answer"]["type"] == "integer"
+    assert output_tool["function"]["strict"] is True
     _assert_metrics_are_valid(chat_span["metrics"], start, end)
 
     # Wrapper agent_run span must not log token metrics (would double-count at rollup).
@@ -1892,6 +1904,25 @@ def calculate(operation: str, a: float, b: float) -> str:
     tool_names = [t["name"] for t in tools if isinstance(t, dict)]
     assert "calculate" in tool_names, f"calculate tool should be in tools list, got: {tool_names}"
 
+    # Tool definitions passed to each leaf model call belong in metadata.tools.
+    chat_spans = [s for s in spans if "chat" in s["span_attributes"]["name"]]
+    assert chat_spans, "chat span not found"
+    for chat_span in chat_spans:
+        model_tools = chat_span["metadata"]["tools"]
+        calculate_tool = next(
+            (
+                tool
+                for tool in model_tools
+                if tool.get("type") == "function" and tool.get("function", {}).get("name") == "calculate"
+            ),
+            None,
+        )
+        assert calculate_tool is not None, f"calculate tool should be in chat metadata.tools, got: {model_tools}"
+        assert set(calculate_tool) == {"type", "function"}
+        assert set(calculate_tool["function"]) == {"name", "description", "parameters", "strict"}
+        assert "operation" in calculate_tool["function"]["parameters"]["properties"]
+        assert calculate_tool["function"]["strict"] is True
+
     # Verify toolsets are NOT in metadata (following the principle: agent.run() accepts it)
     assert "toolsets" not in agent_span["metadata"], "toolsets should NOT be in metadata"
 
@@ -1932,9 +1963,16 @@ def get_weather(city: str) -> str:
     spans = memory_logger.pop()
     agent_span = next((s for s in spans if "agent_run" in s["span_attributes"]["name"]), None)
     tool_span = next((s for s in spans if s["span_attributes"].get("name") == "get_weather"), None)
+    chat_spans = [s for s in spans if "chat" in s["span_attributes"]["name"]]
 
     assert agent_span is not None, "agent_run span not found"
     assert tool_span is not None, "runtime tool span not found"
+    assert chat_spans, "chat span not found"
+    for chat_span in chat_spans:
+        weather_tool = next(
+            tool for tool in chat_span["metadata"]["tools"] if tool["function"]["name"] == "get_weather"
+        )
+        assert weather_tool["function"].get("description") in (None, "")
     assert tool_span["span_attributes"]["type"] == SpanTypeAttribute.TOOL
     assert tool_span["span_parents"] == [agent_span["span_id"]]
     assert tool_span["metadata"].get("tool_call_id")

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

@@ -350,12 +350,53 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def _shape_model_tool_definition(tool: Any) -> dict[str, Any] | None:
+    name = _field_value(tool, "name")
+    if not isinstance(name, str):
+        return None
+
+    description = _field_value(tool, "description")
+
+    parameters = _field_value(tool, "parameters_json_schema")
+    if parameters is _MISSING:
+        parameters = _field_value(tool, "parameters")
+    if parameters is _MISSING:
+        parameters = {"type": "object", "properties": {}, "required": []}
+
+    function = {"name": name, "parameters": parameters}
+    if description is not _MISSING and description is not None:
+        function["description"] = description
+    strict = _field_value(tool, "strict")
+    if strict is not _MISSING and strict is not None:
+        function["strict"] = strict
+
+    return {"type": "function", "function": function}
+
+
+def _extract_model_request_tools(model_request_parameters: Any) -> list[Any]:
+    if model_request_parameters is None:
+        return []
+
+    tools = []
+    for field in ("function_tools", "output_tools"):
+        definitions = _field_value(model_request_parameters, field)
+        if definitions is _MISSING or not definitions:
+            continue
+        for definition in definitions:
+            shaped_definition = _shape_model_tool_definition(definition)
+            if shaped_definition is not None:
+                tools.append(shaped_definition)
+
+    return tools
+
+
 def _build_model_class_input_and_metadata(instance: Any, args: Any, kwargs: Any):
     model_name, provider = _extract_model_info_from_model_instance(instance)
     display_name = model_name or type(instance).__name__
 
     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)
 
@@ -366,6 +407,16 @@ 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