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:
To continue a conversation turn by turn, pass the transcript with Multi-turn sessions and message history.
To avoid repeating model calls when a task retries, use 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:
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 |
|---|---|
|
Every run of the task executes on the same machine, such as a single-machine deployment. The file’s directory must exist. |
|
Tasks run on several machines, or several tasks write the same memory at once. |
|
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:
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:
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=Falsewhen 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.mdand the list of other files sharemax_tokens(about 2,000 tokens by default), andMEMORY.mdis cut atmax_lines(200). The model reaches the rest withread_memoryorsearch_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.