You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "Lee-W (via GitHub)" <gi...@apache.org> on 2026/10/14 07:17:42 UTC
[PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Lee-W opened a new pull request, #74695:
URL: https://github.com/apache/airflow/pull/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, while HookToolset still carried the deprecated tool_name_prefix alias. AgentOperator also annotated output_type narrower than LLMOperator even though both accept the same values at runtime.
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4271166388
##########
providers/common/ai/docs/changelog.rst:
##########
@@ -44,6 +44,14 @@ Changelog
OpenSandbox server configuration, or pass ``SandboxSpec(block_network=False)`` for a sandbox
that may reach the network. See :ref:`the OpenSandbox backend <sandbox-backend-opensandbox>`.
+.. note::
+ ``HookToolset`` no longer accepts the deprecated ``tool_name_prefix`` parameter. Pass
+ ``tool_prefix`` instead: ``tool_name_prefix="s3_"`` becomes ``tool_prefix="s3"``, since
+ ``HookToolset`` adds the underscore itself. ``MCPToolset``'s ``tool_prefix`` now defaults to
+ ``""``, like every other toolset's and like ``MCPHook``'s.
Review Comment:
`MCPHook`'s default changed in this PR too (`str | None = None` to `str = ""`), and several toolsets take no `tool_prefix` at all. Maybe: "`MCPToolset`'s and `MCPHook`'s `tool_prefix` now default to `""` instead of `None`, like the other toolsets that take a prefix."
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -491,7 +494,7 @@ def __init__(
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type = str,
+ output_type: type | UnionType = str,
Review Comment:
`_agent_output_type` has a branch for a list `output_type` (`[str, int]`), and `test_a_list_output_type_is_extended_not_nested` passes one, but this annotation and the docstring only cover a single type or a union. Since this signature is the one being settled, worth adding the list form here and on `LLMOperator`, e.g. `type | UnionType | Sequence[type]`?
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -369,7 +379,11 @@ def _is_ref_allowed(self, catalog: str | None, schema: str | None, table: str) -
@property
def id(self) -> str:
- return f"sql-{self._db_conn_id}"
+ suffix = f"-{self._tool_prefix}" if self._tool_prefix else ""
+ return f"sql-{self._db_conn_id}{suffix}"
+
+ def _tool_name(self, base: str) -> str:
Review Comment:
The hints now follow `tool_prefix`, but the error users hit first steers them the other way: two `SQLToolset`s on one agent fail with pydantic-ai's default `tool_name_conflict_hint`, "wrap the toolset in a `PrefixedToolset`". `.prefixed("pg")` leaves `_tool_prefix` empty, so the model sees `pg_get_schema` while the retry text and the `get_schema` description still name `list_tables` and `get_schema`, which belong to the other database. Could `SQLToolset` and `DataFusionToolset` override `tool_name_conflict_hint` to say "pass `tool_prefix=...`" instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4398260636
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -140,11 +217,215 @@ 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_is_false_without_the_sdk_walker(self):
+ assert is_model_dump_required(A(x=1), [A, int]) 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_keeps_a_container_without_the_sdk_walker(self):
+ output = [A(x=1)]
+
+ assert dump_output_for_xcom(output, list[A], serialize_output=True) is output
+ assert dump_output_for_xcom(output, list[A], serialize_output=False) is 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
+ ):
+ from airflow.sdk.serde import serialize
Review Comment:
Updated it as `MagicMock(spec=RunContext)` and removed the local import.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4323848500
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py:
##########
@@ -152,31 +146,14 @@ def __init__(
*,
allowed_methods: list[str],
tool_prefix: str = "",
- tool_name_prefix: str | None = None,
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}.")
- if tool_name_prefix is not None:
- if tool_prefix and f"{tool_prefix}_" != tool_name_prefix:
- raise ValueError(
- f"tool_prefix={tool_prefix!r} and tool_name_prefix={tool_name_prefix!r} conflict. "
- f"Pass only tool_prefix; tool_name_prefix is its deprecated alias and must equal "
- f"{tool_prefix + '_'!r} when both are given."
- )
- warnings.warn(
- f"Parameter `tool_name_prefix` is deprecated. Use "
- f"`tool_prefix={tool_name_prefix.removesuffix('_')!r}` instead; "
- f"HookToolset adds the '_' itself.",
- category=AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
- name_prefix = tool_name_prefix
- else:
- name_prefix = f"{tool_prefix}_" if tool_prefix else ""
+ name_prefix = f"{tool_prefix}_" if tool_prefix else ""
Review Comment:
Restored `tool_name_prefix` with its deprecation warning, as in 1.0.0, and dropped the 1.0.0 changelog note. `SQLToolset` and `DataFusionToolset` `tool_prefix` and the conflict hints stay.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4394095191
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -140,11 +217,215 @@ 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_is_false_without_the_sdk_walker(self):
+ assert is_model_dump_required(A(x=1), [A, int]) 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_keeps_a_container_without_the_sdk_walker(self):
+ output = [A(x=1)]
+
+ assert dump_output_for_xcom(output, list[A], serialize_output=True) is output
+ assert dump_output_for_xcom(output, list[A], serialize_output=False) is 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
+ ):
+ from airflow.sdk.serde import serialize
Review Comment:
Small one: `serialize` is already imported in the `try` block at the top of the file, so this import can go.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,59 @@
_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`` there is no evidence that a class is unregistered, so nothing is dumped.
+ """
+ if iter_pydantic_models is None:
Review Comment:
After the rebase `AgentOperator` goes through main's `_dump_unregistered_models`, so the two operators now push different values for the same `output_type`. With `[M, int]`:
- a model with `serialize_by_alias=True` arrives as `fullName` from the agent but `full_name` from `LLMOperator`
- a `date` next to the model in a tuple stays a `date` from the agent, but turns into a string here
- before 3.3, the early return here pushes `list[M]` as instances while the agent dumps it (and `structured_output.rst` says the worker registers no class there)
`stability.rst` would lock in both. Could `LLMOperator` and `_convert_output_for_xcom` share one helper in this module? If not, the docs need to say where they differ.
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py:
##########
@@ -107,6 +107,86 @@ def test_tool_definitions_have_descriptions(self):
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"]
+
+
+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()))
Review Comment:
Small one: `MagicMock(spec=RunContext)` here, like the other new tests in this file?
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -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``.
Review Comment:
For `LLMOperator`, `@task.llm` and `LLMFileAnalysisOperator` on 3.3+ this is new: 1.0.0 pushed the instance there (`llm.py:390` at the tag dumps only with `serialize_output`), and the same holds for `ToolOutput(...)`, a `NewType` and a `TypeAliasType`. The upgrade bullets further down and the changelog note only name `AgentOperator`, and the pre-3.3 bullet doesn't hold for the LLM operators, which still push a `list[MyModel]` as instances there. Could the notes name the LLM operators too?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4273157144
##########
providers/common/ai/docs/changelog.rst:
##########
@@ -44,6 +44,14 @@ Changelog
OpenSandbox server configuration, or pass ``SandboxSpec(block_network=False)`` for a sandbox
that may reach the network. See :ref:`the OpenSandbox backend <sandbox-backend-opensandbox>`.
+.. note::
+ ``HookToolset`` no longer accepts the deprecated ``tool_name_prefix`` parameter. Pass
+ ``tool_prefix`` instead: ``tool_name_prefix="s3_"`` becomes ``tool_prefix="s3"``, since
+ ``HookToolset`` adds the underscore itself. ``MCPToolset``'s ``tool_prefix`` now defaults to
+ ``""``, like every other toolset's and like ``MCPHook``'s.
Review Comment:
Reworded as suggested. Thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4340279710
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,56 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to XCom dumped.
+
+ The instances are looked up in the output itself, including the elements of a list, tuple, set or
+ frozenset and the 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`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not (found := list(_iter_models_in_output(output))):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output: bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when ``is_model_dump_required``, else unchanged.
+
+ The dump is JSON-compatible, as the serde ``BaseModel`` serializer's is, so a ``HttpUrl``, ``Path`` or
+ ``Enum`` field is a string that XCom accepts. Field names, not aliases, match ``dump_output_to_json``.
+ """
+ if is_model_dump_required(output, output_type, serialize_output=serialize_output):
+ return to_jsonable_python(output, by_alias=False)
Review Comment:
I asked for `to_jsonable_python` with the newly dumped models in mind, but this also runs on the paths 1.0.0 dumped with python-mode `model_dump()`: `serialize_output=True`, and every model output before Airflow 3.3 (`serialize_output or not _CORE_WALKER`). There a `date`, `datetime`, `Decimal` or `UUID` field now reaches the downstream task as a string instead of the object serde used to round-trip, so code like `inv["due"] < today` starts raising after a minor upgrade of a stable operator. Could the JSON-mode dump apply only to models dumped because the worker can't register their class, keep `model_dump()` for the other two, and add a test that a `date` field survives `serialize_output=True`?
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -64,4 +67,7 @@ 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.
+operator dumps the model with ``pydantic_core.to_jsonable_python`` before pushing to XCom,
+so ``HttpUrl``, ``Path`` and ``Enum`` fields become plain JSON values. The operator also dumps an
+output that is a list, tuple, set or dict holding models, and any model of a shape the worker
+registers no class for (such as a list of types), even without ``serialize_output``.
Review Comment:
This reads as if a list or dict of models is always dumped, but a registered `output_type=list[MyModel]` still pushes the instances unless `serialize_output=True`. Only a model whose class the worker can't register is dumped without it. Could the sentence split those two cases?
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -16,28 +16,53 @@
# 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, datetime
from enum import Enum
+from pathlib import Path
from typing import Annotated, Literal, NamedTuple
+from unittest import mock
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,
+ dump_output_for_xcom,
dump_output_to_json,
+ is_model_dump_required,
rehydrate_pydantic_output,
)
+from airflow.sdk.serde import deserialize, serialize
Review Comment:
This import fails collection in the 2.11.1, 3.0.6 and 3.1.8 compat jobs (`No module named 'airflow.sdk.serde'`, "Interrupted: 2 errors during collection"), so none of the provider's tests run there. `test_llm.py:61` has the same one. Moving both into the `try` block just below works, since the tests that use them are already gated by `requires_walker` / `requires_typed_xcom`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4374366988
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1106,21 +1114,13 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
# So the approved answer is the last answer the agent returned, and there is no need to parse
# it back from that text.
output = self._output_under_review
- if isinstance(output, BaseModel) and (
- self._serialize_model_output or type(output) not in iter_pydantic_models(self.output_type)
- ):
- # 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, so the operator pushes the model as a dict. When
- # iter_pydantic_models could not be imported, _serialize_model_output is always True, so the
- # condition above never reaches the iter_pydantic_models call.
- output = output.model_dump()
+ output = dump_output_for_xcom(
Review Comment:
After the rebase, it goes through `_convert_output_for_xcom`, which dumps an unregistered model with `model_dump(mode="json")`. I removed the comment in `output_type.py` that claimed 1.0.0 pushed Python mode everywhere.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,70 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to XCom dumped.
+
+ The instances are looked up in the output itself, including the elements of a list, tuple, set or
+ frozenset and the 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`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not (found := [model for model, _ in _iter_models_in_output(output)]):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output: bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when ``is_model_dump_required``, else unchanged.
+
+ With ``serialize_output`` the dump holds Python values, as the ``model_dump()`` of 1.0.0 does, so a
+ ``date`` or ``Decimal`` field keeps its type; an output with a set, frozenset or dict key holding a model
+ is dumped to JSON values instead. A model the worker registers no class for, without ``serialize_output``,
+ is dumped to JSON values as well, field names rather than aliases, so a ``HttpUrl``, ``Path`` or ``Enum``
+ field is a string that XCom accepts.
+ """
+ if not is_model_dump_required(output, output_type, serialize_output=serialize_output):
+ return output
+ # Python mode is what 1.0.0 pushed, a minor release must not change it. A model dumped to a dict cannot
+ # be a set member, so python mode raises on a set holding models.
+ if serialize_output and not any(in_set for _, in_set in _iter_models_in_output(output)):
+ if isinstance(output, BaseModel):
+ return output.model_dump()
+ return TypeAdapter(Any).dump_python(output, mode="python")
+ return to_jsonable_python(output, by_alias=False)
+
+
+def _iter_models_in_output(output: Any) -> Iterator[tuple[BaseModel, bool]]:
+ """
+ Yield the ``BaseModel`` instances in an output, each with whether it must be hashable.
+
+ Looks inside list, tuple, set, frozenset and dict. A set member or a dict key is hashable, so it is
+ flagged, and so is whatever is nested in one. Dataclass and NamedTuple contents and model fields
Review Comment:
Fixed. The docstring now says a `NamedTuple` is walked as a tuple. Only dataclass contents and model fields are skipped, and a test covers the `NamedTuple` case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4273153083
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -491,7 +494,7 @@ def __init__(
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type = str,
+ output_type: type | UnionType = str,
Review Comment:
Added `Sequence[type]` to `output_type` on `AgentOperator` and `LLMOperator`, with a docstring.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4269136306
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py:
##########
@@ -84,12 +84,12 @@ def __init__(
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:
self._mcp_conn_id = mcp_conn_id
- self._tool_prefix = tool_prefix
+ self._tool_prefix = tool_prefix or ""
Review Comment:
Yes, we should drop the `None` promise. `tool_prefix` is `str` with a `""` default on both `MCPToolset` and `MCPHook`. removed `or ""` fallbacks, and the docstrings, `mcp.rst` and the changelog mention of `None`.
Also, removed he `type: ignore` tests introduced in another PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4311586204
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -220,7 +221,7 @@ def __init__(
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type | UnionType = str,
+ output_type: type | UnionType | Sequence[type] = str,
Review Comment:
`LLMOperator`, `LLMFileAnalysisOperator`, `AgentOperator` and `rehydrate_pydantic_output` now reuse `is_model_dump_required`, so a model returned under `output_type=[MyModel, int]` is pushed as a `dict` on those return paths, including the `LLMOperator` review path. A single type or a union still returns an instance.
`LLMFileAnalysisOperator` overrides `execute`, so it had its own copy of the check and needed the same change. Added `[Model, int]` tests for each path, and updated the `output_type` docstrings and docs that said the model is pushed unchanged.
Behavior change: `rehydrate_pydantic_output([A, int], ...)` with `serialize_output=False` now returns a `dict` where it used to return the model.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4269120740
##########
providers/common/ai/docs/toolsets/sql.rst:
##########
@@ -329,6 +329,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. ``"pg"`` produces ``"pg_query"``).
Review Comment:
Changed it to say the function sees the prefixed name and compares against `f"{prefix}_query"`, and warns against `endswith("query")` since it also matches `check_query`. The `tool_prefix` entry points back to it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4314032144
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py:
##########
@@ -152,31 +146,14 @@ def __init__(
*,
allowed_methods: list[str],
tool_prefix: str = "",
- tool_name_prefix: str | None = None,
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}.")
- if tool_name_prefix is not None:
- if tool_prefix and f"{tool_prefix}_" != tool_name_prefix:
- raise ValueError(
- f"tool_prefix={tool_prefix!r} and tool_name_prefix={tool_name_prefix!r} conflict. "
- f"Pass only tool_prefix; tool_name_prefix is its deprecated alias and must equal "
- f"{tool_prefix + '_'!r} when both are given."
- )
- warnings.warn(
- f"Parameter `tool_name_prefix` is deprecated. Use "
- f"`tool_prefix={tool_name_prefix.removesuffix('_')!r}` instead; "
- f"HookToolset adds the '_' itself.",
- category=AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
- name_prefix = tool_name_prefix
- else:
- name_prefix = f"{tool_prefix}_" if tool_prefix else ""
+ name_prefix = f"{tool_prefix}_" if tool_prefix else ""
Review Comment:
common.ai 1.0.0 went to PyPI today with `tool_name_prefix` still accepted (deprecated, with its warning), and its `stability.rst` lists `HookToolset` as stable: it "keeps its public parameters, their defaults and the behaviour described below until the next major release". Dropping the alias now turns `HookToolset(..., tool_name_prefix="s3_")` into a `TypeError` at parse on a minor upgrade, and the changelog note for it lands in the already-released 1.0.0 section. Could the alias and its warning stay until 2.0, with the note moved out of 1.0.0? The new `tool_prefix` on `SQLToolset` and `DataFusionToolset` and the conflict hints are fine for 1.1.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +45,24 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether a ``BaseModel`` output must be pushed to XCom as 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 those shapes are dumped.
+ Without ``iter_pydantic_models`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not isinstance(output, BaseModel):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ return type(output) not in iter_pydantic_models(output_type)
Review Comment:
The callers dump this with a python-mode `model_dump()`, so an `HttpUrl`, `Path` or plain `Enum` field stays an object serde can't encode, and `output_type=[Page, str]` with `url: HttpUrl` now fails at the XCom push. 1.0.0 pushed the instance there and serde's pydantic serializer handled it with `mode="json"`. Separately, `[list[Item], str]` never gets here since the output isn't a `BaseModel`, so those models still go out unregistered. Could this dump with `pydantic_core.to_jsonable_python(output)` in one shared helper, with a test that feeds the result to `serde.serialize`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4374365349
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,70 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to XCom dumped.
+
+ The instances are looked up in the output itself, including the elements of a list, tuple, set or
+ frozenset and the 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`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not (found := [model for model, _ in _iter_models_in_output(output)]):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output: bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when ``is_model_dump_required``, else unchanged.
+
+ With ``serialize_output`` the dump holds Python values, as the ``model_dump()`` of 1.0.0 does, so a
+ ``date`` or ``Decimal`` field keeps its type; an output with a set, frozenset or dict key holding a model
+ is dumped to JSON values instead. A model the worker registers no class for, without ``serialize_output``,
+ is dumped to JSON values as well, field names rather than aliases, so a ``HttpUrl``, ``Path`` or ``Enum``
+ field is a string that XCom accepts.
+ """
+ if not is_model_dump_required(output, output_type, serialize_output=serialize_output):
+ return output
+ # Python mode is what 1.0.0 pushed, a minor release must not change it. A model dumped to a dict cannot
+ # be a set member, so python mode raises on a set holding models.
+ if serialize_output and not any(in_set for _, in_set in _iter_models_in_output(output)):
+ if isinstance(output, BaseModel):
+ return output.model_dump()
+ return TypeAdapter(Any).dump_python(output, mode="python")
Review Comment:
`dump_output_for_xcom` now calls `model_dump()` only when the output itself is a model and `serialize_output=True`, as in 1.0.0. A `list[Link]` or `dict[str, Link]` of registered models is left to serde, and tests push one through `serialize`. Only a model the worker cannot register is dumped, to JSON values.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4273149206
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -369,7 +379,11 @@ def _is_ref_allowed(self, catalog: str | None, schema: str | None, table: str) -
@property
def id(self) -> str:
- return f"sql-{self._db_conn_id}"
+ suffix = f"-{self._tool_prefix}" if self._tool_prefix else ""
+ return f"sql-{self._db_conn_id}{suffix}"
+
+ def _tool_name(self, base: str) -> str:
Review Comment:
`SQLToolset` and `DataFusionToolset` now override `tool_name_conflict_hint` to tell users to pass `tool_prefix=...`, and to say that wrapping the toolset in `PrefixedToolset` renames only the tools. `HookToolset` and `MCPToolset` are unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4374368841
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -214,12 +221,14 @@ 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``, a Pydantic ``BaseModel`` output, or a
+ list, tuple, set or dict holding models, is dumped to JSON-compatible
+ values with ``pydantic_core.to_jsonable_python`` before being pushed to
Review Comment:
Updated the `serialize_output` docstrings of `LLMOperator` and `AgentOperator`, `structured_output.rst`, the operator pages and `stability.rst`. All of them now describe the same three cases. `stability.rst` no longer says Agent follows the same rules as `@task.llm`, since Agent does not walk sets or dict keys. The `LLMOperator` row now mentions that `[MyModel, int]` is pushed as a dict.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4402370481
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,59 @@
_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`` there is no evidence that a class is unregistered, so nothing is dumped.
+ """
+ if iter_pydantic_models is None:
Review Comment:
`AgentOperator` and `@task.agent` now push their output through the same `dump_output_for_xcom` helper as the LLM operators, and its private `_dump_unregistered_models` is gone, so one `output_type` gives one XCom value (field names rather than aliases, `date` next to a model in a tuple arrives as a string, a tuple arrives as a list). Before task-sdk 3.3 both still dump only a top-level model and push `list[MyModel]` as instances, as 1.0.0 does
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4269151188
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -429,10 +444,11 @@ async def execute_tool(
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- if name not in ("list_tables", "get_schema", "query", "check_query"):
+ base = name.removeprefix(f"{self._tool_prefix}_") if self._tool_prefix else 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)
Review Comment:
These are now built from the the toolset's own tool names.
* retry text
* `allowed_tables` refusals
* DataFusion retry
* `get_schema` description and hints
`build_schema_result` takes the `get_schema` name as a keyword argument, and the unprefixed text is unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4348587597
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,56 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to XCom dumped.
+
+ The instances are looked up in the output itself, including the elements of a list, tuple, set or
+ frozenset and the 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`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not (found := list(_iter_models_in_output(output))):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output: bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when ``is_model_dump_required``, else unchanged.
+
+ The dump is JSON-compatible, as the serde ``BaseModel`` serializer's is, so a ``HttpUrl``, ``Path`` or
+ ``Enum`` field is a string that XCom accepts. Field names, not aliases, match ``dump_output_to_json``.
+ """
+ if is_model_dump_required(output, output_type, serialize_output=serialize_output):
+ return to_jsonable_python(output, by_alias=False)
Review Comment:
`dump_output_for_xcom` now dumps with `model_dump()` when `serialize_output=True` (which includes task-sdk versions without the walker), as 1.0.0 did, so `date`, `Decimal` and `UUID` fields keep their types. `to_jsonable_python` is used only for a model whose class the worker cannot register.
One exception: an output with a set, frozenset or dict key holding a model is dumped to JSON values, because pydantic cannot hold a dumped model in a set or as a key.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4348576951
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -16,28 +16,53 @@
# 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, datetime
from enum import Enum
+from pathlib import Path
from typing import Annotated, Literal, NamedTuple
+from unittest import mock
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,
+ dump_output_for_xcom,
dump_output_to_json,
+ is_model_dump_required,
rehydrate_pydantic_output,
)
+from airflow.sdk.serde import deserialize, serialize
Review Comment:
Moved `deserialize` and `serialize` into the `try` block. The tests that use them are already gated by `requires_walker` / `requires_typed_xcom`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4398318249
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -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``.
Review Comment:
`structured_output.rst`, `stability.rst` and the changelog note now mentions `LLMOperator`, `@task.llm` and `LLMFileAnalysisOperator` next to `AgentOperator`, and the pre-3.3 bullet mentions the LLM operators still push `list[MyModel]` as instances.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4289924989
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -220,7 +221,7 @@ def __init__(
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type | UnionType = str,
+ output_type: type | UnionType | Sequence[type] = str,
Review Comment:
With a list now part of the signature, `output_type=[Summary, int]` breaks whichever task reads this output. The worker allow-lists only the classes `iter_pydantic_models(op.output_type)` finds, and it finds none in a list, so with the default `serialize_output=False` the pushed `Summary` fails to deserialize downstream ("not found in allow list"), while `Summary | int` works. That hits both operators' normal return paths, and LLMOperator's review path now too, since `rehydrate_pydantic_output` used to hand back the raw string for a list and now returns the model.
`AgentOperator`'s review branch already guards for exactly this (`type(output) not in iter_pydantic_models(self.output_type)`, then `model_dump()`). Applying the same check on the other return paths of both operators, plus a `[Model, int]` test, would close it.
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py:
##########
@@ -148,6 +148,108 @@ def test_tools_declare_string_return_schema(self):
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 _call_prefixed(ts: SQLToolset, name: str, args: dict):
+ return asyncio.run(ts.call_tool(name, args, ctx=MagicMock(), tool=MagicMock()))
Review Comment:
Small one: `_call_prefixed` passes bare `MagicMock()`s, while `_query` above and the dispatch test just before this use `spec=RunContext` and `spec=ToolsetTool`. Could it use the same specs, or could `_query` take the tool name instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4311086203
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py:
##########
@@ -148,6 +148,108 @@ def test_tools_declare_string_return_schema(self):
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 _call_prefixed(ts: SQLToolset, name: str, args: dict):
+ return asyncio.run(ts.call_tool(name, args, ctx=MagicMock(), tool=MagicMock()))
Review Comment:
`_call_prefixed` now passes `MagicMock(spec=RunContext)` and `MagicMock(spec=ToolsetTool)`, the same specs `_query`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4323872505
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +45,24 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether a ``BaseModel`` output must be pushed to XCom as 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 those shapes are dumped.
+ Without ``iter_pydantic_models`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not isinstance(output, BaseModel):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ return type(output) not in iter_pydantic_models(output_type)
Review Comment:
Switched to a shared `dump_output_for_xcom` that dumps with `pydantic_core.to_jsonable_python`, so `HttpUrl`, `Path` and `Enum` fields serialize. A test feeds the dumped model to `serde.serialize`.
Models inside a list, tuple, set or dict output are dumped too when the worker registers no class for them, so `[list[Item], str]` is covered. Field names are dumped, not aliases, like `dump_output_to_json`. Dict keys, dataclass and NamedTuple contents are not walked.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4268554807
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py:
##########
@@ -84,12 +84,12 @@ def __init__(
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:
self._mcp_conn_id = mcp_conn_id
- self._tool_prefix = tool_prefix
+ self._tool_prefix = tool_prefix or ""
Review Comment:
`tool_prefix` is annotated `str` now, but the docstring, `mcp.rst` and the changelog still promise `None` is accepted, and the new tests need `# type: ignore[arg-type]` to exercise that. Since the point here is to settle this signature, should the `None` promise go along with the `or ""` here and in `MCPHook`? `hooks/mcp.py:254` only checks truthiness, so a stray `None` still works without being part of the contract.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -429,10 +444,11 @@ async def execute_tool(
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- if name not in ("list_tables", "get_schema", "query", "check_query"):
+ base = name.removeprefix(f"{self._tool_prefix}_") if self._tool_prefix else 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)
Review Comment:
The retry text just below still says "Use the list_tables and get_schema tools", and the same bare names appear in the `allowed_tables` refusals (lines 612 and 622), the DataFusion retry at `datafusion.py:296`, and the shared `get_schema` description and hints in `utils/query_results.py` (96, 261, 312). With `tool_prefix="pg"` no tool has those names, and in the setup the docstring recommends a prefix for (a second `SQLToolset`), a bare `list_tables` can belong to the other database. Building these strings from `self._tool_name(...)`, and passing the name into the `query_results` helpers, would keep the hints pointing at this toolset's own tools.
##########
providers/common/ai/docs/toolsets/sql.rst:
##########
@@ -329,6 +329,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. ``"pg"`` produces ``"pg_query"``).
Review Comment:
With a prefix set, the `allow_writes` approval recipe a few lines up stops gating writes. `.approval_required()` passes the user's function the toolset's own `ToolDefinition`, so with `tool_prefix="pg"` it sees `pg_query`, `tool_def.name == "query"` is False, and the write runs without asking anyone. Could this entry (or the recipe itself) say the function sees the prefixed name, like `mcp.rst` already does for its tool filter?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4348583200
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -64,4 +67,7 @@ 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.
+operator dumps the model with ``pydantic_core.to_jsonable_python`` before pushing to XCom,
+so ``HttpUrl``, ``Path`` and ``Enum`` fields become plain JSON values. The operator also dumps an
+output that is a list, tuple, set or dict holding models, and any model of a shape the worker
+registers no class for (such as a list of types), even without ``serialize_output``.
Review Comment:
Split into two cases in `structured_output.rst`. With `serialize_output=True` every output holding a model is dumped. Without it, a model whose class the worker registers (`MyModel`, `list[MyModel]`) is pushed as an instance, and only a model of a shape the worker registers no class for (such as `[MyModel, int]`) is dumped to JSON values.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4353422934
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,70 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to XCom dumped.
+
+ The instances are looked up in the output itself, including the elements of a list, tuple, set or
+ frozenset and the 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`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not (found := [model for model, _ in _iter_models_in_output(output)]):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output: bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when ``is_model_dump_required``, else unchanged.
+
+ With ``serialize_output`` the dump holds Python values, as the ``model_dump()`` of 1.0.0 does, so a
+ ``date`` or ``Decimal`` field keeps its type; an output with a set, frozenset or dict key holding a model
+ is dumped to JSON values instead. A model the worker registers no class for, without ``serialize_output``,
+ is dumped to JSON values as well, field names rather than aliases, so a ``HttpUrl``, ``Path`` or ``Enum``
+ field is a string that XCom accepts.
+ """
+ if not is_model_dump_required(output, output_type, serialize_output=serialize_output):
+ return output
+ # Python mode is what 1.0.0 pushed, a minor release must not change it. A model dumped to a dict cannot
+ # be a set member, so python mode raises on a set holding models.
+ if serialize_output and not any(in_set for _, in_set in _iter_models_in_output(output)):
+ if isinstance(output, BaseModel):
+ return output.model_dump()
+ return TypeAdapter(Any).dump_python(output, mode="python")
Review Comment:
In 1.0.0, `model_dump()` ran only when the output itself was a `BaseModel` (`llm.py:390` at the tag), so a `list[Link]` or `dict[str, Link]` went to XCom as instances and serde's pydantic serializer handled each one. Here the whole container is dumped in python mode, which leaves the `HttpUrl` inside a plain dict, and the push fails with `Cannot serialize object of type HttpUrl`. I checked this in breeze with `serialize_output=True`; it also hits Airflow before 3.3 without the flag, since `_serialize_model_output` is always on there. Could this branch keep 1.0.0's top-level-only `model_dump()`, with a `serde.serialize` test for `list[Link]`?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -214,12 +221,14 @@ 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``, a Pydantic ``BaseModel`` output, or a
+ list, tuple, set or dict holding models, is dumped to JSON-compatible
+ values with ``pydantic_core.to_jsonable_python`` before being pushed to
Review Comment:
This docstring and `AgentOperator`'s still say the dump uses `to_jsonable_python`, while `agent.rst` and `structured_output.rst` now say `model_dump()`, so the API reference and the guide disagree. The `LLMOperator` row in `stability.rst` also doesn't mention that a `[MyModel, int]` output is pushed as a dict without `serialize_output`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1106,21 +1114,13 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
# So the approved answer is the last answer the agent returned, and there is no need to parse
# it back from that text.
output = self._output_under_review
- if isinstance(output, BaseModel) and (
- self._serialize_model_output or type(output) not in iter_pydantic_models(self.output_type)
- ):
- # 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, so the operator pushes the model as a dict. When
- # iter_pydantic_models could not be imported, _serialize_model_output is always True, so the
- # condition above never reaches the iter_pydantic_models call.
- output = output.model_dump()
+ output = dump_output_for_xcom(
Review Comment:
1.0.0 dumped an unregistered model on this review path with `model_dump()` (`agent.py:1097-1106` at the tag), so with `output_type=[Invoice, str]` a `date` or `Decimal` field now reaches the downstream task as a string after an approval. JSON mode also fixes `HttpUrl` here, so I'd be fine either way, but the comment at `output_type.py:82` says python mode is what 1.0.0 pushed, which isn't the case for this path.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,70 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to XCom dumped.
+
+ The instances are looked up in the output itself, including the elements of a list, tuple, set or
+ frozenset and the 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`` there is no evidence either way, so the caller's ``serialize_output`` decides.
+ """
+ if not (found := [model for model, _ in _iter_models_in_output(output)]):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output: bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when ``is_model_dump_required``, else unchanged.
+
+ With ``serialize_output`` the dump holds Python values, as the ``model_dump()`` of 1.0.0 does, so a
+ ``date`` or ``Decimal`` field keeps its type; an output with a set, frozenset or dict key holding a model
+ is dumped to JSON values instead. A model the worker registers no class for, without ``serialize_output``,
+ is dumped to JSON values as well, field names rather than aliases, so a ``HttpUrl``, ``Path`` or ``Enum``
+ field is a string that XCom accepts.
+ """
+ if not is_model_dump_required(output, output_type, serialize_output=serialize_output):
+ return output
+ # Python mode is what 1.0.0 pushed, a minor release must not change it. A model dumped to a dict cannot
+ # be a set member, so python mode raises on a set holding models.
+ if serialize_output and not any(in_set for _, in_set in _iter_models_in_output(output)):
+ if isinstance(output, BaseModel):
+ return output.model_dump()
+ return TypeAdapter(Any).dump_python(output, mode="python")
+ return to_jsonable_python(output, by_alias=False)
+
+
+def _iter_models_in_output(output: Any) -> Iterator[tuple[BaseModel, bool]]:
+ """
+ Yield the ``BaseModel`` instances in an output, each with whether it must be hashable.
+
+ Looks inside list, tuple, set, frozenset and dict. A set member or a dict key is hashable, so it is
+ flagged, and so is whatever is nested in one. Dataclass and NamedTuple contents and model fields
Review Comment:
Small one: a NamedTuple is a `tuple`, so line 106 does walk its contents. Only dataclass contents and model fields are skipped.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make common.ai tool_prefix and output_type signatures consistent [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4398254739
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py:
##########
@@ -107,6 +107,86 @@ def test_tool_definitions_have_descriptions(self):
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"]
+
+
+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()))
Review Comment:
Missed this one. Made it `MagicMock(spec=RunContext)`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org