From 48ee55bd129b3f2b92a31933a7ee11910cec7e58 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 4 Sep 2026 14:36:28 -0400 Subject: [PATCH 1/2] feat(pydantic_ai): capture tools in llm metadata Extract function and output tool definitions from model request parameters and normalize them to the backend-compatible OpenAI shape. Extend existing VCR coverage for tool execution and structured output across supported versions. --- .../test_pydantic_ai_integration.py | 29 ++++++++++++ .../integrations/pydantic_ai/tracing.py | 46 +++++++++++++++++++ 2 files changed, 75 insertions(+) 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 42f4052a..9bad2080 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 @@ -906,6 +906,17 @@ 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"} + assert output_tool["function"]["parameters"]["properties"]["answer"]["type"] == "integer" _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 +1903,24 @@ 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"} + assert "operation" in calculate_tool["function"]["parameters"]["properties"] + # 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" diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 3058d104..5de36c20 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -350,12 +350,55 @@ 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") + if not isinstance(description, str) or not description: + description = f"Tool: {name}" + + 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": []} + + return { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": parameters, + }, + } + + +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 +409,9 @@ 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 + tools = _extract_model_request_tools(model_request_parameters) + if tools: + metadata["tools"] = tools return model_name, display_name, input_data, metadata From 9e6e692451b17cd9ecba4c1659599f1e95950d77 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 4 Sep 2026 15:12:10 -0400 Subject: [PATCH 2/2] pr review comments --- .../test_pydantic_ai_integration.py | 13 ++++++++-- .../integrations/pydantic_ai/tracing.py | 25 +++++++++++-------- 2 files changed, 26 insertions(+), 12 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 9bad2080..9946c6e9 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 @@ -915,8 +915,9 @@ def is_descendant(child_span, ancestor_id): None, ) assert output_tool is not None - assert set(output_tool["function"]) == {"name", "description", "parameters"} + 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). @@ -1918,8 +1919,9 @@ def calculate(operation: str, a: float, b: float) -> str: ) 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"} + 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" @@ -1961,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") diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 5de36c20..ddad6564 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/py/src/braintrust/integrations/pydantic_ai/tracing.py @@ -356,8 +356,6 @@ def _shape_model_tool_definition(tool: Any) -> dict[str, Any] | None: return None description = _field_value(tool, "description") - if not isinstance(description, str) or not description: - description = f"Tool: {name}" parameters = _field_value(tool, "parameters_json_schema") if parameters is _MISSING: @@ -365,14 +363,14 @@ def _shape_model_tool_definition(tool: Any) -> dict[str, Any] | None: if parameters is _MISSING: parameters = {"type": "object", "properties": {}, "required": []} - return { - "type": "function", - "function": { - "name": name, - "description": description, - "parameters": parameters, - }, - } + 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]: @@ -409,6 +407,13 @@ 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