apache · kaxil · Jul 30, 2026 · Jul 19, 2026 · Jul 19, 2026 · Jul 20, 2026
@@ -35,7 +35,8 @@
 from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
-from pydantic_core import SchemaValidator, core_schema
+
+from airflow.providers.common.ai.utils.tool_definition import build_args_validator
 
 if TYPE_CHECKING:
     from pydantic_ai._run_context import RunContext
@@ -44,8 +45,6 @@
 
 log = logging.getLogger(__name__)
 
-_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema())
-
 # JSON Schemas for the three DataFusion tools.
 _LIST_TABLES_SCHEMA: dict[str, Any] = {
     "type": "object",
@@ -146,7 +145,7 @@ async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
                 toolset=self,
                 tool_def=tool_def,
                 max_retries=1,
-                args_validator=_PASSTHROUGH_VALIDATOR,
+                args_validator=build_args_validator(schema),
             )
         return tools
 

@@ -26,9 +26,8 @@
 
 from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
-from pydantic_core import SchemaValidator, core_schema
 
-from airflow.providers.common.ai.utils.tool_definition import return_schema_kwargs
+from airflow.providers.common.ai.utils.tool_definition import build_args_validator, return_schema_kwargs
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -37,9 +36,6 @@
 
     from airflow.providers.common.compat.sdk import BaseHook
 
-# Single shared validator — accepts any JSON-decoded dict from the LLM.
-_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema())
-
 # Maps Python types to JSON Schema fragments.
 _TYPE_MAP: dict[type, dict[str, Any]] = {
     str: {"type": "string"},
@@ -127,7 +123,7 @@ async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
                 toolset=self,
                 tool_def=tool_def,
                 max_retries=1,
-                args_validator=_PASSTHROUGH_VALIDATOR,
+                args_validator=build_args_validator(json_schema),
             )
         return tools
 
@@ -152,17 +148,16 @@ async def call_tool(
 def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]:
     """Convert a Python type annotation to a JSON Schema fragment."""
     if annotation is inspect.Parameter.empty or annotation is Any:
-        return {"type": "string"}
+        return {}
+
+    if annotation is type(None):
+        return {"type": "null"}
 
     origin = get_origin(annotation)
     args = get_args(annotation)
 
-    # Optional[X] is Union[X, None] — handle both types.UnionType (3.10+) and typing.Union
     if origin is types.UnionType or origin is Union:
-        non_none = [a for a in args if a is not type(None)]
-        if len(non_none) == 1:
-            return _python_type_to_json_schema(non_none[0])
-        return {"type": "string"}
+        return {"anyOf": [_python_type_to_json_schema(arg) for arg in args]}
 
     # list[X]
     if origin is list:
@@ -175,7 +170,7 @@ def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]:
 
     # Always return a fresh copy — callers may mutate the dict (e.g. adding "description").
     schema = _TYPE_MAP.get(annotation)
-    return dict(schema) if schema else {"type": "string"}
+    return dict(schema) if schema else {}
 
 
 def _build_json_schema_from_signature(method: Callable[..., Any]) -> dict[str, Any]:
@@ -189,12 +184,15 @@ def _build_json_schema_from_signature(method: Callable[..., Any]) -> dict[str, A
 
     properties: dict[str, Any] = {}
     required: list[str] = []
+    allows_additional_properties = False
 
     for name, param in sig.parameters.items():
         if name in ("self", "cls"):
             continue
-        # Skip **kwargs and *args
-        if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD):
+        if param.kind is param.VAR_POSITIONAL:
+            continue
+        if param.kind is param.VAR_KEYWORD:
+            allows_additional_properties = True
             continue
 
         annotation = hints.get(name, param.annotation)
