kaxil commented on code in PR #74389:
URL: https://github.com/apache/airflow/pull/74389#discussion_r4467374142
##########
providers/common/ai/tests/unit/common/ai/decorators/test_agent.py:
##########
@@ -114,27 +114,30 @@ def my_prompt():
prompt, usage_limits=None, run_id="ti-1", cancellation_token=ANY,
usage=ANY
)
+
@patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review",
autospec=True)
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
- def test_sequence_prompt_with_hitl_review_raises_before_run_sync(self,
mock_hook_cls):
- """Sequence prompt + enable_hitl_review=True fails before the agent
runs."""
+ def test_sequence_prompt_with_hitl_review_reaches_the_review(
+ self, mock_hook_cls, mock_run_hitl, make_mock_run_result
+ ):
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
if not AIRFLOW_V_3_1_PLUS:
pytest.skip("enable_hitl_review requires Airflow >= 3.1.0")
mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value = make_mock_run_result("Initial
output")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ mock_run_hitl.return_value = "Initial output"
op = _AgentDecoratedOperator(
task_id="test",
python_callable=lambda: ["x",
ImageUrl(url="https://example.com/x.png";)],
llm_conn_id="my_llm",
enable_hitl_review=True,
)
- with pytest.raises(TypeError, match="enable_hitl_review=True"):
- op.execute(context={})
- mock_agent.run_sync.assert_not_called()
+ assert op.execute(context=_make_context()) == "Initial output"
+ mock_run_hitl.assert_called_once()
Review Comment:
This only checks that the review ran. Could it assert that `run_sync` got
the sequence prompt, like the operator test does? The
`mock_run_hitl.return_value` above can go too, since `execute` returns
`_output_under_review` and never reads what `run_hitl_review` returns.
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -238,7 +241,7 @@ def _poll_loop(
status=SessionStatus.MAX_ITERATIONS_EXCEEDED,
iteration=session.iteration,
max_iterations=self.max_hitl_iterations,
- prompt=self.prompt,
+ prompt=session.prompt,
Review Comment:
The new sequence-prompt test approves on the second round with
`max_hitl_iterations=5`, so it never reaches this push, and
`test_max_iterations_reached_raises` only uses a str prompt. Could the new test
also run into the cap (`max_hitl_iterations = 2` plus a second
`changes_requested`)? If this went back to `self.prompt`, a multimodal prompt
would raise a pydantic `ValidationError` here instead of
`HITLMaxIterationsError`, and no test would notice.
##########
providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py:
##########
@@ -267,6 +269,36 @@ def
test_changes_requested_calls_regenerate_and_pushes_new_output(
assert "airflow_hitl_review_agent_output_2" in output_pushes
assert output_pushes["airflow_hitl_review_agent_output_2"] ==
"Revised: Add more detail"
+ @patch.object(FakeAgenticOperator, "regenerate_with_feedback")
+ @patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep",
autospec=True)
+ def test_session_renders_a_sequence_prompt(
+ self, mock_sleep, mock_regenerate, mock_ti, context,
mock_supervisor_comms
+ ):
+ mock_regenerate.return_value = ("Revised", None)
+ mock_supervisor_comms.send.side_effect = [
+ XComResult(
+ key=XCOM_HUMAN_ACTION,
+ value=HumanActionData(
+ action="changes_requested", feedback="More detail",
iteration=1
+ ).model_dump(mode="json"),
+ ),
+ XComResult(
+ key=XCOM_HUMAN_ACTION,
+ value=HumanActionData(action="approve",
iteration=2).model_dump(mode="json"),
+ ),
+ ]
+ prompt = ["Describe this:", BinaryContent(data=b"\x89PNG",
media_type="image/png")]
+ op = FakeAgenticOperator(prompt=prompt) # type: ignore[arg-type]
Review Comment:
Nit: widening `FakeAgenticOperator`'s `prompt: str` to `str |
Sequence[UserContent]` would drop this ignore and keep the fake in line with
the protocol.
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -48,7 +51,7 @@ class HITLReviewProtocol(Protocol):
max_hitl_iterations: int
hitl_timeout: timedelta | None
hitl_poll_interval: float
- prompt: str
+ prompt: str | Sequence[Any]
Review Comment:
Would a read-only property work here instead of the `# type: ignore[misc]`
this needed in `AgentOperator._complete_run`? A plain protocol attribute is
invariant, so `AgentOperator`'s `prompt: str` stops matching. `@property def
prompt(self) -> str | Sequence[Any]: ...` accepts `str` hosts, and the ignore
then stops hiding drift in the other protocol members.
--
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/x.png"](https://example.com/x.png"){kind=link}