Multi-turn sessions and message history

By default each agent run is a cold, single-turn conversation. To carry a conversation across runs – a chat or iterative agent where “and the third one?” must resolve against an earlier answer – pass message_history.

When message_history is set, the operator seeds the run with those prior turns and, after the run, pushes the full updated transcript (result.all_messages()) to XCom under the key message_history. The next run reads it back to resume the conversation. None (the default) keeps the single-turn behavior unchanged.

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 JSON file in object storage by session_id (use s3:// / gs:// in a deployment); the first run starts from an empty "[]".

airflow/providers/common/ai/example_dags/example_agent.py[source]

@dag(tags=["example"], params={"session_id": "demo-session"})
def example_agent_session():
    """Resume a conversation across runs via ``message_history``.

    The agent step seeds itself with the prior transcript and re-emits the
    updated transcript to XCom (key ``message_history``). Loading and storing
    that transcript under a session key is the DAG's job -- here, a JSON file in
    object storage keyed by ``session_id``. Swap the path for ``s3://`` /
    ``gs://`` in a deployment.
    """
    sessions_root = ObjectStoragePath("file:///tmp/airflow_agent_sessions")

    @task
    def load_history(session_id: str) -> str:
        path = sessions_root / f"{session_id}.json"
        # First turn: no file yet -> start a fresh session (empty transcript).
        return path.read_text() if path.exists() else "[]"

    @task.agent(
        llm_conn_id="pydanticai_default",
        system_prompt="You are a helpful assistant. Use the earlier turns for context.",
        # The XComArg both wires the dependency and resolves to the JSON transcript.
        message_history=load_history("{{ params.session_id }}"),
    )
    def ask(question: str) -> str:
        return question

    @task
    def save_history(session_id: str, transcript: str) -> None:
        # Local/fsspec object storage does not auto-create parent dirs on write.
        sessions_root.mkdir(parents=True, exist_ok=True)
        (sessions_root / f"{session_id}.json").write_text(transcript)

    answer = ask("And what did I ask you a moment ago?")
    saved = save_history(
        "{{ params.session_id }}",
        # The agent step pushes the post-run transcript under this XCom key.
        "{{ ti.xcom_pull(task_ids='ask', key='message_history') }}",
    )
    # save runs after the agent so the pulled transcript is the fresh one.
    answer >> saved

message_history accepts a list of pydantic-ai ModelMessage objects or their JSON form (str / bytes), so the value emitted to XCom feeds straight back in on the next run. When pulling it via a template, pass default='[]' (as above) so the first run – which has no XCom yet – starts a fresh session instead of trying to parse the string "None".

Note

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.

Message format, and limiting session size

The transcript is stored as pydantic-ai’s JSON message format. It is cumulative: each turn appends to it, so it grows for the life of the session. The operator does not trim it itself, but capabilities= is passed to the pydantic-ai Agent, so the ProcessHistory capability from pydantic-ai’s documentation works on AgentOperator as is:

from pydantic_ai.capabilities import ProcessHistory
from pydantic_ai.messages import ModelMessage, ModelRequest, RetryPromptPart, ToolReturnPart


def keep_recent(messages: list[ModelMessage]) -> list[ModelMessage]:
    tail = messages[-10:]
    # Start at a user turn. A tail that begins at a request holding only tool
    # returns or retries (mid-run) loses the user's question.
    for i, message in enumerate(tail):
        if isinstance(message, ModelRequest) and not any(
            isinstance(part, (ToolReturnPart, RetryPromptPart)) for part in message.parts
        ):
            return tail[i:]
    return messages


AgentOperator(
    task_id="chat",
    ...,
    capabilities=[ProcessHistory(keep_recent)],
)

pydantic-ai writes the processed list back into the run state, and result.all_messages() returns that list. With a history processor set, the transcript pushed to XCom is therefore the trimmed one, and turns the processor drops are no longer part of the stored session. Alternatively, read the message_history XCom in a downstream task, trim it, store it in your own session store, and load it from there on the next run; or configure an object-storage XCom backend for long sessions.

The operator has no built-in cross-session memory. To let an agent use facts from earlier sessions, store them yourself, load them in an upstream task (the same load, run, store pattern as above), and put them in the prompt or the instructions. To let the agent keep its own notes across Dag runs instead, see Agent memory.

Was this entry helpful?