Skip to content

Validate common AI tool-call arguments to enable pydantic-ai retries#70096

Open
zozo123 wants to merge 5 commits into
apache:mainfrom
zozo123:fix-toolset-args-validation
Open

Validate common AI tool-call arguments to enable pydantic-ai retries#70096
zozo123 wants to merge 5 commits into
apache:mainfrom
zozo123:fix-toolset-args-validation

Conversation

@zozo123

@zozo123 zozo123 commented Jul 19, 2026

Copy link
Copy Markdown

Malformed model-generated arguments currently bypass validation in the Common AI toolsets. Missing or mistyped arguments can therefore reach tool implementations and fail with KeyError or TypeError instead of producing the bounded retry that pydantic-ai supports.

Invalid payloads are now rejected at validation; pydantic-ai then regenerates a new tool call — the original malformed payload is never re-executed.

Concretely: if the model calls the query tool but omits the required sql argument, the call previously crashed the run with KeyError: 'sql'; it now raises a ValidationError that pydantic-ai turns into one bounded retry, and the model regenerates the call with sql supplied.

This change derives each tool's argument validator from the same JSON Schema advertised to the model and applies it to SQLToolset, DataFusionToolset, and HookToolset. Hook introspection keeps optional and union parameters accurate, leaves unsupported annotations permissive, drops unexpected arguments for fixed signatures, and preserves additional arguments for methods that accept **kwargs.

Validation:

  • 1,043 Common AI unit tests passed
  • provider mypy passed
  • pre-commit and manual prek checks passed

Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Fable 5) and Codex (GPT-5)

Generated-by: Claude Code (Fable 5) and Codex (GPT-5) following the guidelines

Malformed model calls currently reach tool implementations and can fail the task with KeyError or TypeError instead of giving the model a bounded retry.
@zozo123
zozo123 marked this pull request as ready for review July 19, 2026 12:13

@SameerMesiah97 SameerMesiah97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the title is a bit misleading. It appears that the malformed payload is not retried, but rather rejected during validation, allowing Pydantic AI to regenerate a new tool call. "Retry malformed tool calls" initially reads as though the original invalid payload is executed again, whereas the retry is actually another model generation. It might be worth clarifying this in the title and PR description.

This looks fine to me logically (and the tests are solid) but this will need someone who is more familiar with this provider to deeply vet the validation logic. I see a lot of small changes scattered throughout the diff that I am not sure of how to ascertain the impact of.

CI needs to be triggered too.

Comment thread providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py Outdated
@zozo123 zozo123 changed the title Retry malformed common AI tool calls Validate common AI tool-call arguments to enable pydantic-ai retries Jul 20, 2026
@zozo123

zozo123 commented Jul 20, 2026

Copy link
Copy Markdown
Author

Thanks @SameerMesiah97 — good catch on the title. Renamed to "Validate common AI tool-call arguments to enable pydantic-ai retries" and added a line to the description making explicit that the malformed payload is rejected at validation and pydantic-ai regenerates a fresh tool call (the original is never re-executed). Also replied inline on the nested-object question.

Unify object-fragment handling so build_args_validator and the "object"
branch of _fragment_to_core_schema share one builder. Nested objects with
declared properties are now validated recursively, matching how arrays
already recurse into their items. Objects with no "properties" key (e.g.
from dict[K, V]) still accept any dict, so untyped payloads are unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@zozo123
zozo123 requested a review from SameerMesiah97 July 20, 2026 07:40
@potiuk potiuk added the ready for maintainer review Set after triaging when all criteria pass. label Jul 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants