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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
