Update chanelogs and remove dict output type

Author: gopidesupavan

providers/common/ai/docs/changelog.rst

@@ -53,6 +53,14 @@ and :meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.run_age
 DAG authors using :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`,
 ``@task.agent``, and the other LLM operators are unaffected.
 
+``SQLToolset`` now implements the framework-neutral
+:class:`~airflow.providers.common.ai.hooks.base.BaseToolset` interface instead of
+pydantic-ai's ``AbstractToolset`` interface. DAG authors using ``SQLToolset``
+with ``AgentOperator`` or ``@task.agent`` are unaffected. Direct pydantic-ai
+``Agent(toolsets=[SQLToolset(...)])`` callers should use
+``AgentOperator(toolsets=[SQLToolset(...)])`` or pass the SQL tool callables
+through an Airflow agent hook request.
+
 0.4.0
 .....
 

providers/common/ai/docs/operators/agent.rst

@@ -312,8 +312,11 @@ Parameters
   (``SQLToolset``), pydantic-ai ``AbstractToolset`` implementations
   (``HookToolset``, ``MCPToolset``, ``DataFusionToolset``,
   ``AgentSkillsToolset`` for :ref:`agent-skills`, third-party toolsets),
-  plain Python callables, or native pydantic-ai ``Tool`` objects. Mixed lists
-  are supported.
+  pydantic-ai ``DynamicToolset`` instances, plain Python callables, or native
+  pydantic-ai ``Tool`` objects. Mixed lists are supported. Bare Python
+  callables are treated as callable tools; wrap pydantic-ai ``ToolsetFunc``
+  factories with ``DynamicToolset`` to pass them through as native dynamic
+  toolsets.
 - ``enable_tool_logging``: When ``True`` (default), wraps each tool call with
   real-time logging. For pydantic-ai ``AbstractToolset`` items this is done via
   :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`; for

providers/common/ai/docs/toolsets.rst

@@ -53,12 +53,35 @@ toolset to the correct agent parameter automatically.
     - Any third-party ``AbstractToolset``, including PydanticAI's own MCP
       server classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``,
       ``MCPServerStdio``).
+    - pydantic-ai dynamic toolsets, by wrapping a ``ToolsetFunc`` factory with
+      ``DynamicToolset``.
     - :class:`~airflow.providers.common.ai.hooks.base.BaseToolset`
       subclasses (``SQLToolset``).
     - Plain Python callables (``def my_tool(...): ...``).
     - Native pydantic-ai ``Tool`` objects.
 
     The hook routes each item to the correct agent parameter automatically.
+    Bare Python callables are treated as callable tools. To pass a pydantic-ai
+    ``ToolsetFunc`` factory through as a native dynamic toolset, wrap it with
+    ``DynamicToolset``:
+
+    .. code-block:: python
+
+        from pydantic_ai import RunContext
+        from pydantic_ai.agent import DynamicToolset
+        from pydantic_ai.toolsets import AbstractToolset
+
+
+        def select_toolset(ctx: RunContext) -> AbstractToolset | None:
+            return None
+
+
+        AgentOperator(
+            task_id="agent",
+            prompt="Answer with the tools available for this run.",
+            llm_conn_id="pydanticai_default",
+            toolsets=[DynamicToolset(select_toolset)],
+        )
 
 
 Using Toolsets Directly

providers/common/ai/src/airflow/providers/common/ai/hooks/base.py

