guan404ming commented on code in PR #74389:
URL: https://github.com/apache/airflow/pull/74389#discussion_r4482568293
##########
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 sequence-prompt test now hits the cap with `max_hitl_iterations = 2` and
asserts `HITLMaxIterationsError`.
##########
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:
Made `prompt` a read-only property on the protocol and dropped the ignore in
`AgentOperator`.
##########
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:
Now asserts `run_sync` got the sequence prompt and dropped the unused return
value.
##########
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:
Widened the fake's `prompt` and dropped the ignore.
--
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}