This is an automated email from the ASF dual-hosted git repository.
pankajastro pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 31bb7edc69c Add emit_message_history to record an agent transcript
without seeding it (#74811)
31bb7edc69c is described below
commit 31bb7edc69c410c9f58918f460a0fdf4a02f9c96
Author: Pankaj Singh <[email protected]>
AuthorDate: Tue Oct 20 17:11:07 2026 +0530
Add emit_message_history to record an agent transcript without seeding it
(#74811)
* Add emit_message_history to record an agent transcript without seeding it
Single-turn agents that want their transcript on XCom had to pass
message_history=[], a parameter for seeding a conversation, not for
observing a run. A dedicated flag states that intent directly, and
do_xcom_push=False now also suppresses the transcript push, as it does
the other XCom pushes.
Co-Authored-By: Claude <[email protected]>
* Give emit_message_history its own switch independent of do_xcom_push
Gating the transcript push on do_xcom_push also stopped it for
message_history users, narrowing a stable feature. A tri-state
emit_message_history defaults to whether the run is seeded, so existing
sessions keep pushing, a single-turn run can opt in, and a seeded run
can opt out.
Co-Authored-By: Claude <[email protected]>
---------
Co-authored-by: Claude <[email protected]>
---
providers/common/ai/docs/message_history.rst | 10 ++++-
providers/common/ai/docs/operators/agent.rst | 3 ++
providers/common/ai/docs/stability.rst | 4 ++
providers/common/ai/docs/troubleshooting.rst | 2 +-
.../airflow/providers/common/ai/operators/agent.py | 14 ++++++-
.../tests/unit/common/ai/operators/test_agent.py | 45 ++++++++++++++--------
6 files changed, 59 insertions(+), 19 deletions(-)
diff --git a/providers/common/ai/docs/message_history.rst
b/providers/common/ai/docs/message_history.rst
index 453b4d7a7c3..fa9ee78befe 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 fbb7b542989..9dbf792c9c9 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 5feb431be94..3fad7b4ca4f 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 ea3e8c3fae5..e2b329e74dc 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 06f052cb6c2..2392a1fdd19 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 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator,
HITLReviewMixin, Usa
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 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator,
HITLReviewMixin, Usa
# 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 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator,
HITLReviewMixin, Usa
# 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 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
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 70d7faa9e41..178db398166 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 @@ class TestAgentOperatorMessageHistory:
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 @@ class TestAgentOperatorMessageHistory:
@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."""
