diff --git a/providers/common/ai/docs/message_history.rst b/providers/common/ai/docs/message_history.rst index 453b4d7a7c37e..fa9ee78befee1 100644 --- a/providers/common/ai/docs/message_history.rst +++ b/providers/common/ai/docs/message_history.rst @@ -30,6 +30,13 @@ turns and, after the run, pushes the full updated transcript run reads it back to resume the conversation. ``None`` (the default) keeps the single-turn behavior unchanged. +To record the transcript of a single-turn run, set ``emit_message_history=True``. +The operator starts the run with no prior turns and pushes ``result.all_messages()`` +to the same ``message_history`` XCom key, so the transcript is available to +downstream tasks for inspection, evaluation or storage. ``emit_message_history`` defaults +to ``True`` when ``message_history`` is set; pass ``False`` to keep a seeded run from +pushing it. + The operator does **not** decide *where* a session is stored -- that keying is deployment-specific. The pattern is three tasks: load the prior transcript for the session, run the agent, store the updated transcript. The example keys a @@ -49,7 +56,8 @@ fresh session instead of trying to parse the string ``"None"``. .. note:: - ``message_history`` cannot be combined with ``enable_hitl_review`` -- the + ``message_history`` and ``emit_message_history`` cannot be combined with + ``enable_hitl_review`` -- the operator raises at construction. The post-review (human-approved) transcript is not recoverable today, so emitting the pre-review transcript would silently drop the reviewed turns. diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index fbb7b542989c4..9dbf792c9c9a9 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -451,6 +451,9 @@ Parameters When set, the post-run transcript is pushed to XCom under the key ``message_history`` for the next run to resume. Default ``None`` (single-turn). See :doc:`../message_history`. +- ``emit_message_history``: Push the post-run transcript to XCom under the key + ``message_history``. Defaults to ``True`` when ``message_history`` is set and ``False`` + otherwise; ``True`` records a single-turn run without seeding it. See :doc:`../message_history`. - ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via ``model_dump()`` before being pushed to XCom. Default ``False`` -- the diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst index 5feb431be94e9..3fad7b4ca4f92 100644 --- a/providers/common/ai/docs/stability.rst +++ b/providers/common/ai/docs/stability.rst @@ -69,6 +69,10 @@ Pydantic AI toolsets, but the Pydantic AI class they inherit from can change. - Seeds the run with the given conversation, as a list of messages or its JSON form, and publishes the finished conversation under the ``message_history`` XCom key so a later run can continue it. Cannot be combined with ``enable_hitl_review``. + * - ``emit_message_history`` on ``AgentOperator`` + - Publishes the finished conversation under the ``message_history`` XCom key without + seeding the run. Defaults to ``True`` when ``message_history`` is set, and ``False`` + stops a seeded run from publishing it. Cannot be combined with ``enable_hitl_review``. * - ``durable=True`` on ``AgentOperator`` and ``@task.agent``, and :class:`~airflow.providers.common.ai.durable.AirflowDurability` (:doc:`durable_execution`) diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst index ea3e8c3fae59e..e2b329e74dcfa 100644 --- a/providers/common/ai/docs/troubleshooting.rst +++ b/providers/common/ai/docs/troubleshooting.rst @@ -130,7 +130,7 @@ as a task failure. ``AgentOperator(durable=True)`` adds the capability to the agent, so a second one in ``capabilities`` would record every step twice. Drop one of them. -``message_history and enable_hitl_review=True cannot be used together`` +``message_history and enable_hitl_review=True cannot be used together`` and ``emit_message_history and enable_hitl_review=True cannot be used together`` The post-review transcript is not recoverable today, so the operator refuses rather than silently dropping the reviewed turns. See :doc:`message_history`. diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 06f052cb6c295..2392a1fdd19c0 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -406,6 +406,11 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa object-storage XCom backend or trim old turns. Not supported together with ``enable_hitl_review`` (raises) -- the post-review transcript is not yet recoverable. + :param emit_message_history: Push the post-run transcript to XCom under the key + ``message_history``. Defaults to ``True`` when ``message_history`` is set and ``False`` + otherwise. Set ``True`` to record a single-turn run without seeding it, or ``False`` + to keep a seeded run from pushing it. Not supported together with + ``enable_hitl_review`` (raises). **HITL Review parameters** (requires the ``hitl_review`` plugin): @@ -507,6 +512,7 @@ def __init__( durable: bool = False, cache_prompt: bool = True, message_history: list[ModelMessage] | str | bytes | None = None, + emit_message_history: bool | None = None, # Agent feedback parameters enable_hitl_review: bool = False, max_hitl_iterations: int = 5, @@ -537,6 +543,9 @@ def __init__( # No validation here -- see coerce_usage_limits() docstring for why. self.usage_limits = usage_limits self.message_history = message_history + self.emit_message_history = ( + message_history is not None if emit_message_history is None else emit_message_history + ) self.durable = durable self.cache_prompt = cache_prompt @@ -581,6 +590,9 @@ def __init__( # surface the final message history. raise ValueError("message_history and enable_hitl_review=True cannot be used together.") + if emit_message_history and enable_hitl_review: + raise ValueError("emit_message_history and enable_hitl_review=True cannot be used together.") + if durable or enable_hitl_review: self._reject_sandbox_without_continuity(durable=durable, enable_hitl_review=enable_hitl_review) @@ -1078,7 +1090,7 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag if self._usage_budget and (run_usage := self._run_usage) is not None: self._log_cumulative_usage(run_usage) - if self.message_history is not None: + if self.emit_message_history: self._emit_message_history(context, result) output = result.output diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 70d7faa9e4121..178db39816679 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -2127,22 +2127,32 @@ def test_none_is_single_turn_no_history_no_emit(self, mock_hook_cls, make_mock_r assert "message_history" not in pushed_keys assert pushed_keys == {"run_id", "usage", MODEL_NAME_XCOM_KEY} + @pytest.mark.parametrize( + ("kwargs", "emits"), + [ + pytest.param({"message_history": []}, True, id="seeded"), + pytest.param({"emit_message_history": True}, True, id="opt-in"), + pytest.param({"message_history": [], "emit_message_history": False}, False, id="opt-out"), + ], + ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_transcript_emitted_to_xcom_when_history_set(self, mock_hook_cls, make_mock_run_result): - """When message_history is set, the post-run transcript is pushed to XCom and round-trips.""" + def test_transcript_emitted_to_xcom(self, mock_hook_cls, kwargs, emits, make_mock_run_result): + """The transcript push follows emit_message_history, which defaults to whether the run is seeded.""" mock_agent = _make_mock_agent("ok", make_mock_run_result) mock_agent.run_sync.return_value.all_messages.return_value = _sample_history() mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c", message_history=[]) + op = AgentOperator(task_id="t", prompt="run", llm_conn_id="c", **kwargs) context = _make_context() op.execute(context=context) + # emit_message_history never seeds the run. + assert ("message_history" in mock_agent.run_sync.call_args.kwargs) is ("message_history" in kwargs) ti = context["task_instance"] pushes = {c.kwargs["key"]: c.kwargs["value"] for c in ti.xcom_push.call_args_list} - assert set(pushes) == {"run_id", "usage", "message_history", MODEL_NAME_XCOM_KEY} - restored = ModelMessagesTypeAdapter.validate_json(pushes["message_history"]) - assert len(restored) == 2 + assert ("message_history" in pushes) is emits + if emits: + assert len(ModelMessagesTypeAdapter.validate_json(pushes["message_history"])) == 2 @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_usage_limits_still_forwarded_with_history( @@ -2165,16 +2175,19 @@ def test_usage_limits_still_forwarded_with_history( @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) - def test_message_history_with_hitl_review_raises(self): - """message_history cannot be combined with HITL review (post-review transcript is lost).""" - with pytest.raises(ValueError, match="message_history and enable_hitl_review"): - AgentOperator( - task_id="t", - prompt="run", - llm_conn_id="c", - message_history=[], - enable_hitl_review=True, - ) + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + pytest.param({"message_history": []}, "message_history and enable_hitl_review", id="seeded"), + pytest.param( + {"emit_message_history": True}, "emit_message_history and enable_hitl_review", id="opt-in" + ), + ], + ) + def test_message_history_with_hitl_review_raises(self, kwargs, match): + """Transcript emission cannot be combined with HITL review (post-review transcript is lost).""" + with pytest.raises(ValueError, match=match): + AgentOperator(task_id="t", prompt="run", llm_conn_id="c", enable_hitl_review=True, **kwargs) def test_durable_path_also_seeds_message_history(self): """The durable branch forwards message_history into the run, and keeps its conversation id."""