From dd06abe88fa7c7d80c12c4093760104cdf4f284d Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Wed, 14 Oct 2026 02:36:13 +0100 Subject: [PATCH 1/2] Document agent memory across Dag runs for common.ai Agents in common.ai run on Pydantic AI, so they can use pydantic-ai-harness Memory through capabilities= today. Show it for a recurring AgentOperator (SQLite) and for per-customer memory in your own agent (Postgres), and say how Strands and ADK agents use their own memory. --- providers/common/ai/docs/features.rst | 3 + providers/common/ai/docs/memory.rst | 163 ++++++++++++++++++ .../ai/example_dags/example_agent_memory.py | 135 +++++++++++++++ 3 files changed, 301 insertions(+) create mode 100644 providers/common/ai/docs/memory.rst create mode 100644 providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_memory.py diff --git a/providers/common/ai/docs/features.rst b/providers/common/ai/docs/features.rst index b7df68cf6fe35..3885569eced39 100644 --- a/providers/common/ai/docs/features.rst +++ b/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 Message history + Agent memory Capabilities and guardrails Code mode Approve outputs diff --git a/providers/common/ai/docs/memory.rst b/providers/common/ai/docs/memory.rst new file mode 100644 index 0000000000000..6576dac338d93 --- /dev/null +++ b/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 `__ +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.md + + sync_crm hits HTTP 429 at 02:00 (first seen 2026-10-13) + + +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 `__. + +.. _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//`` on the worker's disk by default, and +``BedrockKnowledgeBaseStore``, which keeps memory in an Amazon Bedrock knowledge base. +See the `Strands memory module `__. + +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 `__. diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_memory.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_memory.py new file mode 100644 index 0000000000000..dd28017bc2e74 --- /dev/null +++ b/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() From d2a587ceffa882fc6cc2ae10926c306b51ff6e61 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Wed, 14 Oct 2026 03:26:02 +0100 Subject: [PATCH 2/2] Skip the agent memory example when pydantic-ai-harness is missing Every Dag in example_agent_memory.py needs pydantic-ai-harness, which only the code-mode extra installs, so in CI the file defines no Dags and the import test fails. --- airflow-core/tests/unit/always/test_example_dags.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airflow-core/tests/unit/always/test_example_dags.py b/airflow-core/tests/unit/always/test_example_dags.py index ec942b27810d2..3c7a69331593a 100644 --- a/airflow-core/tests/unit/always/test_example_dags.py +++ b/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.