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 Strands and ADK.

Memory is not the right tool when you need the conversation itself:

Install

Memory is in pydantic-ai-harness 0.24.0 and later, the version the code-mode extra already requires:

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:

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:

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

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))],
        )

On the next run, Memory adds what earlier runs wrote to the model’s request, next to the prompt:

<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:

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

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[Customer](
                                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?"},
            ]
        )

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 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.

Strands and ADK

Strands and ADK agents run in your own @task (see Agent frameworks), 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.

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.

Was this entry helpful?