This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2448be40082 Make common.ai tool_prefix and output_type signatures
consistent (#74695)
2448be40082 is described below
commit 2448be40082ffbb14bcc3a10db42e9cd743a087e
Author: Wei Lee <[email protected]>
AuthorDate: Thu Nov 12 14:59:34 2026 +0900
Make common.ai tool_prefix and output_type signatures consistent (#74695)
Public signatures freeze at 1.0, and changing a parameter's type, default
or availability afterwards is a breaking change. tool_prefix was str on some
toolsets, str | None on MCPToolset and MCPHook, and missing on SQLToolset and
DataFusionToolset. AgentOperator also annotated output_type narrower than
LLMOperator even though both accept the same values at runtime. HookToolset
keeps its deprecated tool_name_prefix alias until 2.0.
SQLToolset and DataFusionToolset also override tool_name_conflict_hint so a
name conflict points users to tool_prefix. rehydrate_pydantic_output now
accepts a list output_type by validating it as a union, so a list output_type
no longer crashes after a HITL review.
Models returned under a list output_type are dumped through a shared
dump_output_for_xcom helper. With serialize_output=True, and on task-sdk
versions without the walker, the dump keeps model_dump() as in 1.0.0. A model
whose class the worker cannot register is dumped with to_jsonable_python, so
HttpUrl, Path and Enum fields serialize.
---
providers/common/ai/docs/changelog.rst | 36 +++
providers/common/ai/docs/concepts.rst | 2 +-
providers/common/ai/docs/operators/agent.rst | 13 +-
providers/common/ai/docs/operators/llm.rst | 1 +
.../common/ai/docs/operators/llm_file_analysis.rst | 13 +-
providers/common/ai/docs/stability.rst | 11 +-
providers/common/ai/docs/structured_output.rst | 48 ++-
providers/common/ai/docs/toolsets/datafusion.rst | 3 +
providers/common/ai/docs/toolsets/mcp.rst | 4 +-
providers/common/ai/docs/toolsets/sql.rst | 10 +-
.../airflow/providers/common/ai/decorators/llm.py | 3 +-
.../src/airflow/providers/common/ai/hooks/mcp.py | 4 +-
.../airflow/providers/common/ai/operators/agent.py | 70 ++---
.../airflow/providers/common/ai/operators/llm.py | 37 ++-
.../common/ai/operators/llm_file_analysis.py | 10 +-
.../providers/common/ai/toolsets/datafusion.py | 33 +-
.../airflow/providers/common/ai/toolsets/hook.py | 6 +-
.../airflow/providers/common/ai/toolsets/mcp.py | 6 +-
.../providers/common/ai/toolsets/object_storage.py | 17 +-
.../providers/common/ai/toolsets/sandbox.py | 24 +-
.../airflow/providers/common/ai/toolsets/sql.py | 35 ++-
.../providers/common/ai/utils/output_type.py | 81 ++++-
.../providers/common/ai/utils/query_results.py | 20 +-
.../providers/common/ai/utils/tool_prefix.py | 56 ++++
.../ai/tests/unit/common/ai/hooks/test_mcp.py | 18 ++
.../tests/unit/common/ai/operators/test_agent.py | 82 ++++-
.../ai/tests/unit/common/ai/operators/test_llm.py | 78 ++++-
.../common/ai/operators/test_llm_file_analysis.py | 28 ++
.../unit/common/ai/toolsets/test_datafusion.py | 90 ++++++
.../ai/tests/unit/common/ai/toolsets/test_mcp.py | 9 +-
.../unit/common/ai/toolsets/test_object_storage.py | 13 +-
.../tests/unit/common/ai/toolsets/test_sandbox.py | 3 +
.../ai/tests/unit/common/ai/toolsets/test_sql.py | 114 +++++++
.../tests/unit/common/ai/utils/test_output_type.py | 343 ++++++++++++++++++++-
.../unit/common/ai/utils/test_query_results.py | 27 ++
.../tests/unit/common/ai/utils/test_tool_prefix.py | 78 +++++
36 files changed, 1240 insertions(+), 186 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index 17fb111bd10..fab49897489 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -44,6 +44,42 @@ Changelog
with ``result.to_state()`` and ``state.approve()`` after
``run_durably_sync()`` should call
``Runner.run`` and handle ``result.interruptions`` itself.
+.. note::
+ **Significant change: output serialization of the LLM operators**
+
+ ``LLMOperator``, ``@task.llm`` and ``LLMFileAnalysisOperator`` now push a
model as a
+ JSON-compatible dict, where 1.1.0 pushed it as a model instance, in these
cases:
+
+ * On Airflow 3.3 and later, the model is itself the output,
``serialize_output`` is not set,
+ and ``output_type`` is one the worker cannot register a class for (a list
of output types
+ such as ``[Summary, int]``, an output marker, an output function, a
``TypeAliasType`` or a
+ ``NewType``).
+ * On Airflow 3.3 and later, the model is inside a list, tuple, set or dict
in the output and
+ the worker cannot register its class, whether or not ``serialize_output``
is set.
+ * On Airflow < 3.3, the model is inside a list, tuple, set or dict in the
output, whatever the
+ ``output_type`` (for example ``list[MyModel]``).
+
+ When a model inside a list, tuple, set or dict is dumped, the whole output
is dumped with
+ it, so a tuple in the output arrives as a list and a ``date`` next to the
model arrives as a
+ string. A downstream task indexes the dict (``result["field"]``) instead of
using attribute
+ access. As in 1.1.0, a model that is itself the output is dumped with
``model_dump()`` when
+ ``serialize_output=True`` is set or on Airflow < 3.3.
+
+ ``AgentOperator`` and ``@task.agent`` already pushed these models as dicts
in 1.1.0.
+ Compared with 1.1.0, when they dump to JSON values, a model that sets
+ ``serialize_by_alias=True`` now arrives with its field names (``full_name``,
not
+ ``fullName``), a tuple in the output arrives as a list, and a ``date`` next
to the model
+ arrives as a string. A model inside a set is now dumped too, and the set
arrives as a list.
+ A model that is itself the output and is dumped with ``model_dump()`` keeps
its aliases, as
+ in 1.1.0. See :ref:`When a model arrives as a dict
<structured-output-dict-shape>`.
+
+.. note::
+ **Signature change: MCP ``tool_prefix``**
+
+ The ``tool_prefix`` parameter of ``MCPHook`` and ``MCPToolset`` now defaults
to ``""`` and is
+ typed ``str``; in 1.1.0 it defaulted to ``None`` and was typed ``str |
None``. This changes
+ only the signature and the annotated type, not the behaviour.
+
1.1.0
.....
diff --git a/providers/common/ai/docs/concepts.rst
b/providers/common/ai/docs/concepts.rst
index 27c64b05426..c63f73e872a 100644
--- a/providers/common/ai/docs/concepts.rst
+++ b/providers/common/ai/docs/concepts.rst
@@ -110,7 +110,7 @@ Results flow through XCom
Every operator pushes its result to XCom like any other task. A plain string
arrives as a
string. A Pydantic ``output_type`` arrives as the model instance, typed, so a
downstream
-task can use attribute access. :doc:`structured_output` explains how the class
is
+task can use attribute access (a model inside a list of types arrives as a
``dict``). :doc:`structured_output` explains how the class is
registered for deserialization and where that stops working.
Related providers
diff --git a/providers/common/ai/docs/operators/agent.rst
b/providers/common/ai/docs/operators/agent.rst
index c8b1ea31dd4..f0506d1a494 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -112,7 +112,8 @@ Structured output
Set ``output_type`` to a Pydantic ``BaseModel`` subclass to get structured data
back. The model instance is pushed to XCom unchanged so downstream tasks can
type-hint the class directly (``def downstream(result: MyModel)``) and use
-attribute access (``result.field``).
+attribute access (``result.field``). Inside a list of types (``[MyModel,
int]``) the
+model is pushed as a ``dict`` instead.
:doc:`../structured_output` explains the XCom deserialization rules, the
cross-Dag gap and
``serialize_output``.
@@ -454,11 +455,11 @@ Parameters
- ``emit_message_history``: Push the post-run transcript to XCom under the key
``message_history``. Defaults to ``True`` when ``message_history`` is set
and ``False``
otherwise; ``True`` records a single-turn run without seeding it. See
:doc:`../message_history`.
-- ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic
- ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
- ``model_dump()`` before being pushed to XCom. Default ``False`` -- the
- Pydantic instance flows through XCom unchanged. Set to ``True`` when a
- downstream consumer needs the dict shape.
+- ``serialize_output``: If ``True``, an output that is itself a Pydantic
``BaseModel`` is dumped with
+ ``model_dump()`` before being pushed to XCom, so ``date`` and ``Decimal``
fields keep their types.
+ Default ``False``: a model whose class the worker registers from
``output_type`` flows through XCom
+ unchanged. Some outputs are pushed as dicts either way; see
+ :ref:`When a model arrives as a dict <structured-output-dict-shape>`.
**HITL review parameters**: ``enable_hitl_review``, ``max_hitl_iterations``,
``hitl_timeout`` and ``hitl_poll_interval`` turn on and bound the iterative
review
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index e6e1040518b..83b134c439b 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -46,6 +46,7 @@ Structured output
Set ``output_type`` to a Pydantic ``BaseModel`` subclass and the model
instance is pushed to
XCom unchanged, so downstream tasks can type-hint the class and use attribute
access.
+Inside a list of types (``[MyModel, int]``) the model is pushed as a ``dict``
instead.
:doc:`../structured_output` explains the XCom deserialization rules, the
cross-Dag gap and
``serialize_output``.
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index e4efec17dae..1e58a977123 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -77,7 +77,8 @@ Structured Output
Set ``output_type`` to a Pydantic ``BaseModel`` when you want a typed response
back from the LLM instead of a plain string. The model instance is pushed to
-XCom unchanged so downstream tasks can type-hint the class directly. The
+XCom unchanged so downstream tasks can type-hint the class directly. Inside a
+list of types (``[MyModel, int]``) the model is pushed as a ``dict`` instead.
The
declared ``output_type`` (and any ``BaseModel`` reachable from
``Union``/``Optional``/``list`` shapes) is registered for deserialization by
the
worker when it loads the Dag. Define the class at **module scope** and bind it
to
@@ -157,11 +158,11 @@ Parameters
- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` (or a templated
``dict`` of
the same fields) enforced on the run; the task fails when a budget is
exceeded.
Default ``None``. See :ref:`Usage Limits <howto/operator:llm_usage_limits>`.
-- ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic
- ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
- ``model_dump()`` before being pushed to XCom. Default ``False`` -- the
- Pydantic instance flows through XCom unchanged. Set to ``True`` when a
- downstream consumer needs the dict shape.
+- ``serialize_output``: If ``True``, an output that is itself a Pydantic
``BaseModel`` is dumped with
+ ``model_dump()`` before being pushed to XCom, so ``date`` and ``Decimal``
fields keep their types.
+ Default ``False``: a model whose class the worker registers from
``output_type`` flows through XCom
+ unchanged. Some outputs are pushed as dicts either way; see
+ :ref:`When a model arrives as a dict <structured-output-dict-shape>`.
``decision_policy`` is not supported here: the operator runs its own
``execute``
without the confidence gate and rejects a policy with a bar at construction.
diff --git a/providers/common/ai/docs/stability.rst
b/providers/common/ai/docs/stability.rst
index 733ccedc052..22a4cdb7807 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -46,11 +46,13 @@ Pydantic AI toolsets, but the Pydantic AI class they
inherit from can change.
* - ``@task.llm`` and
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
- Sends the prompt to the model from ``llm_conn_id`` and returns the
output as the
task's return value: a string by default, or an instance of
``output_type``,
- dumped to a dict when ``serialize_output=True``. Before Airflow 3.3 an
instance is
- always dumped to a dict. ``decision_policy`` is experimental; see below.
+ dumped to a dict when ``serialize_output=True``. Some models arrive as
dicts without it;
+ see :ref:`When a model arrives as a dict
<structured-output-dict-shape>`.
+ ``decision_policy`` is experimental; see below.
* - ``@task.agent`` and
:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`
- Runs an agent with the model from ``llm_conn_id`` and the given
toolsets, and
- returns its output under the same rules as ``@task.llm``. Code mode is
+ returns its output as the task's return value. Some models arrive as
dicts; see
+ :ref:`When a model arrives as a dict <structured-output-dict-shape>`.
Code mode is
experimental; see below.
* - Per-tool approval on ``AgentOperator`` (:doc:`tool_approval`)
- On Airflow 3.3 and later, a call to a tool marked with Pydantic AI's
approval API
@@ -142,7 +144,8 @@ Pydantic AI toolsets, but the Pydantic AI class they
inherit from can change.
:doc:`connections/pydantic_ai_azure` keep their meaning, so an existing
connection
keeps producing the same model.
* - :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset`
- - Exposes ``list_tables``, ``get_schema``, ``query`` and ``check_query``.
Read-only
+ - Exposes ``list_tables``, ``get_schema``, ``query`` and ``check_query``,
each named
+ ``<tool_prefix>_<name>`` when ``tool_prefix`` is set. Read-only
unless ``allow_writes=True``. When ``allowed_tables`` is set, other
tables are
refused. ``query`` returns at most ``max_rows`` rows, and keeps the
rows and column
names within ``max_result_bytes`` bytes. When ``query`` fails or
refuses a
diff --git a/providers/common/ai/docs/structured_output.rst
b/providers/common/ai/docs/structured_output.rst
index e8ca62d062a..1bd2c9cd087 100644
--- a/providers/common/ai/docs/structured_output.rst
+++ b/providers/common/ai/docs/structured_output.rst
@@ -24,6 +24,7 @@ Set ``output_type`` to a Pydantic ``BaseModel`` subclass. The
LLM is instructed
to return structured data, and the model instance is pushed to XCom unchanged
so downstream tasks can type-hint the class directly
(``def downstream(result: MyModel)``) and use attribute access
(``result.field``).
+Inside a list of types (``[MyModel, int]``) the model is pushed as a ``dict``
instead.
The declared ``output_type`` (and any ``BaseModel`` reachable from
``Union``/``Optional``/``list[...]`` types) is registered for XCom
deserialization by
@@ -34,6 +35,8 @@ classes nested inside a function or ``@dag``-decorated body,
parameterized
generics, and dynamically-built classes whose ``__name__`` does not match the
attribute they are bound to cannot be re-imported, so they are skipped with a
warning at worker startup and the value fails to deserialize at the consumer.
+A list of types (``output_type=[MyModel, int]``) registers no class, so the
operator pushes the
+model as a ``dict`` even when ``serialize_output`` is ``False``.
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
:language: python
@@ -53,7 +56,8 @@ cache that bypasses operator construction.
The worker finds no class to register when ``output_type`` is a list of output
types such as ``[Summary, int]``, an output marker such as
``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
-``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+``NewType``. For these output types, ``LLMOperator``, ``@task.llm``,
+``LLMFileAnalysisOperator``, ``AgentOperator`` and ``@task.agent`` push
each model in the output as a JSON-compatible dict, including models inside a
list, tuple or dict. A model is pushed as a dict even if its class is listed in
``[core] allowed_deserialization_classes``, or if another task in the same Dag
@@ -69,10 +73,7 @@ A model at the top level is still dumped with plain
``model_dump()`` when
condition is enough. A plain ``Enum`` member (one that does not also subclass
``str`` or ``int``), an ``HttpUrl`` or a ``Path`` in that model then stays a
Python object. The default XCom backend cannot serialize these objects, so the
-XCom push fails with ``TypeError``. Before Airflow 3.3 the worker registers no
-model class. ``AgentOperator`` and ``@task.agent`` therefore push a model
-inside a list, tuple or dict as a JSON-compatible dict for every
-``output_type``, not only for the output types listed above.
+XCom push fails with ``TypeError``.
Upgrading from common.ai 1.0.0 to 1.1.0 can change the type that a downstream
task receives:
@@ -108,6 +109,37 @@ class is not auto-registered. Add the class qualified name
to
``[core] allowed_deserialization_classes`` (or a glob that matches it) to make
that pattern work.
-If a downstream consumer needs the dict shape (e.g. forwarding to an external
-system that expects JSON-style payloads), pass ``serialize_output=True`` and
the
-operator calls ``model_dump()`` before pushing to XCom.
+.. _structured-output-dict-shape:
+
+When a model arrives as a dict
+------------------------------
+
+``LLMOperator``, ``@task.llm``, ``LLMFileAnalysisOperator``, ``AgentOperator``
and ``@task.agent``
+decide in the same way whether a model is pushed as an instance or dumped. Set
+``serialize_output=True`` when a downstream consumer needs the dict shape (for
example to forward
+it to a system that expects JSON-style payloads).
+
+* With ``serialize_output=True``, an output that is itself a model is dumped
with ``model_dump()``, so
+ field types such as ``date`` and ``Decimal`` are kept.
+* A model inside a list, tuple, set or dict is not dumped when the worker
registers its class
+ (``output_type=list[MyModel]`` on Airflow 3.3 and later). The XCom
serializer pushes it as an
+ instance, with or without ``serialize_output``, unless the output also holds
a model whose class
+ the worker does not register (see the next bullets).
+* On Airflow 3.3 and later, a model whose class the worker cannot register is
dumped to JSON values
+ when it is itself the output and ``serialize_output`` is ``False``. With
``serialize_output=True`` it is dumped with
+ ``model_dump()`` as the first bullet describes. The worker registers the
classes it finds by
+ reading ``output_type`` and finds none in a list of types such as
``output_type=[MyModel, int]``.
+* The operator dumps the whole output with ``to_jsonable_python``, using field
names rather than
+ aliases (also for a model that sets ``serialize_by_alias=True``), as soon as
it holds such a model
+ in a list, tuple, set or dict, or as a set member or a dict key. A ``date``
or ``Decimal`` next to
+ the model becomes a string, a set becomes a list, and a model used as a dict
key becomes its
+ string form.
+* Before Airflow 3.3 the worker registers no class. A model that is itself the
output is then dumped
+ with ``model_dump()`` whether or not ``serialize_output`` is set, and an
output that holds a model
+ in a list, tuple, set or dict is dumped as the previous bullet describes,
for every
+ ``output_type``.
+
+In a dump to JSON values, ``HttpUrl``, ``Path`` and ``Enum`` fields become
plain JSON values. A model
+that is itself the output with ``serialize_output=True`` keeps its Python
field values instead, so an
+``HttpUrl``, a ``Path`` or a plain ``Enum`` member in it stays an object that
the default XCom backend
+cannot serialize.
diff --git a/providers/common/ai/docs/toolsets/datafusion.rst
b/providers/common/ai/docs/toolsets/datafusion.rst
index fd70c00532c..854913d8a94 100644
--- a/providers/common/ai/docs/toolsets/datafusion.rst
+++ b/providers/common/ai/docs/toolsets/datafusion.rst
@@ -144,6 +144,9 @@ Parameters
See :ref:`bounded-query-results` and :ref:`bounded-schema-results`.
- ``max_columns``: Maximum columns ``get_schema`` returns in full. Default
``100``.
Above it the result becomes a bounded summary. See
:ref:`bounded-schema-results`.
+- ``tool_prefix``: Optional prefix joined to each tool name with an underscore
(e.g. ``"lake"`` produces ``"lake_query"``).
+ Must be a valid Python identifier. Set it when one agent has another toolset
with the
+ same tool names, since duplicate tool names are rejected.
- ``max_retries``: How many times the model may correct a failed call to these
tools. Default ``None``, the agent's ``retries``. See
:ref:`toolset-retry-budget`.
diff --git a/providers/common/ai/docs/toolsets/mcp.rst
b/providers/common/ai/docs/toolsets/mcp.rst
index 65ca204f769..30cb615acf0 100644
--- a/providers/common/ai/docs/toolsets/mcp.rst
+++ b/providers/common/ai/docs/toolsets/mcp.rst
@@ -55,9 +55,9 @@ Parameters
connection runs its ``Extra.command`` on the worker, so whoever picks the
connection picks the command. A ``token_provider`` or ``env_provider`` is
shared by every connection the template renders to.
-- ``tool_prefix``: Optional prefix prepended to tool names to avoid
+- ``tool_prefix``: Optional prefix joined to each tool name with an underscore
to avoid
collisions when using multiple MCP servers (e.g. ``"weather"`` produces
- ``"weather_get_forecast"``).
+ ``"weather_get_forecast"``). Default ``""`` (no prefix).
- ``token_provider``: Optional zero-argument callable returning a bearer token.
When set, it overrides the connection's static ``password`` for the
``Authorization`` header. Called once, the first time this toolset
diff --git a/providers/common/ai/docs/toolsets/sql.rst
b/providers/common/ai/docs/toolsets/sql.rst
index 1e44b4b87ad..399262d5f37 100644
--- a/providers/common/ai/docs/toolsets/sql.rst
+++ b/providers/common/ai/docs/toolsets/sql.rst
@@ -317,7 +317,11 @@ Parameters
Default ``False`` -- only SELECT-family and read-only metadata
(``DESCRIBE``/``SHOW``) statements are permitted. To have a person approve a
``query`` call before it runs, wrap the toolset with
``.approval_required()`` and
- return ``tool_def.name == "query"`` from its function. Reads and writes both
go
+ return ``tool_def.name == f"{prefix}_query"`` from its function, where
``prefix`` is
+ your ``tool_prefix`` (with none set, compare against ``"query"``). The
function sees
+ the prefixed name, so a bare ``"query"`` comparison never matches once a
prefix is
+ set and writes run without approval. Compare the full name with ``==``; do
not use
+ ``endswith("query")``, which also matches ``check_query``. Reads and writes
both go
through ``query``, so every query pauses for a decision. On Airflow 3.3 a
task
instance can pause only once per Dag run, so there this fits an agent that
runs one
query, such as a single write; see :doc:`../tool_approval`.
@@ -329,6 +333,10 @@ Parameters
See :ref:`bounded-query-results` and :ref:`bounded-schema-results`.
- ``max_columns``: Maximum columns ``get_schema`` returns in full. Default
``100``.
Above it the result becomes a bounded summary. See
:ref:`bounded-schema-results`.
+- ``tool_prefix``: Optional prefix joined to each tool name with an underscore
(e.g. ``"pg"`` produces ``"pg_query"``).
+ Must be a valid Python identifier. Set it when one agent has another toolset
with the
+ same tool names, since duplicate tool names are rejected. A function passed
to
+ ``.approval_required()`` sees these prefixed names; see ``allow_writes``
above.
- ``max_retries``: How many times the model may correct a failed call to these
tools. Default ``None``, the agent's ``retries``. See
:ref:`toolset-retry-budget`.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
index 32514c62513..60824be3ed6 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
@@ -20,7 +20,8 @@ TaskFlow decorator for general-purpose LLM calls.
The user writes a function that **returns the prompt string**. The decorator
handles hook creation, agent configuration, and the LLM call. When
``output_type`` is a Pydantic ``BaseModel`` subclass, the model instance is
-returned to XCom unchanged so downstream tasks can type-hint it directly.
+returned to XCom unchanged so downstream tasks can type-hint it directly,
except
+inside a list of types (``[MyModel, int]``), where it is returned as a
``dict``.
The class must be defined at module scope.
"""
diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
index 810b0fccbc4..dc39a8ddc5e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
@@ -66,7 +66,7 @@ class MCPHook(BaseHook):
:param mcp_conn_id: Airflow connection ID for the MCP server.
:param tool_prefix: Optional prefix prepended to tool names
- (e.g. ``"weather"`` → ``"weather_get_forecast"``).
+ (e.g. ``"weather"`` → ``"weather_get_forecast"``). Defaults to ``""``
(no prefix).
:param token_provider: Optional zero-argument callable returning a bearer
token string. When set, it overrides the connection ``password`` for
the
``Authorization`` header on HTTP/SSE transports. Called once, the first
@@ -88,7 +88,7 @@ class MCPHook(BaseHook):
def __init__(
self,
mcp_conn_id: str = default_conn_name,
- tool_prefix: str | None = None,
+ tool_prefix: str = "",
*,
token_provider: Callable[[], str] | None = None,
env_provider: Callable[[], dict[str, str]] | None = None,
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 05da22c2760..71b2b626564 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -28,7 +28,7 @@ from datetime import timedelta
from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar, Literal, NoReturn
-from pydantic import BaseModel, TypeAdapter
+from pydantic import TypeAdapter
from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDenied
from pydantic_ai.capabilities import AbstractCapability, Toolset,
WrapperCapability
from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse,
ToolCallPart
@@ -60,6 +60,7 @@ from airflow.providers.common.ai.utils.logging import (
log_run_summary,
log_run_usage,
)
+from airflow.providers.common.ai.utils.output_type import dump_output_for_xcom
from airflow.providers.common.ai.utils.prompt_cache import PromptCaching
from airflow.providers.common.ai.utils.toolset_base import ensure_masked
from airflow.providers.common.ai.utils.toolsets import iter_toolsets
@@ -93,18 +94,18 @@ try:
# See LLMOperator: Newer ``apache-airflow-task-sdk`` versions register
declared ``output_type`` classes
# from a worker-side DAG walk, so the model instance flows through XCom;
older
# ``apache-airflow-task-sdk`` versions without the walk dump to a dict
instead.
- from airflow.sdk.serde import (
- SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER,
- iter_pydantic_models,
- )
+ from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as
_CORE_WALKER
except ImportError: # pragma: no cover - missing ``apache-airflow-task-sdk``
walker
_CORE_WALKER = False
if TYPE_CHECKING:
+ from types import UnionType
+
import jinja2
from pydantic_ai import Agent
from pydantic_ai.capabilities import AgentCapability
from pydantic_ai.messages import ModelMessage
+ from pydantic_ai.output import OutputSpec
from pydantic_ai.usage import UsageLimits
from airflow.providers.common.ai.durable.base import DurableStorageProtocol
@@ -241,26 +242,6 @@ def _declares_agent_template_fields(toolset: Any) -> bool:
)
-def _dump_unregistered_models(value: Any, registered: set[type]) -> Any:
- """Return *value* with each model whose class is not in *registered*
dumped to a JSON-compatible dict."""
- if isinstance(value, BaseModel):
- if type(value) in registered:
- return value
- # A python-mode model_dump() leaves values such as plain Enum members,
HttpUrl and Path in the
- # dict, and the default XCom serializer raises TypeError for those
values.
- return value.model_dump(mode="json")
- # The default XCom serializer already pushes a subclass of list, tuple or
dict, such as a NamedTuple
- # or an OrderedDict, as the plain type. Building the plain type here does
not change what a
- # downstream task receives.
- if isinstance(value, list):
- return [_dump_unregistered_models(item, registered) for item in value]
- if isinstance(value, tuple):
- return tuple(_dump_unregistered_models(item, registered) for item in
value)
- if isinstance(value, dict):
- return {key: _dump_unregistered_models(item, registered) for key, item
in value.items()}
- return value
-
-
# CancellableAgentRunMixin must precede BaseOperator so its on_kill overrides
BaseOperator's
# no-op. The other mixins only add methods, so they can trail BaseOperator.
See the MRO guard
# test in tests/unit/common/ai/mixins/test_cancellable_run.py.
@@ -296,10 +277,16 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook`
for how blank entries in the list are dropped.
:param system_prompt: System-level instructions for the agent.
- :param output_type: Expected output type. Default ``str``. Set to a
Pydantic
- ``BaseModel`` subclass for structured output; the model instance is
+ :param output_type: Expected output type. Default ``str``. A single type
+ (``str``, ``int``, ``list[str]``, a Pydantic ``BaseModel`` subclass)
or a
+ union (``A | B``). A list of types (``[str, int]``) is also accepted
and treated
+ as alternatives (a union). pydantic-ai output markers
(``ToolOutput(...)``,
+ ``NativeOutput(...)``) and output functions are also accepted.
+ For a ``BaseModel`` subclass, the model instance is
returned to XCom unchanged so downstream tasks can type-hint it
- directly. The class must be defined at module scope -- nested classes
+ directly, except inside a list of types, where it is returned as a
``dict``
+ because the worker registers only the classes it finds in a single
type or a union.
+ The class must be defined at module scope -- nested classes
cannot be deserialized from XCom.
:param toolsets: List of pydantic-ai toolsets the agent can use
(e.g. ``SQLToolset``, ``HookToolset``). The connection IDs of
@@ -491,11 +478,11 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
:param tool_approval_assigned_users: Users allowed to decide. ``None``
(default)
leaves it to anyone who can act on the task's Required Actions.
- :param serialize_output: If ``True`` and ``output_type`` is a Pydantic
- ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
- ``model_dump()`` before being pushed to XCom. Default ``False`` --
- the Pydantic instance flows through XCom unchanged. Set to ``True``
- when a downstream consumer needs the dict shape.
+ :param serialize_output: If ``True``, an output that is itself a Pydantic
``BaseModel`` is
+ dumped with ``model_dump()`` before being pushed to XCom, so ``date``
and ``Decimal``
+ fields keep their types. Default ``False``: a model whose class the
worker registers from
+ ``output_type`` flows through XCom unchanged. Some outputs are pushed
as dicts either way;
+ see :ref:`When a model arrives as a dict
<structured-output-dict-shape>`.
"""
deserialization_allowed_class_fields: ClassVar[tuple[str, ...]] =
("output_type",)
@@ -527,7 +514,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator,
HITLReviewMixin, Usa
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type = str,
+ output_type: OutputSpec[Any] | UnionType = str,
toolsets: list[AbstractToolset] | None = None,
capabilities: list[AgentCapability[Any]] | None = None,
enable_tool_logging: bool = True,
@@ -1140,20 +1127,7 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
"""Return the XCom value, with dicts in place of models that a
downstream task could not load."""
if not self.do_xcom_push:
return output
- if self._serialize_model_output and isinstance(output, BaseModel):
- return output.model_dump()
- if not _CORE_WALKER:
- # iter_pydantic_models could not be imported, so the worker
registers no class.
- # _serialize_model_output is always True when that import fails,
so a model at the top level was
- # already dumped above. The class of a model inside a list, tuple
or dict is not registered
- # either, so this call passes an empty set of registered classes
to dump that model too.
- return _dump_unregistered_models(output, set())
- # A downstream task can load a model from XCom only if the worker
registered the model's class.
- # The worker registers the classes that iter_pydantic_models finds in
output_type.
- # iter_pydantic_models finds no class in a list of output types, an
output marker, an output
- # function, a TypeAliasType or a NewType. For these output types the
operator pushes each model
- # as a dict, including models inside a list, tuple or dict.
- return _dump_unregistered_models(output,
set(iter_pydantic_models(self.output_type)))
+ return dump_output_for_xcom(output, self.output_type,
serialize_output=self._serialize_model_output)
def _pause_for_tool_approval(self, context: Context, result: Any, *,
attempt_usage: RunUsage) -> NoReturn:
"""
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index 5c4eceafd2e..0ed6273962f 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -23,7 +23,6 @@ from datetime import timedelta
from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar, Literal
-from pydantic import BaseModel
from pydantic_ai.usage import RunUsage
from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
@@ -56,7 +55,11 @@ from airflow.providers.common.ai.utils.logging import (
log_run_summary,
log_run_usage,
)
-from airflow.providers.common.ai.utils.output_type import dump_output_to_json,
rehydrate_pydantic_output
+from airflow.providers.common.ai.utils.output_type import (
+ dump_output_for_xcom,
+ dump_output_to_json,
+ rehydrate_pydantic_output,
+)
from airflow.providers.common.ai.utils.prompt import describe_prompt
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
from airflow.providers.common.ai.utils.usage_budget import copy_run_usage,
subtract_run_usage
@@ -82,6 +85,7 @@ if TYPE_CHECKING:
from pydantic_ai import Agent
from pydantic_ai.agent import AgentRunResult
+ from pydantic_ai.output import OutputSpec
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import UsageLimits
@@ -104,6 +108,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
via a Pydantic ``BaseModel``. When ``output_type`` is a ``BaseModel``
subclass,
the model instance is returned to XCom unchanged so downstream tasks can
type-hint it directly (e.g. ``def downstream(result: MyModel) -> None``).
+ Inside a list of types (``[MyModel, int]``) the model is returned as a
``dict`` instead.
The class is auto-registered for deserialization in each process that
parses
the DAG, so no edit to ``[core] allowed_deserialization_classes`` is
required.
The Pydantic class must be defined at module scope: classes nested inside
@@ -143,9 +148,14 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
:param system_prompt: System-level instructions for the LLM agent.
:param output_type: Expected output type. Default ``str``. A single type
(``str``, ``int``, ``list[str]``, a Pydantic ``BaseModel`` subclass)
or a
- union (``A | B``). For a ``BaseModel`` subclass, the model instance is
+ union (``A | B``). A list of types (``[str, int]``) is also accepted
and treated
+ as alternatives (a union). pydantic-ai output markers
(``ToolOutput(...)``,
+ ``NativeOutput(...)``) and output functions are also accepted.
+ For a ``BaseModel`` subclass, the model instance is
returned to XCom unchanged so downstream tasks can type-hint it
- directly. The class must be defined at module scope -- nested classes
+ directly, except inside a list of types, where it is returned as a
``dict``
+ because the worker registers only the classes it finds in a single
type or a union.
+ The class must be defined at module scope -- nested classes
cannot be deserialized from XCom.
:param model_settings: Pydantic AI
:class:`~pydantic_ai.settings.ModelSettings`
for the model call, such as ``temperature`` or ``max_tokens``,
including
@@ -228,12 +238,11 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
connection does not silently switch off a control the author set.
Independent of
``require_approval``, which always asks. ``on_uncertain="review"``
needs Airflow 3.1+,
like ``require_approval``. Default ``None``: no gate, behaviour
unchanged.
- :param serialize_output: If ``True`` and ``output_type`` is a Pydantic
- ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
- ``model_dump()`` before being pushed to XCom. Default ``False`` --
- the Pydantic instance flows through XCom unchanged. Set to ``True``
- when a downstream consumer needs the dict shape (e.g. sending to an
- external system that expects JSON-style payloads).
+ :param serialize_output: If ``True``, an output that is itself a Pydantic
``BaseModel`` is
+ dumped with ``model_dump()`` before being pushed to XCom, so ``date``
and ``Decimal``
+ fields keep their types. Default ``False``: a model whose class the
worker registers from
+ ``output_type`` flows through XCom unchanged. Some outputs are pushed
as dicts either way;
+ see :ref:`When a model arrives as a dict
<structured-output-dict-shape>`.
"""
deserialization_allowed_class_fields: ClassVar[tuple[str, ...]] =
("output_type",)
@@ -263,7 +272,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type | UnionType = str,
+ output_type: OutputSpec[Any] | UnionType = str,
model_settings: ModelSettings | None = None,
output_retries: int | None = None,
usage_limits: UsageLimits | dict[str, Any] | None = None,
@@ -431,11 +440,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
)
self.defer_for_approval(context, output, body=body,
decision=record) # type: ignore[misc]
- if self._serialize_model_output and isinstance(output, BaseModel):
- # ``serialize_output=True``, or an ``apache-airflow-task-sdk``
version without the worker-side
- # deserialization-class walk: dump to a dict so XCom carries a
plain
- # JSON payload that deserializes without an allow-list entry.
- output = output.model_dump()
+ output = dump_output_for_xcom(output, self.output_type,
serialize_output=self._serialize_model_output)
self._clear_usage_budget(context)
return output
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
index 08b37247df8..1bb2997a4e0 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
@@ -21,10 +21,9 @@ from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, ClassVar
-from pydantic import BaseModel
-
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.utils.file_analysis import
build_file_analysis_request
+from airflow.providers.common.ai.utils.output_type import dump_output_for_xcom
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
if TYPE_CHECKING:
@@ -174,12 +173,7 @@ class LLMFileAnalysisOperator(LLMOperator):
if self.require_approval:
self.defer_for_approval(context, output) # type: ignore[misc]
- if self._serialize_model_output and isinstance(output, BaseModel):
- # Honour ``serialize_output=True`` and the older-Airflow fallback,
- # matching ``LLMOperator.execute``. ``LLMFileAnalysisOperator``
- # overrides ``execute`` rather than calling ``super().execute()``,
- # so the dump step has to be repeated here.
- output = output.model_dump()
+ output = dump_output_for_xcom(output, self.output_type,
serialize_output=self._serialize_model_output)
self._clear_usage_budget(context)
return output
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
index 2eda37a280e..b8fa3a2cf77 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
@@ -39,12 +39,13 @@ from airflow.providers.common.ai.utils.masking import
dumps_masked
from airflow.providers.common.ai.utils.query_results import (
DEFAULT_MAX_COLUMNS,
DEFAULT_MAX_RESULT_BYTES,
- GET_SCHEMA_TOOL_DESCRIPTION as _GET_SCHEMA_DESCRIPTION,
QUERY_TOOL_DESCRIPTION as _QUERY_DESCRIPTION,
+ build_get_schema_description,
build_query_result,
build_schema_result,
)
from airflow.providers.common.ai.utils.tool_definition import
build_args_validator
+from airflow.providers.common.ai.utils.tool_prefix import ToolPrefixMixin
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
validate_max_retries
if TYPE_CHECKING:
@@ -92,7 +93,7 @@ _RETRYABLE_QUERY_ERROR_PATTERNS = (
)
-class DataFusionToolset(AirflowToolset):
+class DataFusionToolset(ToolPrefixMixin, AirflowToolset):
"""
Curated toolset that gives an LLM agent SQL access to object-storage data
via Apache DataFusion.
@@ -137,6 +138,10 @@ class DataFusionToolset(AirflowToolset):
the full list is replaced by a bounded summary (column count, a type
histogram, a
sample of columns) that points the agent at the ``name_contains``
filter, so a
several-thousand-column table cannot exhaust the context before a
query is written.
+ :param tool_prefix: Optional prefix joined to each tool name with an
underscore
+ (e.g. ``"lake"`` produces ``"lake_query"``). Must be a valid Python
identifier. Set it
+ when one agent has another toolset with the same tool names, such as a
``SQLToolset``,
+ since duplicate tool names are rejected.
:param max_retries: How many times the model may correct a failed call to
one of these
tools before the run fails. ``None`` (the default) uses the agent's
tool retry
budget, its ``retries``, as pydantic-ai's own toolsets do.
@@ -150,9 +155,11 @@ class DataFusionToolset(AirflowToolset):
max_rows: int = 50,
max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES,
max_columns: int = DEFAULT_MAX_COLUMNS,
+ tool_prefix: str = "",
max_retries: int | None = None,
) -> None:
self._max_retries = validate_max_retries(max_retries)
+ self._set_tool_prefix(tool_prefix)
if not datasource_configs:
raise ValueError("datasource_configs must contain at least one
DataSourceConfig")
self._datasource_configs = datasource_configs
@@ -165,7 +172,7 @@ class DataFusionToolset(AirflowToolset):
@property
def id(self) -> str:
suffix = "_".join(config.table_name.replace("-", "_") for config in
self._datasource_configs)
- return f"sql_datafusion_{suffix}"
+ return f"sql_datafusion_{suffix}{self._get_tool_prefix_id_suffix('_')}"
def _get_engine(self) -> DataFusionEngine:
"""Lazily create and configure a DataFusionEngine from
*datasource_configs*."""
@@ -180,11 +187,16 @@ class DataFusionToolset(AirflowToolset):
max_retries = self._get_tool_max_retries(ctx)
tools: dict[str, ToolsetTool[Any]] = {}
- for name, description, schema in (
+ for base, description, schema in (
("list_tables", "List available table names.",
_LIST_TABLES_SCHEMA),
- ("get_schema", _GET_SCHEMA_DESCRIPTION, _GET_SCHEMA_SCHEMA),
+ (
+ "get_schema",
+
build_get_schema_description(tool_name=self._build_tool_name("get_schema")),
+ _GET_SCHEMA_SCHEMA,
+ ),
("query", _QUERY_DESCRIPTION, _QUERY_SCHEMA),
):
+ name = self._build_tool_name(base)
tool_def = ToolDefinition(
name=name,
description=description,
@@ -207,13 +219,14 @@ class DataFusionToolset(AirflowToolset):
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- if name == "list_tables":
+ base = self._find_base_tool_name(name)
+ if base == "list_tables":
return await self.run_blocking(self._list_tables)
- if name == "get_schema":
+ if base == "get_schema":
return await self.run_blocking(
self._get_schema, tool_args["table_name"],
tool_args.get("name_contains")
)
- if name == "query":
+ if base == "query":
return await self.run_blocking(self._query, tool_args["sql"])
raise ValueError(f"Unknown tool: {name!r}")
@@ -245,6 +258,7 @@ class DataFusionToolset(AirflowToolset):
max_columns=self._max_columns,
max_result_bytes=self._max_result_bytes,
name_contains=name_contains,
+ tool_name=self._build_tool_name("get_schema"),
)
def _query(self, sql: str) -> str:
@@ -277,7 +291,8 @@ class DataFusionToolset(AirflowToolset):
except QueryExecutionException as ex:
if self._is_retryable_query_error(ex):
raise ModelRetry(
- f"error: {ex!s}, Use get_schema and list_tables tools for
more details."
+ f"error: {ex!s}, Use {self._build_tool_name('get_schema')}
and "
+ f"{self._build_tool_name('list_tables')} tools for more
details."
) from ex
return dumps_masked({"error": str(ex), "query": sql})
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 7c7f57c6c4b..a3b03d931dc 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
@@ -39,6 +39,7 @@ from airflow.providers.common.ai.utils.tool_definition import
(
return_schema_kwargs,
serialize_for_llm,
)
+from airflow.providers.common.ai.utils.tool_prefix import validate_tool_prefix
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
validate_max_retries
if TYPE_CHECKING:
@@ -156,10 +157,7 @@ class HookToolset(AirflowToolset):
pinned_arguments: dict[str, Any] | None = None,
max_retries: int | None = None,
) -> None:
- if tool_prefix and not tool_prefix.isidentifier():
- # The prefixed names are rendered as Python function signatures
under
- # code mode, so a name that is not an identifier breaks there.
- raise ValueError(f"tool_prefix must be a valid Python identifier,
got {tool_prefix!r}.")
+ validate_tool_prefix(tool_prefix)
if tool_name_prefix is not None:
if tool_prefix and f"{tool_prefix}_" != tool_name_prefix:
raise ValueError(
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
index 43b860995ea..71c2d4f9493 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py
@@ -67,8 +67,8 @@ class MCPToolset(AirflowToolset):
:param mcp_conn_id: Airflow connection ID for the MCP server. Templated
when
the toolset is passed to ``AgentOperator`` / ``@task.agent``.
- :param tool_prefix: Optional prefix prepended to tool names
- (e.g. ``"weather"`` → ``"weather_get_forecast"``).
+ :param tool_prefix: Optional prefix joined to each tool name with an
underscore
+ (e.g. ``"weather"`` → ``"weather_get_forecast"``). Defaults to ``""``
(no prefix).
:param token_provider: Optional zero-argument callable returning a bearer
token string, overriding the connection ``password`` for HTTP/SSE auth.
Called once, the first time this toolset establishes a connection.
@@ -86,7 +86,7 @@ class MCPToolset(AirflowToolset):
self,
mcp_conn_id: str,
*,
- tool_prefix: str | None = None,
+ tool_prefix: str = "",
token_provider: Callable[[], str] | None = None,
env_provider: Callable[[], dict[str, str]] | None = None,
) -> None:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py
index 53e2217485d..ee72c6ff01e 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py
@@ -44,6 +44,7 @@ from airflow.providers.common.ai.utils.tool_definition import
(
return_schema_kwargs,
serialize_for_llm,
)
+from airflow.providers.common.ai.utils.tool_prefix import ToolPrefixMixin
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
validate_max_retries
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, ObjectStoragePath
@@ -109,7 +110,7 @@ _COLUMNAR_FORMATS: tuple[Literal["parquet", "avro"], ...] =
("parquet", "avro")
_MEDIA_FORMATS = frozenset({"jpeg", "jpg", "pdf", "png"})
-class ObjectStorageToolset(AirflowToolset):
+class ObjectStorageToolset(ToolPrefixMixin, AirflowToolset):
"""
Give an agent read-only access to the files under one object-storage path.
@@ -171,28 +172,22 @@ class ObjectStorageToolset(AirflowToolset):
):
if value < 1:
raise ValueError(f"{name} must be at least 1, got {value}.")
- if tool_prefix and not tool_prefix.isidentifier():
- raise ValueError(f"tool_prefix must be a valid Python identifier,
got {tool_prefix!r}.")
+ self._set_tool_prefix(tool_prefix)
self._path = path
self._conn_id = conn_id
self._max_files = max_files
self._max_read_bytes = max_read_bytes
self._max_output_bytes = max_output_bytes
- self._tool_prefix = tool_prefix
@property
def id(self) -> str:
- suffix = f"-{self._tool_prefix}" if self._tool_prefix else ""
- return f"object-storage-{self._conn_id or 'default'}{suffix}"
-
- def _tool_name(self, base: str) -> str:
- return f"{self._tool_prefix}_{base}" if self._tool_prefix else base
+ return f"object-storage-{self._conn_id or
'default'}{self._get_tool_prefix_id_suffix()}"
async def get_tools(self, ctx: RunContext[Any]) -> dict[str,
ToolsetTool[Any]]:
max_retries = self._get_tool_max_retries(ctx)
tools: dict[str, ToolsetTool[Any]] = {}
for base, schema in _SCHEMAS.items():
- name = self._tool_name(base)
+ name = self._build_tool_name(base)
tools[name] = ToolsetTool(
toolset=self,
tool_def=ToolDefinition(
@@ -214,7 +209,7 @@ class ObjectStorageToolset(AirflowToolset):
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> str:
- base = name.removeprefix(f"{self._tool_prefix}_") if self._tool_prefix
else name
+ base = self._find_base_tool_name(name)
relative = tool_args.get("path") or ""
try:
if base == LIST_FILES:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
index a03c968b390..4424c5df932 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
@@ -55,6 +55,7 @@ from airflow.providers.common.ai.utils.tool_definition import
(
code_arg_kwargs,
return_schema_kwargs,
)
+from airflow.providers.common.ai.utils.tool_prefix import ToolPrefixMixin
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
from airflow.providers.common.compat.sdk import ObjectStoragePath,
get_current_context
@@ -167,7 +168,7 @@ _DESCRIPTIONS = {
}
-class SandboxToolset(AirflowToolset):
+class SandboxToolset(ToolPrefixMixin, AirflowToolset):
"""
Give an agent shell and file access inside a disposable sandbox, off the
Airflow worker.
@@ -291,10 +292,7 @@ class SandboxToolset(AirflowToolset):
f"default_command_timeout ({default_command_timeout}) must not
exceed "
f"max_command_timeout ({max_command_timeout})."
)
- if tool_prefix and not tool_prefix.isidentifier():
- # The prefixed names are rendered as Python function signatures
under
- # code mode, so a name that is not an identifier breaks there.
- raise ValueError(f"tool_prefix must be a valid Python identifier,
got {tool_prefix!r}.")
+ self._set_tool_prefix(tool_prefix)
if attach_to is not None:
if not isinstance(backend, AttachableSandboxBackend):
raise ValueError(
@@ -355,7 +353,6 @@ class SandboxToolset(AirflowToolset):
self._exports = dict(exports or {})
self._export_conn_id = export_conn_id
self._max_export_bytes = int(max_export_bytes)
- self._tool_prefix = tool_prefix
self._sandbox: str | None = None
self._create_future: concurrent.futures.Future[str] | None = None
# Set while attached: who this run claimed the sandbox as, when the
sandbox ends
@@ -378,19 +375,10 @@ class SandboxToolset(AirflowToolset):
@property
def id(self) -> str:
- suffix = f"-{self._tool_prefix}" if self._tool_prefix else ""
- return f"sandbox-{self._backend.name}{suffix}"
-
- def _tool_name(self, base: str) -> str:
- return f"{self._tool_prefix}_{base}" if self._tool_prefix else base
+ return
f"sandbox-{self._backend.name}{self._get_tool_prefix_id_suffix()}"
def _base_name(self, tool_name: str) -> str | None:
- if not self._tool_prefix:
- return tool_name if tool_name in _SCHEMAS else None
- prefix = f"{self._tool_prefix}_"
- if not tool_name.startswith(prefix):
- return None
- base = tool_name[len(prefix) :]
+ base = self._find_base_tool_name(tool_name)
return base if base in _SCHEMAS else None
async def for_run(self, ctx: RunContext[Any]) -> AbstractToolset[Any]:
@@ -860,7 +848,7 @@ class SandboxToolset(AirflowToolset):
# run_command additionally carries code_arg metadata so code mode
# leaves it native instead of folding a shell surface into
run_code.
extra = code_arg_kwargs("command", "shell") if base == RUN_COMMAND
else {}
- name = self._tool_name(base)
+ name = self._build_tool_name(base)
description = _DESCRIPTIONS[base]
if base == RUN_COMMAND:
description = f"{description}
{self._network_note}{self._export_note}"
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
index 1311c9ee157..5da4eaaf48e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
@@ -44,12 +44,13 @@ from airflow.providers.common.ai.utils.masking import
dumps_masked
from airflow.providers.common.ai.utils.query_results import (
DEFAULT_MAX_COLUMNS,
DEFAULT_MAX_RESULT_BYTES,
- GET_SCHEMA_TOOL_DESCRIPTION as _GET_SCHEMA_DESCRIPTION,
QUERY_TOOL_DESCRIPTION as _QUERY_DESCRIPTION,
+ build_get_schema_description,
build_query_result,
build_schema_result,
)
from airflow.providers.common.ai.utils.tool_definition import
build_args_validator, return_schema_kwargs
+from airflow.providers.common.ai.utils.tool_prefix import ToolPrefixMixin
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
validate_max_retries
from airflow.providers.common.compat.sdk import BaseHook
@@ -173,7 +174,7 @@ class _CappedFetch:
return rows
-class SQLToolset(AirflowToolset):
+class SQLToolset(ToolPrefixMixin, AirflowToolset):
"""
Curated toolset that gives an LLM agent safe access to a SQL database.
@@ -266,6 +267,10 @@ class SQLToolset(AirflowToolset):
the full list is replaced by a bounded summary (column count, a type
histogram, a
sample of columns) that points the agent at the ``name_contains``
filter, so a
several-thousand-column table cannot exhaust the context before a
query is written.
+ :param tool_prefix: Optional prefix joined to each tool name with an
underscore
+ (e.g. ``"pg"`` produces ``"pg_query"``). Must be a valid Python
identifier. Set it
+ when one agent has another toolset with the same tool names, such as a
second
+ ``SQLToolset``, since duplicate tool names are rejected.
:param max_retries: How many times the model may correct a failed call to
one of these
tools before the run fails. ``None`` (the default) uses the agent's
tool retry
budget, its ``retries``, as pydantic-ai's own toolsets do.
@@ -286,9 +291,11 @@ class SQLToolset(AirflowToolset):
max_rows: int = 50,
max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES,
max_columns: int = DEFAULT_MAX_COLUMNS,
+ tool_prefix: str = "",
max_retries: int | None = None,
) -> None:
self._max_retries = validate_max_retries(max_retries)
+ self._set_tool_prefix(tool_prefix)
self._allowed_tables: frozenset[str] | None
if allowed_tables is _UNSET:
self._allowed_tables = None
@@ -369,7 +376,7 @@ class SQLToolset(AirflowToolset):
@property
def id(self) -> str:
- return f"sql-{self._db_conn_id}"
+ return f"sql-{self._db_conn_id}{self._get_tool_prefix_id_suffix()}"
# ------------------------------------------------------------------
# Lazy hook resolution
@@ -395,9 +402,13 @@ class SQLToolset(AirflowToolset):
max_retries = self._get_tool_max_retries(ctx)
tools: dict[str, ToolsetTool[Any]] = {}
- for name, description, schema in (
+ for base, description, schema in (
("list_tables", "List available table names in the database.",
_LIST_TABLES_SCHEMA),
- ("get_schema", _GET_SCHEMA_DESCRIPTION, _GET_SCHEMA_SCHEMA),
+ (
+ "get_schema",
+
build_get_schema_description(tool_name=self._build_tool_name("get_schema")),
+ _GET_SCHEMA_SCHEMA,
+ ),
("query", _QUERY_DESCRIPTION, _QUERY_SCHEMA),
("check_query", "Validate SQL syntax without executing it.",
_CHECK_QUERY_SCHEMA),
):
@@ -406,6 +417,7 @@ class SQLToolset(AirflowToolset):
# blocking calls of the other toolsets that use it.
# return_schema is "string": every tool returns a JSON-encoded
string
# (json.dumps), so code mode renders `-> str` instead of `-> Any`.
+ name = self._build_tool_name(base)
tool_def = ToolDefinition(
name=name,
description=description,
@@ -429,10 +441,11 @@ class SQLToolset(AirflowToolset):
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- if name not in ("list_tables", "get_schema", "query", "check_query"):
+ base = self._find_base_tool_name(name)
+ if base not in ("list_tables", "get_schema", "query", "check_query"):
raise ValueError(f"Unknown tool: {name!r}")
try:
- return await self.run_blocking(self._run_tool, name, tool_args)
+ return await self.run_blocking(self._run_tool, base, tool_args)
except Exception as e:
# Hand the database's own error back to the agent as a retry so it
can
# read the message and fix its SQL within the run. pydantic-ai
bounds
@@ -441,7 +454,8 @@ class SQLToolset(AirflowToolset):
# for Airflow to retry, rather than being silently worked around.
raise ModelRetry(
f"The {name} tool failed: {e}\n"
- "Use the list_tables and get_schema tools to inspect the
database, "
+ f"Use the {self._build_tool_name('list_tables')} and
{self._build_tool_name('get_schema')} tools "
+ "to inspect the database, "
"then fix the query and try again."
) from e
@@ -504,6 +518,7 @@ class SQLToolset(AirflowToolset):
max_columns=self._max_columns,
max_result_bytes=self._max_result_bytes,
name_contains=name_contains,
+ tool_name=self._build_tool_name("get_schema"),
)
def _dialect_for_validation(self) -> str | None:
@@ -593,7 +608,7 @@ class SQLToolset(AirflowToolset):
raise SQLSafetyError(
f"Query uses a data source that cannot be checked against
allowed_tables: "
f"{'; '.join(scan.unverifiable_sources)}. Query the allowed
tables directly: "
- f"use list_tables to see them."
+ f"use {self._build_tool_name('list_tables')} to see them."
)
disallowed = [
".".join(part for part in (catalog, schema, table) if part)
@@ -603,5 +618,5 @@ class SQLToolset(AirflowToolset):
if disallowed:
raise SQLSafetyError(
f"Query references tables that are not in the allowed tables
list: "
- f"{', '.join(sorted(set(disallowed)))}. Use list_tables to see
the allowed tables."
+ f"{', '.join(sorted(set(disallowed)))}. Use
{self._build_tool_name('list_tables')} to see the allowed tables."
)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
index 70de2cd73b4..03ce698608a 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
@@ -20,14 +20,22 @@ from __future__ import annotations
import inspect
import logging
+import operator
+from collections.abc import Iterator
+from functools import reduce
from typing import Any, get_origin
from pydantic import BaseModel, TypeAdapter, ValidationError
from pydantic.errors import PydanticUserError
-from pydantic_core import CoreSchema, PydanticSerializationError
+from pydantic_core import CoreSchema, PydanticSerializationError,
to_jsonable_python
from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
+try:
+ from airflow.sdk.serde import iter_pydantic_models
+except ImportError: # pragma: no cover - missing ``apache-airflow-task-sdk``
walker
+ iter_pydantic_models = None # type: ignore[assignment]
+
log = logging.getLogger(__name__)
_MAX_SHOWN_CHARS = 200
@@ -38,6 +46,59 @@ _MAX_SHOWN_CHARS = 200
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError,
KeyError, IndexError)
+def _is_model_dump_required(output: Any, output_type: Any) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to
XCom dumped.
+
+ The instances are looked up in the output itself, as
``_iter_models_in_output`` walks it: the elements
+ of a list, tuple (a NamedTuple included), set or frozenset, and the keys
and values of a dict. A
+ downstream task can load a model from XCom only if the worker registered
the model's class, and the
+ worker registers the classes ``iter_pydantic_models`` finds in
+ ``output_type``. It finds none in a list of output types, an output
marker, an output function, a
+ ``TypeAliasType`` or a ``NewType``, so an instance of a class it did not
find is dumped. Without
+ ``iter_pydantic_models`` (Airflow < 3.3 or an older task-sdk) the worker
registers no class, so every
+ model in the output counts as unregistered and is dumped.
+ """
+ registered = set() if iter_pydantic_models is None else
set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in
_iter_models_in_output(output))
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output:
bool) -> Any:
+ """
+ Return the output as it is pushed to XCom.
+
+ With ``serialize_output`` a model output is dumped with ``model_dump()``,
as 1.0.0 does, so a ``date``
+ or ``Decimal`` field keeps its type. Any other output is dumped, whole and
to JSON values, only when it
+ holds a model the worker registers no class for, with field names rather
than aliases, so a ``HttpUrl``,
+ ``Path`` or ``Enum`` field is a string that XCom accepts. Registered
models, in a container or not, are
+ left for the serde serializer.
+ """
+ if serialize_output and isinstance(output, BaseModel):
+ return output.model_dump()
+ if not _is_model_dump_required(output, output_type):
+ return output
+ return to_jsonable_python(output, by_alias=False)
+
+
+def _iter_models_in_output(output: Any) -> Iterator[BaseModel]:
+ """
+ Yield the ``BaseModel`` instances in an output.
+
+ Looks inside list, tuple (a NamedTuple included), set, frozenset and dict.
Dataclass contents and
+ model fields are not walked.
+ """
+ stack = [output]
+ while stack:
+ item = stack.pop()
+ if isinstance(item, BaseModel):
+ yield item
+ elif isinstance(item, (set, frozenset, list, tuple)):
+ stack.extend(item)
+ elif isinstance(item, dict):
+ stack.extend(item)
+ stack.extend(item.values())
+
+
def dump_output_to_json(output: Any) -> str:
"""Serialize an output as JSON that ``rehydrate_pydantic_output``
validates back, or ``str()`` if pydantic cannot."""
# ``round_trip`` drops computed fields, and ``by_alias=False`` keeps field
names even when a model
@@ -72,7 +133,7 @@ def rehydrate_pydantic_output(
"""
Turn the reviewed string back into ``output_type``.
- ``raw`` is JSON, or bare text for types pydantic builds from a string.
``str`` outputs, output
+ ``raw`` is JSON, or bare text for types pydantic builds from a string. A
list of candidate types is validated as their union. ``str`` outputs, output
functions, pydantic-ai output markers such as ``ToolOutput`` and types
pydantic has no schema for
return ``raw`` unchanged, and so does an unedited output that does not
validate back (with a warning).
@@ -82,11 +143,19 @@ def rehydrate_pydantic_output(
"""
if output_type is str:
return raw
+ declared_type = output_type
+ if isinstance(output_type, (list, tuple)):
+ # pydantic-ai treats a list of output types as "any of these".
Unioning an empty list or output
+ # markers is a ``TypeError``, and leaves the output as it is.
+ try:
+ output_type = reduce(operator.or_, output_type)
+ except TypeError:
+ return raw
+ # A pydantic-ai output marker (``ToolOutput``, ``NativeOutput``, ...) is
not a type, and
+ # ``TypeAdapter`` fails on it with an ``AttributeError``.
try:
adapter: TypeAdapter[Any] = TypeAdapter(output_type)
is_output_function = _is_output_function(adapter.core_schema)
- # A pydantic-ai output marker (``ToolOutput``, ``NativeOutput``, ...) is
not a type, and
- # ``TypeAdapter`` fails on it with an ``AttributeError``.
except (PydanticUserError, AttributeError):
return raw
if is_output_function:
@@ -105,9 +174,7 @@ def rehydrate_pydantic_output(
log.warning("%s Returning the output unchanged.", message)
return raw
raise ReviewedOutputValidationError(message) from json_error
- if serialize_output and isinstance(rehydrated, BaseModel):
- return rehydrated.model_dump()
- return rehydrated
+ return dump_output_for_xcom(rehydrated, declared_type,
serialize_output=serialize_output)
def _validate_bare_text(adapter: TypeAdapter[Any], raw: str) -> Any:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py
index 09b03c17cf4..eff89771fdc 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py
@@ -86,19 +86,24 @@ QUERY_TOOL_DESCRIPTION = (
#: unlike query rows they cannot be silently dropped: the description states
the ``name_contains``
#: filter and the summary shape so a truncated result is read as "narrow your
request", not as a
#: small table.
-GET_SCHEMA_TOOL_DESCRIPTION = (
+_GET_SCHEMA_DESCRIPTION_TEMPLATE = (
"Get a table's columns. Returns JSON of the form "
- '{"columns": [{"name": ..., "type": ...}, ...], "column_count": N}. Pass
`name_contains` to '
+ '{{"columns": [{{"name": ..., "type": ...}}, ...], "column_count": N}}.
Pass `name_contains` to '
"return only the columns whose name contains that substring
(case-insensitive) -- use it to "
"find the columns relevant to your question on a wide table. When a table
has more columns "
"than can be returned at once, the full list is replaced by a summary
(`column_count`, and "
"where the budget allows a `type_histogram` and a `sample_columns`
preview) with `truncated` "
- "set, `truncated_by` naming the limit that was hit, and a `hint`. Call
get_schema again with "
+ "set, `truncated_by` naming the limit that was hit, and a `hint`. Call
{tool_name} again with "
"a `name_contains` substring to retrieve the specific columns you need
instead of the whole "
"table."
)
+def build_get_schema_description(*, tool_name: str = "get_schema") -> str:
+ """Render the ``get_schema`` tool description naming the tool as the model
sees it."""
+ return _GET_SCHEMA_DESCRIPTION_TEMPLATE.format(tool_name=tool_name)
+
+
def _dumps(payload: Any) -> str:
return dumps_masked(payload, **_DUMP_KWARGS)
@@ -222,6 +227,7 @@ def build_schema_result(
max_columns: int,
max_result_bytes: int,
name_contains: str | None = None,
+ tool_name: str = "get_schema",
) -> str:
"""
Render a table's columns as a bounded JSON tool result.
@@ -238,6 +244,8 @@ def build_schema_result(
:param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
columns are considered and the value is echoed back so a filtered
subset is never mistaken
for the whole table.
+ :param tool_name: Name of the ``get_schema`` tool as the model sees it
(including any
+ ``tool_prefix``), so hints point at this toolset's own tool.
"""
name_contains = name_contains or None
total_columns = len(columns)
@@ -258,7 +266,7 @@ def build_schema_result(
)
else:
advice = (
- f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"No columns match name_contains={name_contains!r}. Call
{tool_name} without "
f"name_contains to list all {total_columns} column{plural}."
)
return _dumps(
@@ -284,6 +292,7 @@ def build_schema_result(
max_columns=max_columns,
max_result_bytes=max_result_bytes,
name_contains=name_contains,
+ tool_name=tool_name,
)
@@ -294,6 +303,7 @@ def _summarize_schema(
max_columns: int,
max_result_bytes: int,
name_contains: str | None,
+ tool_name: str,
) -> str:
"""Build the bounded summary returned when the full column list does not
fit."""
# Count is checked before bytes: it is the cheaper, more explainable
bound. A bytes-only
@@ -309,7 +319,7 @@ def _summarize_schema(
move = (
"Use a more specific name_contains substring to narrow to the columns
you need."
if name_contains is not None
- else "Call get_schema with a name_contains substring to return only
the columns you need."
+ else f"Call {tool_name} with a name_contains substring to return only
the columns you need."
)
output: dict[str, Any] = {
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_prefix.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_prefix.py
new file mode 100644
index 00000000000..d7c2a3ee7fc
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_prefix.py
@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The ``tool_prefix`` handling shared by the toolsets that namespace their
tool names."""
+
+from __future__ import annotations
+
+
+def validate_tool_prefix(tool_prefix: str) -> None:
+ """Raise ``ValueError`` if a non-empty ``tool_prefix`` is not a valid
Python identifier."""
+ if tool_prefix and not tool_prefix.isidentifier():
+ # The prefixed names are rendered as Python function signatures under
+ # code mode, so a name that is not an identifier breaks there.
+ raise ValueError(f"tool_prefix must be a valid Python identifier, got
{tool_prefix!r}.")
+
+
+class ToolPrefixMixin:
+ """Build, and map back from, the tool names and the toolset id that a
``tool_prefix`` produces."""
+
+ _tool_prefix: str = ""
+
+ tool_name_conflict_hint: str = (
+ "Pass a distinct `tool_prefix=...` to this toolset so its tool names
and "
+ "model-facing hints are renamed together; wrapping it in
`PrefixedToolset` "
+ "renames only the tools."
+ )
+
+ def _set_tool_prefix(self, tool_prefix: str) -> None:
+ validate_tool_prefix(tool_prefix)
+ self._tool_prefix = tool_prefix
+
+ def _build_tool_name(self, base: str) -> str:
+ return f"{self._tool_prefix}_{base}" if self._tool_prefix else base
+
+ def _find_base_tool_name(self, name: str) -> str | None:
+ """Return ``name`` without the prefix, or ``None`` when a prefix is
set and ``name`` does not carry it."""
+ if not self._tool_prefix:
+ return name
+ prefix = f"{self._tool_prefix}_"
+ return name.removeprefix(prefix) if name.startswith(prefix) else None
+
+ def _get_tool_prefix_id_suffix(self, separator: str = "-") -> str:
+ return f"{separator}{self._tool_prefix}" if self._tool_prefix else ""
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
index d0efd7f64dd..e7e67c82e2a 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
@@ -44,6 +44,9 @@ class TestMCPHookInit:
hook = MCPHook(mcp_conn_id="my_mcp")
assert hook.mcp_conn_id == "my_mcp"
+ def test_tool_prefix_defaults_to_empty_string(self):
+ assert MCPHook(mcp_conn_id="my_mcp").tool_prefix == ""
+
def test_tool_prefix(self):
hook = MCPHook(mcp_conn_id="my_mcp", tool_prefix="weather")
assert hook.tool_prefix == "weather"
@@ -99,6 +102,21 @@ class TestMCPHookGetConn:
headers={"Authorization": "Bearer my-secret-token"},
)
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_HTTP_TRANSPORT, autospec=True)
+ def test_empty_tool_prefix_does_not_wrap(self, mock_transport_cls,
mock_toolset_cls):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ host="http://localhost:3001/mcp";,
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ result = hook.get_conn()
+
+ mock_toolset_cls.return_value.prefixed.assert_not_called()
+ assert result is mock_toolset_cls.return_value
+
@patch(_MCP_TOOLSET, autospec=True)
@patch(_HTTP_TRANSPORT, autospec=True)
def test_passes_tool_prefix(self, mock_transport_cls, mock_toolset_cls):
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 66ed0f025da..2e9a92f1340 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -18,12 +18,13 @@ from __future__ import annotations
import asyncio
import dataclasses
+import inspect
import json
import logging
import sys
from collections import OrderedDict
from contextlib import nullcontext
-from datetime import timedelta
+from datetime import date, timedelta
from decimal import Decimal
from enum import Enum
from types import ModuleType, SimpleNamespace
@@ -32,7 +33,7 @@ from unittest.mock import ANY, MagicMock, PropertyMock, call,
patch
from uuid import UUID
import pytest
-from pydantic import BaseModel, ConfigDict, HttpUrl
+from pydantic import BaseModel, ConfigDict, Field, HttpUrl
from pydantic_ai import Agent, DeferredToolRequests, RunContext, Tool,
ToolOutput
from pydantic_ai.capabilities import (
AbstractCapability,
@@ -73,6 +74,7 @@ from airflow.providers.common.ai.durable.storage import
DurableStorage
from airflow.providers.common.ai.exceptions import UnsupportedToolDeferralError
from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.operators.agent import AgentOperator,
HITLReviewLink
+from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.sandbox.base import (
HOLDER_TAG,
OWNER_TAG,
@@ -141,6 +143,11 @@ class SummaryKind(Enum):
SHORT = "short"
+class SerializedByAlias(BaseModel):
+ model_config = ConfigDict(serialize_by_alias=True)
+ full_name: str = Field(serialization_alias="fullName")
+
+
class LinkedSummary(BaseModel):
kind: SummaryKind
url: HttpUrl
@@ -789,6 +796,21 @@ class TestAgentOperatorExecute:
"What is the answer?", usage_limits=None, run_id="ti-1",
cancellation_token=ANY, usage=ANY
)
+ def test_output_type_is_annotated_like_llm_operator(self):
+ annotation =
inspect.signature(AgentOperator.__init__).parameters["output_type"].annotation
+ assert annotation ==
inspect.signature(LLMOperator.__init__).parameters["output_type"].annotation
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_passes_a_union_output_type_to_the_agent(self,
mock_hook_cls, make_mock_run_result):
+ mock_agent = _make_mock_agent(Summary(text="ok"), make_mock_run_result)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm",
output_type=Summary | int)
+ op.execute(context=_make_context())
+
+ expected = [Summary | int, DeferredToolRequests] if AIRFLOW_V_3_3_PLUS
else Summary | int
+ assert
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs["output_type"]
== expected
+
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_keys_the_run_by_attempt_without_a_task_instance_id(
self, mock_hook_cls, make_mock_run_result
@@ -1124,7 +1146,7 @@ class TestAgentOperatorExecute:
pytest.param(
ToolOutput(tuple[Summary, int]),
(Summary(text="ok"), 1),
- ({"text": "ok", "score": 0.0}, 1),
+ [{"text": "ok", "score": 0.0}, 1],
False,
id="tuple-in-output-marker",
),
@@ -1138,14 +1160,14 @@ class TestAgentOperatorExecute:
pytest.param(
ToolOutput(list[tuple[dict[str, list[Summary]], int]]),
[({"a": [Summary(text="ok")]}, 1)],
- [({"a": [{"text": "ok", "score": 0.0}]}, 1)],
+ [[{"a": [{"text": "ok", "score": 0.0}]}, 1]],
False,
id="nested-containers-in-output-marker",
),
pytest.param(
ToolOutput(SummaryPair),
SummaryPair(Summary(text="ok"), 1),
- ({"text": "ok", "score": 0.0}, 1),
+ [{"text": "ok", "score": 0.0}, 1],
False,
id="named-tuple-in-output-marker",
),
@@ -1238,9 +1260,9 @@ class TestAgentOperatorExecute:
pytest.param(list[Summary], id="list"),
],
)
- @patch("airflow.providers.common.ai.operators.agent._CORE_WALKER", False)
+
@patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
- def test_execute_dumps_nested_models_without_the_worker_walk(
+ def test_execute_pushes_nested_models_as_dicts_without_the_worker_walk(
self, mock_hook_cls, make_mock_run_result, output_type
):
mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
@@ -1253,6 +1275,51 @@ class TestAgentOperatorExecute:
result = op.execute(context=MagicMock())
assert result == [{"text": "ok", "score": 0.0}]
+ assert type(result[0]) is dict
+
+ @pytest.mark.parametrize(
+ ("output_type", "output", "serialize_output"),
+ [
+ pytest.param(
+ [SerializedByAlias, int], SerializedByAlias(full_name="Ann"),
False, id="serialize-by-alias"
+ ),
+ pytest.param(
+ ToolOutput(tuple[Summary, date]),
+ (Summary(text="ok"), date(2026, 1, 2)),
+ False,
+ id="date-next-to-model-in-tuple",
+ ),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_pushes_the_same_value_as_the_llm_operator(
+ self,
+ mock_agent_hook_cls,
+ mock_llm_hook_cls,
+ make_mock_run_result,
+ output_type,
+ output,
+ serialize_output,
+ ):
+ mock_agent_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ output, make_mock_run_result
+ )
+ mock_llm_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ output, make_mock_run_result
+ )
+ kwargs = {
+ "prompt": "Analyze this",
+ "llm_conn_id": "my_llm",
+ "output_type": output_type,
+ "serialize_output": serialize_output,
+ }
+
+ agent_result = AgentOperator(task_id="agent",
**kwargs).execute(context=MagicMock())
+ llm_result = LLMOperator(task_id="llm",
**kwargs).execute(context=MagicMock())
+
+ assert agent_result == llm_result
+ assert type(agent_result) is type(llm_result)
@pytest.mark.parametrize(
"output",
@@ -1452,6 +1519,7 @@ class TestAgentOperatorExecute:
[Summary(text="ok")],
[{"text": "ok", "score": 0.0}],
id="list-in-output-marker",
+ marks=requires_typed_xcom,
),
pytest.param(
ToolOutput(LinkedSummary),
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index 5cac9f17d38..618bd144ccd 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -18,13 +18,13 @@ from __future__ import annotations
import contextlib
import logging
-from datetime import timedelta
+from datetime import date, timedelta
from decimal import Decimal
from unittest.mock import ANY, MagicMock, patch
from uuid import uuid4
import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, HttpUrl
from pydantic_ai import Agent
from pydantic_ai.exceptions import UnexpectedModelBehavior, UsageLimitExceeded
from pydantic_ai.messages import BinaryContent, ModelMessage, ModelResponse,
TextPart
@@ -48,7 +48,11 @@ from tests_common.test_utils.compat import
OperatorSerialization
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
try:
- from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as
_CORE_WALKER
+ from airflow.sdk.serde import (
+ SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER,
+ deserialize,
+ serialize,
+ )
except ImportError:
_CORE_WALKER = False
@@ -85,6 +89,14 @@ class Entities(BaseModel):
names: list[str]
+class Dated(BaseModel):
+ d: date
+
+
+class Link(BaseModel):
+ url: HttpUrl
+
+
class Summary(BaseModel):
text: str
@@ -494,6 +506,63 @@ class TestLLMOperator:
assert result == {"names": ["A", "B"]}
assert not isinstance(result, Entities)
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_serialize_output_keeps_python_field_types(self,
mock_hook_cls, make_mock_run_result):
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value =
make_mock_run_result(Dated(d=date(2026, 1, 2)))
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=Dated, serialize_output=True)
+ result = op.execute(context=MagicMock())
+
+ assert type(result) is dict
+ assert type(result["d"]) is date
+
+ @requires_typed_xcom
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def
test_execute_dumps_a_model_the_worker_cannot_register_from_a_list_of_types(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ """The worker registers no class from a list of types, so the model
goes to XCom as a dict."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value =
make_mock_run_result(Entities(names=["A", "B"]))
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=[Entities, int])
+ result = op.execute(context=MagicMock())
+
+ assert result == {"names": ["A", "B"]}
+
+ @requires_typed_xcom
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_dumps_models_inside_a_list_output_under_a_list_of_types(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ """A list output of ``[list[Entities], str]`` is no model, yet its
models need dumping for the worker."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value =
make_mock_run_result([Entities(names=["A"])])
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=[list[Entities], str])
+ result = op.execute(context=MagicMock())
+
+ assert result == [{"names": ["A"]}]
+ assert deserialize(serialize(result)) == [{"names": ["A"]}]
+
+ @requires_typed_xcom
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_dumps_http_url_fields_so_xcom_can_serialize_them(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value =
make_mock_run_result(Link(url="https://example.com/a";))
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=[Link, str])
+ result = op.execute(context=MagicMock())
+
+ assert deserialize(serialize(result)) == {"url":
"https://example.com/a"}
+
def _result(make_mock_run_result, output, details, *, cost=None):
result = make_mock_run_result(output, cost=cost)
@@ -1659,8 +1728,9 @@ class TestLLMOperatorApproval:
(int, "5", 5),
(list[str], '["a","b"]', ["a", "b"]),
pytest.param(Summary, '{"text":"hello"}', Summary(text="hello"),
marks=requires_typed_xcom),
+ pytest.param([Summary, int], '{"text":"hello"}', {"text":
"hello"}, marks=requires_typed_xcom),
],
- ids=["int", "list", "basemodel"],
+ ids=["int", "list", "basemodel", "basemodel-in-a-list-of-types"],
)
def test_execute_complete_restores_non_str_output_type(self, output_type,
generated_output, expected):
op = LLMOperator(
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
index 61feeafa241..29fc50fad5d 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
@@ -275,6 +275,34 @@ class TestLLMFileAnalysisOperator:
assert result == {"findings": ["error spike"]}
assert not isinstance(result, Summary)
+ @requires_typed_xcom
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ @patch(
+
"airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request",
autospec=True
+ )
+ def
test_execute_dumps_a_model_the_worker_cannot_register_from_a_list_of_types(
+ self, mock_build_request, mock_hook_cls
+ ):
+ mock_build_request.return_value = FileAnalysisRequest(
+ user_content="prepared prompt",
+ resolved_paths=["/tmp/app.log"],
+ total_size_bytes=10,
+ )
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value =
_make_mock_run_result(Summary(findings=["error spike"]))
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMFileAnalysisOperator(
+ task_id="test",
+ prompt="Summarize the file",
+ llm_conn_id="my_llm",
+ file_path="/tmp/app.log",
+ output_type=[Summary, int],
+ )
+ result = op.execute(context={})
+
+ assert result == {"findings": ["error spike"]}
+
@patch(
"airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request",
autospec=True
)
diff --git
a/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
index 9e6372a96f9..f149912da23 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
@@ -107,6 +107,96 @@ class TestDataFusionToolsetGetTools:
assert tool.tool_def.description
+class TestDataFusionToolsetToolPrefix:
+ def test_prefix_is_joined_to_every_tool_name_with_an_underscore(self):
+ ts = DataFusionToolset([_make_mock_datasource_config()],
tool_prefix="lake")
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock(spec=RunContext)))
+ assert set(tools.keys()) == {"lake_list_tables", "lake_get_schema",
"lake_query"}
+ assert all(name == tool.tool_def.name for name, tool in tools.items())
+
+ def test_prefix_is_part_of_the_toolset_id(self):
+ ts = DataFusionToolset([_make_mock_datasource_config("orders")],
tool_prefix="lake")
+ assert ts.id == "sql_datafusion_orders_lake"
+
+ @pytest.mark.parametrize("tool_prefix", ["my-lake", "1abc", "has space"])
+ def test_prefix_must_be_identifier(self, tool_prefix):
+ with pytest.raises(ValueError, match="tool_prefix must be a valid
Python identifier"):
+ DataFusionToolset([_make_mock_datasource_config()],
tool_prefix=tool_prefix)
+
+ def test_prefixed_tool_dispatches_to_the_unprefixed_implementation(self):
+ ts = DataFusionToolset([_make_mock_datasource_config()],
tool_prefix="lake")
+ ts._engine = _make_mock_engine(registered_tables={"sales":
"s3://bucket/sales/"})
+
+ result = asyncio.run(
+ ts.call_tool(
+ "lake_list_tables", {}, ctx=MagicMock(spec=RunContext),
tool=MagicMock(spec=ToolsetTool)
+ )
+ )
+
+ assert json.loads(result) == ["sales"]
+
+ def test_unprefixed_name_is_unknown_when_a_prefix_is_set(self):
+ ts = DataFusionToolset([_make_mock_datasource_config()],
tool_prefix="lake")
+
+ with pytest.raises(ValueError, match="Unknown tool: 'list_tables'"):
+ asyncio.run(
+ ts.call_tool(
+ "list_tables", {}, ctx=MagicMock(spec=RunContext),
tool=MagicMock(spec=ToolsetTool)
+ )
+ )
+
+
+class TestDataFusionToolsetNameConflictHint:
+ def test_hint_points_to_tool_prefix(self):
+ ts = DataFusionToolset([_make_mock_datasource_config()])
+ assert "tool_prefix" in ts.tool_name_conflict_hint
+
+
+class TestDataFusionToolsetToolPrefixModelFacingText:
+ def test_get_schema_description_names_the_prefixed_tool(self):
+ ts = DataFusionToolset([_make_mock_datasource_config()],
tool_prefix="lake")
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock(spec=RunContext)))
+ assert "Call lake_get_schema again" in
tools["lake_get_schema"].tool_def.description
+
+ def test_get_schema_hint_names_the_prefixed_tool(self):
+ ts = DataFusionToolset([_make_mock_datasource_config()],
tool_prefix="lake", max_columns=2)
+ ts._engine = _make_mock_engine(
+ registered_tables={"sales": "s3://bucket/sales/"},
+ schema_fields=[(f"c{i}", "Int64") for i in range(5)],
+ )
+
+ result = asyncio.run(
+ ts.call_tool(
+ "lake_get_schema",
+ {"table_name": "sales"},
+ ctx=MagicMock(spec=RunContext),
+ tool=MagicMock(spec=ToolsetTool),
+ )
+ )
+
+ assert "Call lake_get_schema with a name_contains substring" in
json.loads(result)["hint"]
+
+ def test_query_retry_names_the_prefixed_tools(self, tmp_path: Path):
+ csv_path = tmp_path / "trips.csv"
+ csv_path.write_text("pickup_date,id\n2024-01-01,1\n")
+ ts = DataFusionToolset(
+ [DataSourceConfig(conn_id="", table_name="trips",
uri=f"file://{csv_path}", format="csv")],
+ tool_prefix="lake",
+ )
+
+ with pytest.raises(ModelRetry) as exc_info:
+ asyncio.run(
+ ts.call_tool(
+ "lake_query",
+ {"sql": "SELECT name FROM trips"},
+ ctx=MagicMock(spec=RunContext),
+ tool=MagicMock(spec=ToolsetTool),
+ )
+ )
+
+ assert "Use lake_get_schema and lake_list_tables tools" in
exc_info.value.message
+
+
class TestDataFusionToolsetArgsValidation:
@pytest.mark.parametrize(
("tool_name", "valid_args"),
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
index 5391d62e0e0..edaad7de4c5 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_mcp.py
@@ -36,6 +36,9 @@ class TestMCPToolsetInit:
ts = MCPToolset("my_mcp_server", tool_prefix="weather")
assert ts._tool_prefix == "weather"
+ def test_tool_prefix_defaults_to_empty_string(self):
+ assert MCPToolset("my_mcp_server")._tool_prefix == ""
+
class TestMCPToolsetGetServer:
@patch(_HOOK_PATH, autospec=True)
@@ -47,7 +50,7 @@ class TestMCPToolsetGetServer:
server = ts._get_server()
mock_hook_cls.assert_called_once_with(
- mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=None,
env_provider=None
+ mcp_conn_id="mcp_conn", tool_prefix="", token_provider=None,
env_provider=None
)
mock_hook_cls.return_value.get_conn.assert_called_once()
assert server is mock_server
@@ -74,7 +77,7 @@ class TestMCPToolsetGetServer:
ts._get_server()
mock_hook_cls.assert_called_once_with(
- mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=provider,
env_provider=None
+ mcp_conn_id="mcp_conn", tool_prefix="", token_provider=provider,
env_provider=None
)
@patch(_HOOK_PATH, autospec=True)
@@ -88,7 +91,7 @@ class TestMCPToolsetGetServer:
ts._get_server()
mock_hook_cls.assert_called_once_with(
- mcp_conn_id="mcp_conn", tool_prefix=None, token_provider=None,
env_provider=env_provider
+ mcp_conn_id="mcp_conn", tool_prefix="", token_provider=None,
env_provider=env_provider
)
@patch(_HOOK_PATH, autospec=True)
diff --git
a/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py
index 0dc15cdd5b4..83a08f451a0 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py
@@ -22,13 +22,14 @@ import io
import json
import uuid
from typing import Any
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
import pytest
from fsspec.implementations.memory import MemoryFileSystem
from pydantic_ai import RunContext
from pydantic_ai.exceptions import ToolFailed
from pydantic_ai.models.test import TestModel
+from pydantic_ai.toolsets.abstract import ToolsetTool
from pydantic_ai.usage import RunUsage
from airflow.providers.common.ai.toolsets.object_storage import
ObjectStorageToolset
@@ -111,6 +112,16 @@ class TestTools:
)
assert listing["path"] == "/"
+ def test_unprefixed_name_is_unknown_when_a_prefix_is_set(self, storage):
+ toolset = ObjectStorageToolset(f"file://{storage}",
tool_prefix="reports")
+ ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage())
+
+ with pytest.raises(ValueError, match="Unknown tool: 'list_files'"):
+ asyncio.run(toolset.call_tool("list_files", {}, ctx,
MagicMock(spec=ToolsetTool)))
+
+ def test_hint_points_to_tool_prefix(self, storage):
+ assert "tool_prefix" in
ObjectStorageToolset(f"file://{storage}").tool_name_conflict_hint
+
@pytest.mark.parametrize(
("kwargs", "match"),
[
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
index 4c291cd2f76..562e978d2f4 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
@@ -193,6 +193,9 @@ class TestInit:
with pytest.raises(ValueError, match="tool_prefix"):
SandboxToolset(_RecordingBackend(), tool_prefix=bad)
+ def test_hint_points_to_tool_prefix(self):
+ assert "tool_prefix" in
SandboxToolset(_RecordingBackend()).tool_name_conflict_hint
+
def test_id_includes_backend_name_and_prefix(self):
assert SandboxToolset(_RecordingBackend()).id == "sandbox-rec"
assert SandboxToolset(_RecordingBackend(), tool_prefix="local").id ==
"sandbox-rec-local"
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
index dd16a1c1b17..57256030a57 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
@@ -148,6 +148,120 @@ class TestSQLToolsetGetTools:
assert tool.tool_def.return_schema == {"type": "string"}
+class TestSQLToolsetToolPrefix:
+ def test_prefix_is_joined_to_every_tool_name_with_an_underscore(self):
+ ts = SQLToolset("pg_default", tool_prefix="pg")
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+ assert set(tools.keys()) == {"pg_list_tables", "pg_get_schema",
"pg_query", "pg_check_query"}
+ assert all(name == tool.tool_def.name for name, tool in tools.items())
+
+ def test_prefix_is_part_of_the_toolset_id(self):
+ assert SQLToolset("my_pg", tool_prefix="pg").id == "sql-my_pg-pg"
+
+ @pytest.mark.parametrize("tool_prefix", ["my-db", "1abc", "has space"])
+ def test_prefix_must_be_identifier(self, tool_prefix):
+ with pytest.raises(ValueError, match="tool_prefix must be a valid
Python identifier"):
+ SQLToolset("my_pg", tool_prefix=tool_prefix)
+
+ def test_prefixed_tool_dispatches_to_the_unprefixed_implementation(self):
+ ts = SQLToolset("pg_default", tool_prefix="pg")
+ ts._hook = _make_mock_db_hook(table_names=["users"])
+
+ result = asyncio.run(
+ ts.call_tool(
+ "pg_list_tables", {}, ctx=MagicMock(spec=RunContext),
tool=MagicMock(spec=ToolsetTool)
+ )
+ )
+
+ assert json.loads(result) == ["users"]
+
+ def test_unprefixed_name_is_unknown_when_a_prefix_is_set(self):
+ ts = SQLToolset("pg_default", tool_prefix="pg")
+
+ with pytest.raises(ValueError, match="Unknown tool: 'list_tables'"):
+ asyncio.run(
+ ts.call_tool(
+ "list_tables", {}, ctx=MagicMock(spec=RunContext),
tool=MagicMock(spec=ToolsetTool)
+ )
+ )
+
+
+def _call_prefixed(ts: SQLToolset, name: str, args: dict):
+ return asyncio.run(
+ ts.call_tool(name, args, ctx=MagicMock(spec=RunContext),
tool=MagicMock(spec=ToolsetTool))
+ )
+
+
+class TestSQLToolsetNameConflictHint:
+ def test_hint_points_to_tool_prefix(self):
+ assert "tool_prefix" in
SQLToolset("pg_default").tool_name_conflict_hint
+
+
+class TestSQLToolsetToolPrefixModelFacingText:
+ """Text shown to the model names this toolset's own tools, not the bare
base names."""
+
+ def test_get_schema_description_names_the_prefixed_tool(self):
+ tools = asyncio.run(SQLToolset("pg_default",
tool_prefix="pg").get_tools(ctx=MagicMock()))
+ description = tools["pg_get_schema"].tool_def.description
+ assert "Call pg_get_schema again" in description
+
+ def test_get_schema_description_without_prefix_is_unchanged(self):
+ tools =
asyncio.run(SQLToolset("pg_default").get_tools(ctx=MagicMock()))
+ assert "Call get_schema again" in
tools["get_schema"].tool_def.description
+
+ def test_tool_failure_retry_names_the_prefixed_tools(self):
+ ts = SQLToolset("pg_default", tool_prefix="pg")
+ ts._hook = _make_mock_db_hook()
+ ts._hook.run.side_effect = RuntimeError("boom")
+
+ with pytest.raises(ModelRetry) as exc_info:
+ _call_prefixed(ts, "pg_query", {"sql": "SELECT 1"})
+
+ assert "Use the pg_list_tables and pg_get_schema tools" in
exc_info.value.message
+
+ @pytest.mark.parametrize(
+ ("sql", "expected"),
+ [
+ ("SELECT * FROM dblink('host=evil', 'SELECT 1') AS t(x int)", "use
pg_list_tables to see them"),
+ ("SELECT * FROM secret_salaries", "Use pg_list_tables to see the
allowed tables"),
+ ],
+ ids=["unverifiable_source", "disallowed_table"],
+ )
+ def test_allowed_tables_refusal_names_the_prefixed_tool(self, sql,
expected):
+ ts = SQLToolset("pg_default", allowed_tables=["orders"],
tool_prefix="pg")
+ ts._hook = _make_mock_db_hook()
+ ts._hook.dialect_name = "postgresql"
+
+ with pytest.raises(ModelRetry) as exc_info:
+ _call_prefixed(ts, "pg_query", {"sql": sql})
+
+ assert expected in exc_info.value.message
+
+ @pytest.mark.parametrize(
+ ("table_schema", "args", "expected"),
+ [
+ (
+ [{"name": f"c{i}", "type": "INT"} for i in range(5)],
+ {"table_name": "users", "name_contains": "zzz"},
+ "Call pg_get_schema without name_contains",
+ ),
+ (
+ [{"name": f"c{i}", "type": "INT"} for i in range(300)],
+ {"table_name": "users"},
+ "Call pg_get_schema with a name_contains substring",
+ ),
+ ],
+ ids=["no_columns_match", "summary"],
+ )
+ def test_get_schema_hint_names_the_prefixed_tool(self, table_schema, args,
expected):
+ ts = SQLToolset("pg_default", tool_prefix="pg")
+ ts._hook = _make_mock_db_hook(table_schema=table_schema)
+
+ hint = json.loads(_call_prefixed(ts, "pg_get_schema", args))["hint"]
+
+ assert expected in hint
+
+
class TestSQLToolsetListTables:
def test_returns_all_tables(self):
ts = SQLToolset("pg_default")
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
b/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
index cd22f989c3e..213e6fc8bab 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
@@ -16,32 +16,103 @@
# under the License.
from __future__ import annotations
+import json
from collections import deque
from collections.abc import Sequence
from dataclasses import dataclass
+from datetime import UTC, date, datetime
+from decimal import Decimal
from enum import Enum
+from pathlib import Path
from typing import Annotated, Literal, NamedTuple
+from unittest import mock
+from uuid import UUID
import pytest
-from pydantic import BaseModel, ConfigDict, Field, computed_field,
model_validator
+from pydantic import BaseModel, ConfigDict, Field, HttpUrl, computed_field,
model_validator
from pydantic_ai import NativeOutput, PromptedOutput, ToolOutput
from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
from airflow.providers.common.ai.utils.output_type import (
_is_container_mismatch,
+ _is_model_dump_required,
+ dump_output_for_xcom,
dump_output_to_json,
rehydrate_pydantic_output,
)
+try:
+ from airflow.sdk.serde import (
+ SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER,
+ deserialize,
+ serialize,
+ )
+except ImportError:
+ _CORE_WALKER = False
+
+requires_walker = pytest.mark.skipif(
+ not _CORE_WALKER,
+ reason="Requires an ``apache-airflow-task-sdk`` version with the
worker-side deserialization-class walk.",
+)
+
class A(BaseModel):
x: int
+class SubA(A):
+ pass
+
+
+class Item(BaseModel):
+ n: int
+
+
class B(BaseModel):
y: str
+class Stamped(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ d: date
+ at: datetime
+ amount: Decimal
+ ident: UUID
+
+
+STAMPED_JSON = {
+ "d": "2026-01-02",
+ "at": "2026-01-02T03:04:05Z",
+ "amount": "1.5",
+ "ident": "12345678-1234-5678-1234-567812345678",
+}
+STAMPED = Stamped(
+ d=date(2026, 1, 2),
+ at=datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC),
+ amount=Decimal("1.5"),
+ ident=UUID("12345678-1234-5678-1234-567812345678"),
+)
+
+
+class Link(BaseModel):
+ url: HttpUrl
+
+
+LINK = Link(url=HttpUrl("https://example.com/a";))
+LINK_SERIALIZED = {
+ "__classname__": f"{Link.__module__}.Link",
+ "__version__": 1,
+ "__data__": {"url": "https://example.com/a"},
+}
+
+
+class Key(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ n: int
+
+
class FromAttributesModel(BaseModel):
model_config = ConfigDict(from_attributes=True)
x: int
@@ -56,6 +127,12 @@ class Color(str, Enum):
RED = "red"
+class Page(BaseModel):
+ url: HttpUrl
+ kind: Color
+ path: Path
+
+
class ValidationAlias(BaseModel):
user_name: str = Field(alias="userName")
@@ -140,17 +217,246 @@ class UnresolvedForwardReference:
value: Missing # type: ignore[name-defined] # noqa: F821
+class ItemPair(NamedTuple):
+ item: Item
+
+
class Point(NamedTuple):
x: int
y: str
+class TestIsModelDumpRequired:
+ @pytest.mark.parametrize(
+ ("output", "output_type", "expected"),
+ [
+ (A(x=1), [A, int], True),
+ (A(x=1), (A, int), True),
+ (A(x=1), ToolOutput(A), True),
+ (A(x=1), A, False),
+ (A(x=1), A | int, False),
+ (A(x=1), list[A], False),
+ (SubA(x=1), A, True),
+ (5, [A, int], False),
+ ],
+ ids=[
+ "list-of-types",
+ "tuple-of-types",
+ "output-marker",
+ "single-model",
+ "union",
+ "generic-list",
+ "subclass-instance",
+ "not-a-model",
+ ],
+ )
+ @requires_walker
+ def test_is_true_when_the_worker_registers_no_class_for_the_output(self,
output, output_type, expected):
+ result = _is_model_dump_required(output, output_type)
+
+ assert result is expected
+
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def test_requires_a_dump_without_the_sdk_walker(self):
+ assert _is_model_dump_required(A(x=1), [A, int]) is True
+ assert _is_model_dump_required([A(x=1)], list[A]) is True
+ assert _is_model_dump_required(5, list[A]) is False
+
+
+class TestIsModelDumpRequiredForContainerOutputs:
+ @pytest.mark.parametrize(
+ ("output", "output_type"),
+ [
+ ([A(x=1)], list[A]),
+ ({"k": A(x=1)}, dict[str, A]),
+ (None, A | None),
+ ("text", [list[Item], str]),
+ ],
+ ids=["registered-list", "registered-dict", "none-for-optional",
"no-model-in-output"],
+ )
+ @requires_walker
+ def
test_is_false_for_a_registered_shape_or_an_output_without_a_model(self, output,
output_type):
+ assert _is_model_dump_required(output, output_type) is False
+
+ @pytest.mark.parametrize(
+ ("output", "output_type"),
+ [
+ ([Item(n=1)], [list[Item], str]),
+ ({"k": Item(n=1)}, [dict[str, Item], str]),
+ ((Item(n=1),), [tuple[Item, ...], str]),
+ (ItemPair(item=Item(n=1)), [ItemPair, str]),
+ ],
+ ids=["list", "dict", "tuple", "named-tuple"],
+ )
+ @requires_walker
+ def
test_is_true_for_models_inside_a_container_the_worker_registers_no_class_for(
+ self, output, output_type
+ ):
+ assert _is_model_dump_required(output, output_type) is True
+
+
+class TestDumpOutputForXcom:
+ @requires_walker
+ def
test_dumped_model_with_http_url_enum_path_serializes_through_serde(self):
+ page = Page(url="https://example.com/a";, kind=Color.RED,
path=Path("/tmp/a"))
+
+ dumped = dump_output_for_xcom(page, [Page, str],
serialize_output=False)
+
+ assert deserialize(serialize(dumped)) == {
+ "url": "https://example.com/a";,
+ "kind": "red",
+ "path": "/tmp/a",
+ }
+
+ @requires_walker
+ def
test_dumps_models_inside_a_list_when_output_type_is_a_list_of_types(self):
+ dumped = dump_output_for_xcom([Item(n=1)], [list[Item], str],
serialize_output=False)
+
+ assert dumped == [{"n": 1}]
+ assert deserialize(serialize(dumped)) == [{"n": 1}]
+
+ @requires_walker
+ def test_keeps_a_registered_container_as_model_instances(self):
+ output = {"k": A(x=1)}
+
+ dumped = dump_output_for_xcom(output, dict[str, A],
serialize_output=False)
+
+ assert dumped is output
+ assert isinstance(dumped["k"], A)
+
+ @pytest.mark.parametrize(
+ "output",
+ [datetime(2026, 1, 1, tzinfo=UTC), 5, "s", None, {"a": 1}],
+ ids=["datetime", "int", "str", "none", "dict"],
+ )
+ def test_non_model_outputs_are_returned_unchanged(self, output):
+ dumped = dump_output_for_xcom(output, [A, int], serialize_output=True)
+
+ assert dumped is output
+
+ @requires_walker
+ def
test_dumps_the_whole_output_when_only_some_of_its_models_are_unregistered(self):
+ output = [A(x=1), Item(n=1)]
+
+ assert _is_model_dump_required(output, A) is True
+ assert dump_output_for_xcom(output, A, serialize_output=False) ==
[{"x": 1}, {"n": 1}]
+
+ @pytest.mark.parametrize("model_cls", [SerializationAlias,
SerializedByAlias])
+ @requires_walker
+ def test_dumps_field_names_not_aliases_like_dump_output_to_json(self,
model_cls):
+ model = model_cls(user_name="kax")
+
+ dumped = dump_output_for_xcom(model, [model_cls, str],
serialize_output=False)
+
+ assert dumped == {"user_name": "kax"}
+ assert dumped == json.loads(dump_output_to_json(model))
+
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def test_dumps_a_container_without_the_sdk_walker(self):
+ output = [A(x=1)]
+
+ for serialize_output in (True, False):
+ dumped = dump_output_for_xcom(output, list[A],
serialize_output=serialize_output)
+
+ assert dumped == [{"x": 1}]
+ assert dumped is not output
+
+ @pytest.mark.parametrize(
+ ("output", "output_type", "expected"),
+ [
+ (A(x=1), A, {"x": 1}),
+ ([A(x=1)], list[A], [{"x": 1}]),
+ ],
+ ids=["top-level-model", "list-of-models"],
+ )
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def test_dumps_every_model_without_the_sdk_walker_when_not_serializing(
+ self, output, output_type, expected
+ ):
+ dumped = dump_output_for_xcom(output, output_type,
serialize_output=False)
+
+ assert dumped == expected
+ assert dumped is not output
+
+ def test_serialize_output_keeps_python_field_types(self):
+ dumped = dump_output_for_xcom(STAMPED, Stamped, serialize_output=True)
+
+ assert type(dumped) is dict
+ assert dumped == {"d": STAMPED.d, "at": STAMPED.at, "amount":
STAMPED.amount, "ident": STAMPED.ident}
+ assert [type(value) for value in dumped.values()] == [date, datetime,
Decimal, UUID]
+
+ @pytest.mark.parametrize(
+ ("output", "expected"),
+ [
+ ([STAMPED], [STAMPED_JSON]),
+ ({"k": STAMPED}, {"k": STAMPED_JSON}),
+ ((STAMPED,), [STAMPED_JSON]),
+ ({STAMPED}, [STAMPED_JSON]),
+ ],
+ ids=["list", "dict", "tuple", "set"],
+ )
+ @requires_walker
+ def
test_serialize_output_dumps_an_unregistered_container_to_json_values(self,
output, expected):
+ dumped = dump_output_for_xcom(output, [Stamped, str],
serialize_output=True)
+
+ assert dumped == expected
+
+ @requires_walker
+ def test_dumps_a_model_used_as_a_dict_key(self):
+ output = {Key(n=1): 1, "k": 2}
+
+ dumped = dump_output_for_xcom(output, [Key, str],
serialize_output=True)
+
+ assert dumped == {"n=1": 1, "k": 2}
+ assert json.dumps(dumped)
+
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def
test_serialize_output_dumps_a_top_level_model_without_the_sdk_walker(self):
+ dumped = dump_output_for_xcom(STAMPED, Stamped, serialize_output=True)
+
+ assert type(dumped) is dict
+ assert type(dumped["d"]) is date
+
+ @pytest.mark.parametrize(
+ ("output", "output_type", "expected"),
+ [
+ ([LINK], list[Link], [LINK_SERIALIZED]),
+ ({"k": LINK}, dict[str, Link], {"k": LINK_SERIALIZED}),
+ ],
+ ids=["list", "dict"],
+ )
+ @requires_walker
+ def
test_serialize_output_keeps_registered_models_in_a_container_that_serde_pushes(
+ self, output, output_type, expected
+ ):
+ dumped = dump_output_for_xcom(output, output_type,
serialize_output=True)
+
+ assert dumped is output
+ assert serialize(dumped) == expected
+
+ @requires_walker
+ def test_unregistered_model_is_dumped_to_json_values(self):
+ dumped = dump_output_for_xcom(STAMPED, [Stamped, str],
serialize_output=False)
+
+ assert dumped == {
+ "d": "2026-01-02",
+ "at": STAMPED.at.isoformat().replace("+00:00", "Z"),
+ "amount": "1.5",
+ "ident": str(STAMPED.ident),
+ }
+
+
class TestRehydratePydanticOutput:
+ @requires_walker
def test_returns_model_instance(self):
result = rehydrate_pydantic_output(A, '{"x": 7}',
serialize_output=False)
assert isinstance(result, A)
assert result.x == 7
+
@mock.patch("airflow.providers.common.ai.utils.output_type.iter_pydantic_models",
None)
+ def test_returns_dict_without_the_sdk_walker(self):
+ assert rehydrate_pydantic_output(A, '{"x": 7}',
serialize_output=False) == {"x": 7}
+
def test_returns_dict_when_serialize_output(self):
result = rehydrate_pydantic_output(A, '{"x": 7}',
serialize_output=True)
assert result == {"x": 7}
@@ -187,6 +493,38 @@ class TestRehydratePydanticOutput:
assert result == expected
assert type(result) is type(expected)
+ @pytest.mark.parametrize(
+ ("output_type", "raw", "expected"),
+ [
+ ([str, int], '"x"', "x"),
+ ([int, float], "1.5", 1.5),
+ ((int, float), "2", 2),
+ pytest.param([A, int], '{"x": 7}', {"x": 7},
marks=requires_walker),
+ ],
+ ids=["str-int-string", "int-float", "tuple",
"model-in-a-list-of-types-is-dumped"],
+ )
+ def test_validates_list_output_type_as_union(self, output_type, raw,
expected):
+ assert rehydrate_pydantic_output(output_type, raw,
serialize_output=False) == expected
+
+ @pytest.mark.parametrize(
+ "output_type",
+ [[], [ToolOutput(int), ToolOutput(str)]],
+ ids=["empty-list", "unsupported-operand"],
+ )
+ def test_returns_raw_when_list_cannot_be_unioned(self, output_type):
+ assert rehydrate_pydantic_output(output_type, "1",
serialize_output=False) == "1"
+
+ @mock.patch("airflow.providers.common.ai.utils.output_type.TypeAdapter",
autospec=True)
+ def test_propagates_type_error_from_type_adapter(self, mock_type_adapter):
+ mock_type_adapter.side_effect = TypeError("broken schema")
+
+ with pytest.raises(TypeError, match="broken schema"):
+ rehydrate_pydantic_output(int, "1", serialize_output=False)
+
+ def test_rejects_edit_matching_no_list_candidate(self):
+ with pytest.raises(ReviewedOutputValidationError):
+ rehydrate_pydantic_output([int, float], '"x"',
serialize_output=False)
+
@pytest.mark.parametrize(
"takes_an_order", [False, True], ids=["call-schema",
"definitions-wrapped-schema"]
)
@@ -241,6 +579,7 @@ class TestRehydratePydanticOutput:
@pytest.mark.parametrize(
"raw", ['{"user_name": "kax"}', '{"userName": "kax"}'],
ids=["field-name", "alias"]
)
+ @requires_walker
def test_accepts_a_field_name_or_an_alias(self, raw):
result = rehydrate_pydantic_output(ValidationAlias, raw,
serialize_output=False)
@@ -335,6 +674,7 @@ class TestRehydratePydanticOutput:
"list-of-computed-field-forbid-extra",
],
)
+ @requires_walker
def test_a_dumped_output_validates_back(self, output_type, output):
raw = dump_output_to_json(output)
@@ -418,6 +758,7 @@ class TestRehydratePydanticOutput:
assert rehydrate_pydantic_output(A | str, raw, serialize_output=False)
== raw
+@requires_walker
def
test_a_model_that_turned_off_validation_by_alias_reads_its_dump_back_unchanged():
output = AliasNamedLikeAnotherField(source="SRC", destination="DST")
diff --git
a/providers/common/ai/tests/unit/common/ai/utils/test_query_results.py
b/providers/common/ai/tests/unit/common/ai/utils/test_query_results.py
index eee5a06c746..3ab3a2428ca 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_query_results.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_query_results.py
@@ -24,6 +24,7 @@ import pytest
from airflow.providers.common.ai.utils.query_results import (
_SCHEMA_SAMPLE_SIZE,
+ build_get_schema_description,
build_query_result,
build_schema_result,
)
@@ -334,3 +335,29 @@ class TestSchemaResult:
assert data["sample_columns"] == []
assert "type_histogram" not in data
assert data["hint"]
+
+
+class TestSchemaToolName:
+ def test_description_names_the_given_tool(self):
+ assert "Call pg_get_schema again" in
build_get_schema_description(tool_name="pg_get_schema")
+
+ def test_description_defaults_to_the_bare_name(self):
+ assert "Call get_schema again" in build_get_schema_description()
+
+ def test_hints_name_the_given_tool(self):
+ no_match = json.loads(
+ build_schema_result(
+ _cols(3),
+ max_columns=100,
+ max_result_bytes=65_536,
+ name_contains="zzz",
+ tool_name="pg_get_schema",
+ )
+ )
+ summary = json.loads(
+ build_schema_result(
+ _cols(250), max_columns=100, max_result_bytes=65_536,
tool_name="pg_get_schema"
+ )
+ )
+ assert "Call pg_get_schema without name_contains" in no_match["hint"]
+ assert "Call pg_get_schema with a name_contains substring" in
summary["hint"]
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_tool_prefix.py
b/providers/common/ai/tests/unit/common/ai/utils/test_tool_prefix.py
new file mode 100644
index 00000000000..95b00500379
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_tool_prefix.py
@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import pytest
+
+from airflow.providers.common.ai.utils.tool_prefix import ToolPrefixMixin,
validate_tool_prefix
+
+
+class _Prefixed(ToolPrefixMixin):
+ def __init__(self, tool_prefix: str = "") -> None:
+ self._set_tool_prefix(tool_prefix)
+
+
+class TestValidateToolPrefix:
+ @pytest.mark.parametrize("tool_prefix", ["a-b", "9lives", "has space"])
+ def test_rejects_a_prefix_that_is_not_an_identifier(self, tool_prefix):
+ with pytest.raises(
+ ValueError, match=f"tool_prefix must be a valid Python identifier,
got {tool_prefix!r}"
+ ):
+ validate_tool_prefix(tool_prefix)
+
+ @pytest.mark.parametrize("tool_prefix", ["", "pg", "_pg", "pg2"])
+ def test_accepts_an_empty_or_identifier_prefix(self, tool_prefix):
+ validate_tool_prefix(tool_prefix)
+
+
+class TestToolPrefixMixin:
+ def test_set_tool_prefix_rejects_and_keeps_the_previous_prefix(self):
+ toolset = _Prefixed("pg")
+
+ with pytest.raises(ValueError, match="valid Python identifier"):
+ toolset._set_tool_prefix("a-b")
+
+ assert toolset._tool_prefix == "pg"
+
+ def test_the_prefix_defaults_to_empty(self):
+ assert ToolPrefixMixin._tool_prefix == ""
+
+ @pytest.mark.parametrize(("tool_prefix", "expected"), [("", "query"),
("pg", "pg_query")])
+ def test_build_tool_name(self, tool_prefix, expected):
+ assert _Prefixed(tool_prefix)._build_tool_name("query") == expected
+
+ @pytest.mark.parametrize(
+ ("tool_prefix", "name", "expected"),
+ [
+ ("pg", "pg_query", "query"),
+ ("pg", "query", None),
+ ("pg", "pgquery", None),
+ ("", "query", "query"),
+ ],
+ )
+ def
test_find_base_tool_name_is_strict_about_a_name_without_the_prefix(self,
tool_prefix, name, expected):
+ assert _Prefixed(tool_prefix)._find_base_tool_name(name) == expected
+
+ @pytest.mark.parametrize(
+ ("tool_prefix", "separator", "expected"),
+ [("", "-", ""), ("", "_", ""), ("pg", "-", "-pg"), ("pg", "_", "_pg")],
+ )
+ def test_get_tool_prefix_id_suffix(self, tool_prefix, separator, expected):
+ assert _Prefixed(tool_prefix)._get_tool_prefix_id_suffix(separator) ==
expected
+
+ def test_get_tool_prefix_id_suffix_defaults_to_a_dash(self):
+ assert _Prefixed("pg")._get_tool_prefix_id_suffix() == "-pg"
