fix(pydantic_ai): [MLOB-3476] fix tool call tracing in pydantic-ai >= 0.4.4 [backport 3.11] #14236
Conversation
![dd-octo-sts[bot]](https://avatars.githubusercontent.com/in/1157446?s=60&v=4){kind=small}
|
|
Bootstrap import analysisComparison of import times between this PR and base. SummaryThe average import time from this PR is: 277 ± 8 ms. The average import time from base is: 280 ± 9 ms. The import time difference between this PR and base is: -3.3 ± 0.4 ms. Import time breakdownThe following import paths have shrunk:
|
Performance SLOsCandidate: backport-14171-to-3.11 (31633b6) 🔵 No Baseline Data (24 suites)🔵 coreapiscenario - 12/12 (2 unstable)🔵 No baseline data available for this suite
|
… 0.4.4 (#14171) We received this issue where upgrading to `pydantic-ai-slim >= 0.4.4` breaks tool call tracing for Pydantic AI since the traced tool call method (`Tool.run`) was removed in this version due to the new [Toolsets](https://ai.pydantic.dev/toolsets/) interface being introduced. This PR fixes that issue by conditionally tracing either the `run` method of the Tool class for versions < 0.4.4 or tracing the `ToolManager.handle_call` in newer versions of Pydantic AI. Also, some minor changes were made to support annotating tools defined via toolsets for the agent manifest. I first ensured that conditionally patching `Tool.run` resolved the original error message where we were crashing during startup while trying to enable the Pydantic AI integration. I did this by making sure I had version `pydantic-ai>=0.4.4` and then running the following script with `ddtrace-run python startup.py`. ``` import pydantic_ai from ddtrace.llmobs import LLMObs LLMObs.enable() ``` This resulted in no error message. I then wanted to check that we indeed trace tool calls made either via regular tools or via toolsets. I first used the normal `tools` argument to register my tool and invoke it. ``` from pydantic_ai import Agent def calculate_square(x: int) -> int: return x * x agent = Agent( 'openai:gpt-4o', system_prompt='Use one of the tools provided to calculate the mathematical operation.', tools=[calculate_square] ) result = agent.run_sync('What is the square of 5?') print(result.output) ``` This resulted in this [trace](https://app.datadoghq.com/llm/traces?query=%40ml_app%3Anicole-test%20%40event_type%3Aspan%20%40parent_id%3Aundefined&agg_m=count&agg_m_source=base&agg_t=count&fromUser=true&llmPanels=%5B%7B%22t%22%3A%22sampleDetailPanel%22%2C%22rEID%22%3A%22AwAAAZhcffsFoQctBAAAABhBWmhjZmZzRkFBRDhQWGUzdTMwbEFBQUEAAAAkMTE5ODVjN2UtMGRiMC00MGQ1LWE3ZjAtZDk2ZmIwZjE5NTQxAAAAWQ%22%7D%5D&spanId=14197137118701788221&start=1753897528617&end=1753898428617&paused=false) where you can see the `calculate_square_tool` being invoked and also see the tool definition in the agent manifest. I then tried to use the toolsets argument by modifying the script slightly. ``` from pydantic_ai import Agent from pydantic_ai.toolsets import FunctionToolset def calculate_square(x: int) -> int: return x * x toolset = FunctionToolset(tools=[calculate_square]) agent = Agent( 'openai:gpt-4o', system_prompt='Use one of the tools provided to calculate the mathematical operation.', toolsets=[toolset] ) result = agent.run_sync('What is the square of 5?') print(result.output) ``` This resulted in this [trace](https://app.datadoghq.com/llm/traces?query=%40ml_app%3Anicole-test%20%40event_type%3Aspan%20%40parent_id%3Aundefined&agg_m=count&agg_m_source=base&agg_t=count&fromUser=true&llmPanels=%5B%7B%22t%22%3A%22sampleDetailPanel%22%2C%22rEID%22%3A%22AwAAAZhcga3TRDnHrwAAABhBWmhjZ2EzVEFBQ1dfOU1EX09rTkFBQUEAAAAkZjE5ODVjODEtY2M1Yi00ZjNiLTg5NzYtZDJiNTlmMDE5ODBmAAAAqw%22%7D%5D&spanId=4209762990183049031&start=1753897790055&end=1753898690055&paused=false) where again you see the tool call being traced along with the tool definition in the agent manifest. - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) (cherry picked from commit efcbf0d) Signed-off-by: Nicole Cybul <[email protected]>
backport of https://github.com/DataDog/dd-trace-py/pull/14233/files to unlock 3.11 CI ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
ncybul
force-pushed
the
backport-14171-to-3.11
branch
from
a444862 to
f0e397b
Compare
Backport efcbf0d from #14171 to 3.11.
We received this issue where upgrading to
pydantic-ai-slim >= 0.4.4breaks tool call tracing for Pydantic AI since the traced tool call method (Tool.run) was removed in this version due to the new Toolsets interface being introduced. This PR fixes that issue by conditionally tracing either therunmethod of the Tool class for versions < 0.4.4 or tracing theToolManager.handle_callin newer versions of Pydantic AI.Also, some minor changes were made to support annotating tools defined via toolsets for the agent manifest.
Manual Testing
I first ensured that conditionally patching
Tool.runresolved the original error message where we were crashing during startup while trying to enable the Pydantic AI integration. I did this by making sure I had versionpydantic-ai>=0.4.4and then running the following script withddtrace-run python startup.py.This resulted in no error message.
I then wanted to check that we indeed trace tool calls made either via regular tools or via toolsets.
I first used the normal
toolsargument to register my tool and invoke it.This resulted in this trace where you can see the
calculate_square_toolbeing invoked and also see the tool definition in the agent manifest.I then tried to use the toolsets argument by modifying the script slightly.
This resulted in this trace where again you see the tool call being traced along with the tool definition in the agent manifest.
Checklist
Reviewer Checklist