Document agent memory across Dag runs for common.ai

airflow-core/tests/unit/always/test_example_dags.py

@@ -48,6 +48,8 @@
     # value: a dictionary containing package distributions with an optional version specifier, e.g., >=2.3.4
     # yandexcloud is automatically removed in case botocore is upgraded to latest
     r".*example_yandexcloud.*\.py": {"yandexcloud": None},
+    # Every Dag in it uses pydantic-ai-harness Memory, which only the opt-in code-mode extra installs
+    r".*example_agent_memory\.py": {"pydantic-ai-harness": ">=0.24.0"},
 }
 IGNORE_AIRFLOW_PROVIDER_DEPRECATION_WARNING: tuple[str, ...] = (
     # Certain examples or system tests may trigger AirflowProviderDeprecationWarnings.

providers/common/ai/docs/features.rst

@@ -26,6 +26,8 @@ you use. Each is a parameter on the operator or decorator.
 - :doc:`structured_output`: ``output_type`` returns a typed Pydantic object through XCom
   instead of a string.
 - :doc:`message_history`: ``message_history`` carries a conversation across agent runs.
+- :doc:`memory`: pydantic-ai-harness ``Memory`` gives an agent notes it writes in one Dag
+  run and reads in the next.
 - :doc:`capabilities`: ``capabilities=`` adds pydantic-ai capabilities such as ``Thinking`` and
   ``WebSearch``, and ``pydantic-ai-shields`` guardrails, to an agent.
 - :doc:`code_mode`: the ``CodeMode`` capability lets the model call several tools from one
@@ -50,6 +52,7 @@ Making retries cheap with ``durable=True`` is a reliability feature and lives un
 
     Structured output <structured_output>
     Message history <message_history>
+    Agent memory <memory>
     Capabilities and guardrails <capabilities>
     Code mode <code_mode>
     Approve outputs <approval_gates>

providers/common/ai/docs/memory.rst

@@ -0,0 +1,163 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _howto/memory:
+
+Agent memory
+============
+
+Memory lets an agent keep notes that outlive the Dag run. A triage agent that runs
+every morning can record a cause it diagnosed today and recognize it tomorrow; a
+support agent can remember a customer's setup between tickets. The agent decides
+what to write, and the next run reads it back.
+
+This provider does not ship a memory layer of its own. Agents run on Pydantic AI, so
+they use `Pydantic AI's Memory capability <https://pydantic.dev/docs/ai/harness/memory/>`__
+from ``pydantic-ai-harness``, passed in ``capabilities=`` like any other capability.
+Strands and ADK agents use their framework's own memory; see
+:ref:`howto/memory:other-frameworks`.
+
+Memory is not the right tool when you need the conversation itself:
+
+- To continue a conversation turn by turn, pass the transcript with
+  :doc:`message_history`.
+- To avoid repeating model calls when a task retries, use :doc:`durable_execution`.
+  It replays one task instance's steps and deletes its cache when the task succeeds.
+
+Install
+-------
+
+``Memory`` is in ``pydantic-ai-harness`` 0.24.0 and later, the version the
+``code-mode`` extra already requires:
+
+.. code-block:: bash
+
+    pip install "pydantic-ai-harness>=0.24.0"
+
+``pydantic-ai-harness`` is pre-1.0, and its API can change between minor releases.
+
+How it works
+------------
+
+``Memory`` gives the agent a notebook of Markdown files. At the start of each model
+request it adds a bounded excerpt of the main file, ``MEMORY.md``, to the prompt, and it
+gives the model four tools: ``write_memory``, ``read_memory``, ``search_memory`` and
+``delete_memory``. The notebook is stored in the store you pass to ``Memory``.
+
+Pick the store by where your tasks run. Each task runs in its own process, and with the
+Celery or Kubernetes executor the next run can be on a different machine, so a store on
+the worker's local disk is not there for the next run:
+
+.. list-table::
+    :header-rows: 1
+    :widths: 30 70
+
+    * - Store
+      - Use it when
+    * - ``SqliteMemoryStore(database=...)``
+      - Every run of the task executes on the same machine, such as a single-machine
+        deployment. The file's directory must exist.
+    * - ``PostgresMemoryStore(pool)``
+      - Tasks run on several machines, or several tasks write the same memory at once.
+    * - ``InMemoryStore()``
+      - Tests. It is lost when the task's process exits.
+
+Remember across Dag runs
+------------------------
+
+This Dag runs every day and gives a triage agent the failed tasks an upstream task collects. The
+system prompt tells the agent to check its notes for known causes and to record new
+ones. The notebook lives in a SQLite file, so tomorrow's run sees what today's run
+wrote:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_memory.py
+    :language: python
+    :start-after: [START howto_agent_memory_sqlite]
+    :end-before: [END howto_agent_memory_sqlite]
+
+On the next run, ``Memory`` adds what earlier runs wrote to the model's request, next
+to the prompt:
+
+.. code-block:: text
+
+    <memory>
+    ### MEMORY.md
+
+    sync_crm hits HTTP 429 at 02:00 (first seen 2026-10-13)
+    </memory>
+
+One memory per customer
+-----------------------
+
+When one agent serves many customers, give each one its own notes with a namespace.
+``namespace`` is resolved by your code from the run's dependencies, not by the model, so
+the agent cannot read or write another customer's notes through its tools.
+
+``PostgresMemoryStore`` takes an ``asyncpg`` pool, which you create with ``await`` and
+close when the run is done. A Dag file cannot do either around an ``AgentOperator`` run,
+so this example runs its own Pydantic AI agent in a ``@task``, and creates the pool,
+runs the agent and closes the pool in one event loop. The database credentials come from an Airflow connection:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_memory.py
+    :language: python
+    :start-after: [START howto_agent_memory_postgres]
+    :end-before: [END howto_agent_memory_postgres]
+
+The mapped task runs once per customer, and each run reads and writes only its
+customer's notes. ``asyncpg`` is installed with the Postgres provider. See
+:doc:`frameworks/pydantic_ai` for what a task running its own agent gets from this
+provider and what it does not.
+
+What to watch for
+-----------------
+
+- Memory is written by the model, and it goes back into later prompts. A prompt
+  injection that reaches the agent once can persist in its notes. Set
+  ``inject_memory=False`` when less-trusted input reaches the agent, so the model reads
+  memory only when it calls a tool.
+- Do not put secrets in prompts to an agent with memory: anything the model sees, it can
+  write to the store.
+- Namespaces separate customers in the agent's tools. They are not access control on
+  the database: anyone with the connection's credentials can read every namespace.
+- The stores have no expiry. What is injected is bounded: the guidance, ``MEMORY.md``
+  and the list of other files share ``max_tokens`` (about 2,000 tokens by default), and
+  ``MEMORY.md`` is cut at ``max_lines`` (200). The model reaches the rest with
+  ``read_memory`` or ``search_memory``. Prune notes from the store when they stop being
+  useful.
+
+For every ``Memory`` parameter, the injection budget and the search limits, see the
+`Pydantic AI Memory documentation <https://pydantic.dev/docs/ai/harness/memory/>`__.
+
+.. _howto/memory:other-frameworks:
+
+Strands and ADK
+---------------
+
+Strands and ADK agents run in your own ``@task`` (see :doc:`frameworks/index`), so they
+use the memory their framework ships. The same rule about where the store lives applies.
+
+A Strands agent takes a ``MemoryManager`` with one or more stores in its
+``memory_manager`` argument. Strands ships ``FileMemoryStore``, which writes under
+``./.strands/memory/<store-name>/`` on the worker's disk by default, and
+``BedrockKnowledgeBaseStore``, which keeps memory in an Amazon Bedrock knowledge base.
+See the `Strands memory module <https://github.com/strands-agents/harness-sdk/tree/main/strands-py/src/strands/memory>`__.
+
+An ADK agent gets memory from the ``Runner``: pass a memory service as
+``memory_service``, add sessions to it with ``add_session_to_memory``, and give the
+agent the ``load_memory`` or ``preload_memory`` tool. ``InMemoryMemoryService`` is lost
+when the task's process exits; ``VertexAiMemoryBankService`` keeps memory in Vertex AI.
+See the `ADK memory documentation <https://adk.dev/sessions/memory/>`__.

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_memory.py

@@ -0,0 +1,135 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Give an agent memory that outlives the Dag run, with pydantic-ai-harness ``Memory``.
+
+``Memory`` gives the agent a notebook it reads at the start of a run and writes to with
+tools. The notebook lives in a store you choose, so a run tomorrow sees what the agent
+wrote today.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from dataclasses import dataclass
+
+from pydantic_ai import Agent
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+try:
+    from pydantic_ai_harness import Memory
+    from pydantic_ai_harness.memory import PostgresMemoryStore, SqliteMemoryStore
+except ImportError:
+    Memory = None  # type: ignore[assignment,misc]
+
+try:
+    import asyncpg
+except ImportError:
+    asyncpg = None
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "pydanticai_default")
+MEMORY_DB = os.environ.get("AGENT_MEMORY_DB", "/opt/airflow/agent-memory/triage.db")
+MEMORY_CONN_ID = os.environ.get("AGENT_MEMORY_CONN_ID", "agent_memory_db")
+
+
+# [START howto_agent_memory_sqlite]
+if Memory is not None:
+
+    @dag(schedule="@daily", catchup=False, tags=["example"])
+    def example_agent_memory():
+        @task
+        def collect_failures() -> list[str]:
+            # Replace with a query against your monitoring system.
+            return [
+                "load_orders: psycopg2.OperationalError: could not connect to server",
+                "sync_crm: HTTP 429 Too Many Requests from api.example-crm.com",
+            ]
+
+        collect_failures() >> AgentOperator(
+            task_id="triage_failures",
+            llm_conn_id=LLM_CONN_ID,
+            prompt="Triage the failed tasks from {{ ds }}:\n{{ ti.xcom_pull(task_ids='collect_failures') }}",
+            system_prompt=(
+                "You triage failed data pipeline tasks. Your notes hold the causes you have "
+                "seen before. Say which failures match a known cause and which are new. "
+                "Record each new cause in your notes with the date you first saw it, and "
+                "add the date to a known cause each time it comes back."
+            ),
+            capabilities=[Memory(SqliteMemoryStore(database=MEMORY_DB))],
+        )
+
+    # [END howto_agent_memory_sqlite]
+
+    example_agent_memory()
+
+
+# [START howto_agent_memory_postgres]
+if Memory is not None and asyncpg is not None:
+
+    @dataclass
+    class Customer:
+        customer_id: str
+
+    @dag(schedule=None, tags=["example"])
+    def example_agent_memory_per_customer():
+        @task
+        def answer_ticket(customer_id: str, question: str) -> str:
+            conn = BaseHook.get_connection(MEMORY_CONN_ID)
+
+            async def run() -> str:
+                # The pool belongs to this event loop, so it is created and closed in it.
+                async with asyncpg.create_pool(
+                    host=conn.host,
+                    port=conn.port or 5432,
+                    user=conn.login,
+                    password=conn.password,
+                    database=conn.schema,
+                ) as pool:
+                    agent = Agent(
+                        PydanticAIHook.get_hook(LLM_CONN_ID).get_conn(),
+                        deps_type=Customer,
+                        instructions=(
+                            "You answer support tickets. Your notes hold what you have learned "
+                            "about this customer's setup. Record anything that will help with "
+                            "their next ticket."
+                        ),
+                        capabilities=[
+                            Memory(
+                                PostgresMemoryStore(pool),
+                                namespace=lambda ctx: ctx.deps.customer_id,
+                            )
+                        ],
+                    )
+                    result = await agent.run(question, deps=Customer(customer_id))
+                    return result.output
+
+            return asyncio.run(run())
+
+        answer_ticket.expand_kwargs(
+            [
+                {"customer_id": "acme", "question": "Our nightly export stopped at 02:00. Why?"},
+                {"customer_id": "globex", "question": "Can we move our sync to every 15 minutes?"},
+            ]
+        )
+
+    # [END howto_agent_memory_postgres]
+
+    example_agent_memory_per_customer()