kaxil commented on code in PR #74388: URL: https://github.com/apache/airflow/pull/74388#discussion_r4243883364
########## providers/common/ai/src/airflow/providers/common/ai/utils/prompt.py: ########## @@ -0,0 +1,53 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from pydantic_ai.messages import ( + AudioUrl, + BinaryContent, + DocumentUrl, + FileUrl, + ImageUrl, + TextContent, + VideoUrl, +) + + +def describe_prompt(prompt: str | Sequence[Any]) -> str: + if isinstance(prompt, str): + return prompt + return "\n".join(_describe_part(part) for part in prompt) + + +def _describe_part(part: Any) -> str: + if isinstance(part, FileUrl) and part.url[:5].lower() == "data:": + try: + part = BinaryContent.from_data_uri(part.url) + except ValueError: + return "[data URL]" + if isinstance(part, str): + return part + if isinstance(part, TextContent): + return part.content + if isinstance(part, BinaryContent): + return f"[{part.media_type}, {len(part.data)} bytes]" + if isinstance(part, (ImageUrl, AudioUrl, DocumentUrl, VideoUrl)): + return f"[{part.kind}: {part.url}]" + return f"[{type(part).__name__}]" Review Comment: `pydantic_ai.messages.UserContent` exists at the 2.33.0 floor, so the signature on line 33 (and the `prompt` field on the protocol in `approval.py`) could be `str | Sequence[UserContent]` rather than `Any`. Typed that way, a missing branch shows up: `UploadedFile` is a `UserContent` member and falls through to a bare `[UploadedFile]` here, while URL parts show their URL. Something like `[uploaded-file: {file_id}]` would tell the reviewer which file it was. ########## providers/common/ai/tests/unit/common/ai/operators/test_llm.py: ########## @@ -1104,25 +1090,22 @@ def test_execute_complete_restores_non_str_output_type(self, output_type, genera @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) -class TestLLMOperatorMultimodalPromptGuard: - """LLMOperator.execute raises before agent.run_sync when require_approval is True - and self.prompt is not a string -- covering direct-operator construction and the - native template rendering escape (where a string template renders to a Sequence).""" - +class TestLLMOperatorMultimodalPromptApproval: + @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) + @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_hook_cls): + def test_execute_with_approval_reviews_a_sequence_prompt( + self, mock_hook_cls, mock_upsert, mock_trigger_cls, make_mock_run_result + ): mock_agent = MagicMock(spec=["run_sync"]) + mock_agent.run_sync.return_value = make_mock_run_result("LLM response") mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - op = LLMOperator( - task_id="t", - prompt="placeholder", - llm_conn_id="c", - require_approval=True, - ) - op.prompt = ["x", object()] # simulate post-template-render value + op = LLMOperator(task_id="t", prompt="placeholder", llm_conn_id="c", require_approval=True) Review Comment: With `require_approval=True`, `review_reason` returns `"require_approval"`, so `body` stays `None` and this goes through the mixin's default body, which `test_approval.py` already covers. The `describe_prompt` call on the decision-policy path (`llm.py:344`) never sees a Sequence prompt, so reverting that line to `{self.prompt}` keeps the suite green. Could this drive a `decision_policy` with a below-threshold confidence instead, or parametrize over both paths? Passing the list to the constructor rather than reassigning `op.prompt` afterwards would also be closer to real use. ########## providers/common/ai/tests/unit/common/ai/utils/test_prompt.py: ########## @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest +from pydantic_ai.messages import BinaryContent, CachePoint, DocumentUrl, ImageUrl, TextContent + +from airflow.providers.common.ai.utils.prompt import describe_prompt + + [email protected]( + ("prompt", "expected"), + [ + ("plain text", "plain text"), + (["first", "second"], "first\nsecond"), + ([TextContent(content="typed text")], "typed text"), + ([BinaryContent(data=b"\x89PNG", media_type="image/png")], "[image/png, 4 bytes]"), + ([ImageUrl(url="https://example.com/a.png";)], "[image-url: https://example.com/a.png]";), + ([ImageUrl(url="data:image/png;base64,iVBORw0KGgo=")], "[image/png, 8 bytes]"), + ([DocumentUrl(url="data:text/plain,hello")], "[data URL]"), + ([ImageUrl(url="DATA:image/png;base64,not-base64!")], "[data URL]"), Review Comment: This passes because `from_data_uri` rejects the uppercase `DATA:` scheme before it decodes anything, so the invalid-base64 path isn't what this case exercises. Maybe split it into a lowercase invalid-payload case and an explicit uppercase one (a valid uppercase data URL also renders `[data URL]`), with ids so a failure says which? ########## providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py: ########## @@ -80,12 +81,10 @@ class DeferForApprovalProtocol(Protocol): on_approval_timeout: Literal["fail", "approve", "reject"] approval_notifiers: Sequence[BaseNotifier] approval_assigned_users: list[HITLUser] - prompt: str + prompt: str | Sequence[Any] Review Comment: The class docstring still lists ``prompt`` as ``str`` (line 133), so it should change along with this. ########## providers/common/ai/src/airflow/providers/common/ai/utils/prompt.py: ########## @@ -0,0 +1,53 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from pydantic_ai.messages import ( + AudioUrl, + BinaryContent, + DocumentUrl, + FileUrl, + ImageUrl, + TextContent, + VideoUrl, +) + + +def describe_prompt(prompt: str | Sequence[Any]) -> str: + if isinstance(prompt, str): + return prompt + return "\n".join(_describe_part(part) for part in prompt) + + +def _describe_part(part: Any) -> str: + if isinstance(part, FileUrl) and part.url[:5].lower() == "data:": + try: + part = BinaryContent.from_data_uri(part.url) + except ValueError: + return "[data URL]" + if isinstance(part, str): + return part + if isinstance(part, TextContent): + return part.content + if isinstance(part, BinaryContent): + return f"[{part.media_type}, {len(part.data)} bytes]" + if isinstance(part, (ImageUrl, AudioUrl, DocumentUrl, VideoUrl)): + return f"[{part.kind}: {part.url}]" Review Comment: This copies the full URL into the review body, and that body also goes to every `approval_notifiers` entry as `task.body`. A presigned S3/GCS URL, or one with `user:pass@` in it, would land in Slack or email with its signature intact until it expires. Could we keep scheme, host and path and drop userinfo, query and fragment (`urlsplit` / `urlunsplit`)? That still tells the reviewer which object it is. A test case with a signed URL would pin it. ########## providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py: ########## @@ -31,7 +31,6 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.validation import ( - reject_sequence_with_unsupported_feature, validate_prompt, Review Comment: Nit: with one name left this can collapse to `from airflow.providers.common.ai.utils.validation import validate_prompt`. Same in the other three decorators; the trailing comma stops ruff-format from doing it. ########## providers/common/ai/tests/unit/common/ai/decorators/test_llm.py: ########## @@ -102,24 +100,6 @@ def my_prompt(): assert op.prompt == prompt mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None, cancellation_token=ANY) - @pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="require_approval needs Airflow >= 3.1.0") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_sequence_prompt_with_require_approval_raises_before_run_sync(self, mock_hook_cls): Review Comment: With this and the three sibling `..._raises_before_run_sync` tests gone, no decorator test runs `require_approval=True` with a Sequence prompt, so a guard creeping back into any of the four decorators wouldn't fail anything. One test parametrized over the four decorators that asserts the approval pause and that `run_sync` ran would cover 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
![https://example.com/a.png"](https://example.com/a.png"){kind=link}
![https://example.com/a.png]"](https://example.com/a.png]"){kind=link}