diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py index 63037e1a4f8f9..8cba44d5153a1 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py @@ -24,6 +24,7 @@ import types from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints +from pydantic_ai.exceptions import ModelRetry from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool @@ -61,6 +62,10 @@ class HookToolset(AbstractToolset[Any]): auto-discovery is intentionally not supported for safety. :param tool_name_prefix: Optional prefix prepended to each tool name (e.g. ``"s3_"`` → ``"s3_list_keys"``). + + When a tool call fails, the hook's own error message is returned to the agent as a + retry (:class:`pydantic_ai.ModelRetry`) so the model can correct its arguments within + the run instead of failing the task. """ def __init__( @@ -136,7 +141,10 @@ async def call_tool( ) -> Any: method_name = name.removeprefix(self._tool_name_prefix) if self._tool_name_prefix else name method: Callable[..., Any] = getattr(self._hook, method_name) - result = method(**tool_args) + try: + result = method(**tool_args) + except Exception as e: + raise ModelRetry(f"The {name} tool failed: {e}") from e return _serialize_for_llm(result) diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py index ae2d4c6f86754..eaa8645e30cc6 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py @@ -20,6 +20,7 @@ from unittest.mock import MagicMock import pytest +from pydantic_ai.exceptions import ModelRetry from pydantic_core import ValidationError from airflow.providers.common.ai.toolsets.hook import ( @@ -50,6 +51,10 @@ def read_file(self, key: str) -> str: def no_docstring(self, x: int) -> int: return x * 2 + def failing_method(self, x: int) -> int: + """Always raise, for testing call_tool failure handling.""" + raise RuntimeError("boom") + def request( self, endpoint: str | None = None, data: dict[str, object] | str | None = None, **kwargs: object ) -> dict[str, object]: @@ -203,6 +208,21 @@ def test_dispatches_with_prefix(self): ) assert result == "contents of test.txt" + def test_wraps_failure_in_model_retry(self): + hook = _FakeHook() + ts = HookToolset(hook, allowed_methods=["failing_method"]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + + with pytest.raises(ModelRetry, match="failing_method tool failed: boom"): + asyncio.run( + ts.call_tool( + "failing_method", + {"x": 1}, + ctx=MagicMock(), + tool=tools["failing_method"], + ) + ) + class TestBuildJsonSchemaFromSignature: def test_basic_types(self):