@@ -137,8 +137,7 @@ class AgentRunRequest:
 
     :param prompt: User prompt for this invocation (plain ``str`` or a multimodal
         ``Sequence`` accepted by the backend agent's run API).
-    :param output_type: Expected structured output type or backend-specific JSON schema
-        mapping (default: ``str``).
+    :param output_type: Expected structured output type (default: ``str``).
     :param instructions: System-level instructions for the agent.
     :param toolsets: List of tools/toolsets the agent may call (BaseToolset instances, plain callables, or backend-native tool objects).
     :param usage_limits: Backend-specific usage limits; ignored if the hook does not support them.
@@ -152,7 +151,7 @@ class AgentRunRequest:
     """
 
     prompt: str | Sequence[Any]
-    output_type: type[Any] | dict[str, Any] | None = str
+    output_type: type[Any] | None = str
     instructions: str = ""
     toolsets: list[Any] | None = None
     usage_limits: Any = None

providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py

@@ -298,6 +298,12 @@ def _build_agent(self, request: AgentRunRequest) -> PydanticAgentHandle:
                     [type(toolset).__name__ for toolset in processed],
                 )
 
+        if isinstance(request.output_type, dict):
+            raise ValueError(
+                "PydanticAIHook does not support raw JSON schema mappings for output_type. "
+                "Pass a Python type, such as a Pydantic BaseModel subclass."
+            )
+
         agent_kwargs: dict[str, Any] = {"instructions": request.instructions, **extra_kwargs}
         if request.output_type is not None:
             agent_kwargs["output_type"] = request.output_type

providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py

@@ -41,6 +41,7 @@
 from pydantic_ai.usage import RunUsage
 
 from airflow.providers.common.ai.hooks.base import BaseToolset
+from airflow.providers.common.ai.utils.callables import is_async_callable
 
 if TYPE_CHECKING:
     from collections.abc import Coroutine
@@ -187,15 +188,15 @@ def _build_structured_tool_from_spec(
 
     def _sync_call(**kwargs: Any) -> Any:
         try:
-            if asyncio.iscoroutinefunction(spec.fn):
+            if is_async_callable(spec.fn):
                 return _run_coro_sync(spec.fn(**kwargs))
             return spec.fn(**kwargs)
         except ModelRetry as e:
             return str(e)
 
     async def _async_call(**kwargs: Any) -> Any:
         try:
-            if asyncio.iscoroutinefunction(spec.fn):
+            if is_async_callable(spec.fn):
                 return await spec.fn(**kwargs)
             return spec.fn(**kwargs)
         except ModelRetry as e:

providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py

@@ -19,10 +19,12 @@
 import functools
 import json
 import sys
+from typing import TYPE_CHECKING
 from unittest.mock import MagicMock, patch
 
 import pytest
-from pydantic_ai import Agent
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.agent import DynamicToolset
 from pydantic_ai.messages import ModelResponse, TextPart
 from pydantic_ai.models import Model
 from pydantic_ai.models.test import TestModel
@@ -46,6 +48,9 @@
 )
 from airflow.providers.common.ai.mixins.durable import DurableState
 
+if TYPE_CHECKING:
+    from pydantic_ai.toolsets import AbstractToolset
+
 
 def _test_agent() -> Agent[None, str]:
     return Agent(TestModel())
@@ -453,6 +458,25 @@ def test_create_agent_with_agent_params(self, mock_agent_cls, mock_infer_model):
             retries=3,
         )
 
+    @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True)
+    @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True)
+    def test_create_agent_rejects_raw_json_schema_output_type(self, mock_agent_cls, mock_infer_model):
+        mock_model = MagicMock(spec=Model)
+        mock_infer_model.return_value = mock_model
+
+        hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3")
+        conn = Connection(conn_id="test_conn", conn_type="pydanticai")
+        request = AgentRunRequest(
+            prompt="hi",
+            output_type={"type": "object", "properties": {}},  # type: ignore[arg-type]
+        )
+
+        with patch.object(hook, "get_connection", return_value=conn):
+            with pytest.raises(ValueError, match="raw JSON schema mappings"):
+                hook.create_agent(request)
+
+        mock_agent_cls.assert_not_called()
+
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True)
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True)
     def test_create_agent_rejects_tools_in_agent_params_with_toolsets(self, mock_agent_cls, mock_infer_model):
@@ -650,6 +674,30 @@ def test_create_agent_routes_abstract_toolset_to_toolsets_kwarg(self, mock_agent
         assert "toolsets" in call_kwargs
         assert any(ts is abstract_ts for ts in call_kwargs["toolsets"])
 
+    @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True)
+    @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True)
+    def test_create_agent_routes_dynamic_toolset_to_toolsets_kwarg(self, mock_agent_cls, mock_infer_model):
+        """DynamicToolset-wrapped factories must pass through as native pydantic-ai toolsets."""
+
+        mock_model = MagicMock(spec=Model)
+        mock_infer_model.return_value = mock_model
+
+        def select_toolset(ctx: RunContext) -> AbstractToolset | None:
+            return None
+
+        dynamic_toolset = DynamicToolset(select_toolset)
+
+        hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3")
+        conn = Connection(conn_id="test_conn", conn_type="pydanticai")
+        request = AgentRunRequest(prompt="hi", toolsets=[dynamic_toolset], enable_tool_logging=False)
+        with patch.object(hook, "get_connection", return_value=conn):
+            hook.create_agent(request)
+
+        call_kwargs = mock_agent_cls.call_args[1]
+        assert "tools" not in call_kwargs
+        assert "toolsets" in call_kwargs
+        assert any(ts is dynamic_toolset for ts in call_kwargs["toolsets"])
+
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True)
     @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True)
     def test_create_agent_wraps_abstract_toolset_with_logging(self, mock_agent_cls, mock_infer_model):