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..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 @@ -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") diff --git a/py/src/braintrust/integrations/pydantic_ai/tracing.py b/py/src/braintrust/integrations/pydantic_ai/tracing.py index 3058d104..ddad6564 100644 --- a/py/src/braintrust/integrations/pydantic_ai/tracing.py +++ b/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