@@ -207,6 +205,8 @@ def _build_json_schema_from_signature(method: Callable[..., Any]) -> dict[str, A
     schema: dict[str, Any] = {"type": "object", "properties": properties}
     if required:
         schema["required"] = required
+    if allows_additional_properties:
+        schema["additionalProperties"] = True
     return schema
 
 

@@ -35,6 +35,7 @@
 import concurrent.futures
 from typing import TYPE_CHECKING, Any
 
+from pydantic import ValidationError
 from pydantic_ai import RunContext
 from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.models.test import TestModel
@@ -85,10 +86,19 @@ def airflow_toolset_to_langchain_tools(
     works regardless of how the agent handles tool errors. Raising instead would
     abort the run under ``create_agent``'s default tool-error handling.
 
+    Argument validation failures are handled the same way: each tool validates
+    its arguments with the toolset's ``args_validator`` before dispatch, and a
+    :exc:`pydantic.ValidationError` from that step is fed back to the model as
+    the tool output so it can correct the call. This mirrors pydantic-ai's
+    native two-stage behaviour: only arg-validation ``ValidationError`` is
+    retried; a ``ValidationError`` raised inside ``call_tool`` (for example from
+    a Hook method or MCP client) propagates rather than being fed back, so a
+    non-idempotent tool that already ran a side effect is not re-invoked.
+
     The retry message is bounded by the tool's ``max_retries``: a tool that keeps
-    raising ``ModelRetry`` (for example an unrecoverable connection error) stops
-    being fed back and propagates once the budget is exhausted, so the run fails
-    instead of looping forever. The count resets after a successful call.
+    raising ``ModelRetry`` (or keeps failing arg validation) stops being fed back
+    and propagates once the budget is exhausted, so the run fails instead of
+    looping forever. The count resets after a successful call.
 
     The toolset's ``get_tools`` is invoked eagerly here to enumerate the tools.
 
@@ -153,16 +163,16 @@ def _validate(kwargs: dict[str, Any]) -> dict[str, Any]:
         # the args unchanged; a typed one coerces them (e.g. "5" -> 5).
         return toolset_tool.args_validator.validate_python(kwargs)
 
-    # ModelRetry is a "feed this back to the model and retry" signal, so the bridge
-    # returns its message as the tool output instead of raising (see docstring).
-    # Bound it the way native pydantic-ai does, via the tool's max_retries: a tool
-    # that keeps raising ModelRetry (e.g. an unrecoverable connection error) must
-    # eventually propagate so the run fails rather than looping forever. The count
-    # resets on the first successful call.
+    # Mirror pydantic-ai's ToolManager two-stage handling: ValidationError from
+    # arg validation is a "feed this back and retry" signal; ModelRetry from
+    # call_tool is too; a ValidationError raised inside call_tool is not (it
+    # would otherwise re-invoke a non-idempotent tool that already ran). Bound
+    # retries via the tool's max_retries so a tool that keeps failing eventually
+    # propagates. The count resets on the first successful call.
     max_retries = toolset_tool.max_retries if toolset_tool.max_retries is not None else 1
     retries = {"count": 0}
 
-    def _handle_retry(error: ModelRetry) -> str:
+    def _handle_retry(error: ModelRetry | ValidationError) -> str:
         retries["count"] += 1
         if retries["count"] > max_retries:
             # Reset before propagating so a reused tool starts the next run with a
@@ -173,15 +183,23 @@ def _handle_retry(error: ModelRetry) -> str:
 
     def _sync_call(**kwargs: Any) -> Any:
         try:
-            result = _run_coro_sync(toolset.call_tool(name, _validate(kwargs), ctx, toolset_tool))
+            validated = _validate(kwargs)
+        except ValidationError as e:
+            return _handle_retry(e)
+        try:
+            result = _run_coro_sync(toolset.call_tool(name, validated, ctx, toolset_tool))
         except ModelRetry as e:
             return _handle_retry(e)
         retries["count"] = 0
         return result
 
     async def _async_call(**kwargs: Any) -> Any:
         try:
-            result = await toolset.call_tool(name, _validate(kwargs), ctx, toolset_tool)
+            validated = _validate(kwargs)
+        except ValidationError as e:
+            return _handle_retry(e)
+        try:
+            result = await toolset.call_tool(name, validated, ctx, toolset_tool)
         except ModelRetry as e:
             return _handle_retry(e)
         retries["count"] = 0

@@ -39,16 +39,13 @@
 from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
-from pydantic_core import SchemaValidator, core_schema
 
-from airflow.providers.common.ai.utils.tool_definition import return_schema_kwargs
+from airflow.providers.common.ai.utils.tool_definition import build_args_validator, return_schema_kwargs
 from airflow.providers.common.compat.sdk import BaseHook
 
 if TYPE_CHECKING:
     from pydantic_ai._run_context import RunContext
 
-_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema())
-
 # JSON Schemas for the four SQL tools.
 _LIST_TABLES_SCHEMA: dict[str, Any] = {
     "type": "object",
@@ -272,7 +269,7 @@ async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
                 toolset=self,
                 tool_def=tool_def,
                 max_retries=1,
-                args_validator=_PASSTHROUGH_VALIDATOR,
+                args_validator=build_args_validator(schema),
             )
         return tools
 

@@ -19,9 +19,10 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Any
+from typing import Any, Literal
 
 from pydantic_ai.tools import ToolDefinition
+from pydantic_core import SchemaValidator, core_schema
 
 # ``ToolDefinition.return_schema`` is newer than the provider's pydantic-ai
 # floor. Detect it once so callers can include the kwarg only when supported,
@@ -42,3 +43,74 @@ def return_schema_kwargs(schema: dict[str, Any]) -> dict[str, Any]:
     if _SUPPORTS_RETURN_SCHEMA:
         return {"return_schema": schema}
     return {}
+
+
+def _fragment_to_core_schema(fragment: dict[str, Any]) -> core_schema.CoreSchema:
+    any_of = fragment.get("anyOf")
+    if isinstance(any_of, list):
+        choices: list[core_schema.CoreSchema | tuple[core_schema.CoreSchema, str]] = [
+            _fragment_to_core_schema(choice) for choice in any_of if isinstance(choice, dict)
+        ]
+        return core_schema.union_schema(choices) if choices else core_schema.any_schema()
+
+    schema_type = fragment.get("type")
+    if isinstance(schema_type, list):
+        choices = [
+            _fragment_to_core_schema({**fragment, "type": item})
+            for item in schema_type
+            if isinstance(item, str)
+        ]
+        return core_schema.union_schema(choices) if choices else core_schema.any_schema()
+
+    match schema_type:
+        case "string":
+            return core_schema.str_schema()
+        case "integer":
+            return core_schema.int_schema()
+        case "number":
+            return core_schema.float_schema()
+        case "boolean":
+            return core_schema.bool_schema()
+        case "null":
+            return core_schema.none_schema()
+        case "array":
+            items = fragment.get("items")
+            return core_schema.list_schema(
+                _fragment_to_core_schema(items) if isinstance(items, dict) else None
+            )
+        case "object":
+            return _object_fragment_to_core_schema(fragment)
+        case _:
+            return core_schema.any_schema()
+
+
+def _object_fragment_to_core_schema(fragment: dict[str, Any]) -> core_schema.CoreSchema:
+    """
+    Convert a JSON Schema ``object`` fragment to a core schema.
+
+    A fragment with no ``properties`` key is an untyped object (e.g. from a
+    ``dict[K, V]`` annotation): accept any dict rather than stripping its
+    contents. When ``properties`` is present, build a typed-dict that validates
+    each declared field recursively — nested objects are handled the same way
+    arrays already recurse into ``items``.
+
+    Undeclared keys follow native pydantic-ai: ``forbid`` for fixed signatures
+    (so a mistyped field name becomes a bounded retry), ``allow`` only when the
+    schema sets ``additionalProperties: true`` (methods that accept ``**kwargs``).
+    """
+    if "properties" not in fragment:
+        return core_schema.dict_schema()
+    required = set(fragment.get("required", []))
+    fields = {
+        name: core_schema.typed_dict_field(_fragment_to_core_schema(prop), required=name in required)
+        for name, prop in fragment["properties"].items()
+    }
+    extra_behavior: Literal["allow", "forbid"] = (
+        "allow" if fragment.get("additionalProperties") is True else "forbid"
+    )
+    return core_schema.typed_dict_schema(fields, extra_behavior=extra_behavior)
+
+
+def build_args_validator(parameters_json_schema: dict[str, Any]) -> SchemaValidator:
+    """Build an argument validator from the schema advertised to the model."""
+    return SchemaValidator(_object_fragment_to_core_schema(parameters_json_schema))
@@ -25,6 +25,7 @@
 from pydantic_ai._run_context import RunContext
 from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.toolsets.abstract import ToolsetTool
+from pydantic_core import ValidationError
 
 from airflow.providers.common.ai.toolsets.datafusion import (
     _RETRYABLE_QUERY_ERROR_PATTERNS,
@@ -107,6 +108,24 @@ def test_tool_definitions_have_descriptions(self):
             assert tool.tool_def.description
 
 
+class TestDataFusionToolsetArgsValidation:
+    @pytest.mark.parametrize(
+        ("tool_name", "valid_args"),
+        [
+            ("get_schema", {"table_name": "sales_data"}),
+            ("query", {"sql": "SELECT 1"}),
+        ],
+    )
+    def test_enforces_required_args(self, tool_name, valid_args):
+        cfg = _make_mock_datasource_config()
+        ts = DataFusionToolset([cfg])
+        tools = asyncio.run(ts.get_tools(ctx=MagicMock(spec=RunContext)))
+        validator = tools[tool_name].args_validator
+        assert validator.validate_python(valid_args) == valid_args
+        with pytest.raises(ValidationError):
+            validator.validate_python({})
+
+
 class TestDataFusionToolsetListTables:
     def test_returns_registered_tables(self):
         cfg = _make_mock_datasource_config()