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