pya-memory 0.0.1

Pydantic-AI native message store: persist ModelMessage history to SQLite, Postgres, or DynamoDB with one line.

pya-memory

Pydantic-AI native message store. Persist pydantic_ai.messages.ModelMessage history to SQLite, Postgres, DynamoDB, or in-memory, and plug into any pydantic_ai.Agent with one capability — no boilerplate.

from pydantic_ai import Agent
from pya_memory import SQLiteMessageStore

store = SQLiteMessageStore("chat.sqlite")
agent = Agent(
    "openai-chat:gpt-5.4-nano",
    capabilities=[store.capability(conversation_id="thread-1", owner_id="user-42")],
)
result = await agent.run("Hello")
# history loaded automatically; new messages persisted automatically

Install

uv add pya-memory                  # SQLite + in-memory backends
uv add 'pya-memory[postgres]'      # adds PostgresMessageStore
uv add 'pya-memory[dynamodb]'      # adds DynamoDBMessageStore (+ S3 offload)

Or with pip:

pip install pya-memory
pip install 'pya-memory[postgres]'
pip install 'pya-memory[dynamodb]'

Requires Python 3.11+ and pydantic-ai >= 1.97.

Relationship to pydantic-ai-harness

pya-memory is an unaffiliated, third-party package. The Pydantic team is adding an official SessionPersistence capability to their own pydantic-ai-harness ("batteries for your Pydantic AI agent") package — see PR #176, with SQLite / Redis / Postgres backends on the roadmap.

If you want the officially-blessed path, use that. pya-memory exists for people who need production-grade SQL persistence today — multi-process-safe SQLite (WAL + BEGIN IMMEDIATE + retry), Postgres with advisory-locked migrations and a connection pool, DynamoDB with transparent S3 offload for oversize messages — plus schema migrations, owner-mismatch protection for multi-tenant safety, transactional O(1) aggregate counters, and conversation forking. Pick whichever fits your stage.

How it works

The capability hooks pydantic-ai's lifecycle:

• before_model_request (fires once per agent run): loads stored history and prepends it to the request. Idempotent across multi-step runs (tool loops don't duplicate history).
• after_run (fires once per agent.run() / run_sync() / run_stream()): calls result.new_messages() and appends to the store transactionally.

Both happen inside the existing pydantic-ai extension points — no monkeypatching, no wrappers around Agent.

Capping what the model sees per turn

StoreCapability loads the full conversation on every run. For long chats that's wasted cost — combine it with pydantic-ai's first-class ProcessHistory capability to trim the loaded history before it reaches the model. The store keeps everything; the model sees only what you allow.

from pydantic_ai import Agent
from pydantic_ai.capabilities import ProcessHistory
from pya_memory import SQLiteMessageStore

store = SQLiteMessageStore("chat.sqlite")

def keep_last_n(messages):
    return messages[-20:]

agent = Agent(
    "openai-chat:gpt-5.4-nano",
    capabilities=[
        store.capability(conversation_id="thread-1"),  # loads full history
        ProcessHistory(keep_last_n),                    # trims before model sees
    ],
)

Ordering matters: StoreCapability first (it loads), ProcessHistory second (it trims). See examples/with_history_window.py for a token-budget variant.

Caveats. The naive messages[-n:] shown above is fine for plain text chat but has sharp edges you should design around:

• Tool-call pairing. If the cutoff lands between a ToolCallPart and its matching ToolReturnPart, the model sees an orphaned return. Providers may error or hallucinate. A safe trimmer walks back to a request/response boundary, not an arbitrary index.
• Request/response alternation. pydantic-ai expects history to alternate request → response. Slicing without checking can yield [response, response] after a tool loop and confuse the model.
• Instructions and system prompts. ModelRequest.instructions lives on the request that introduced it — if you trim that message away, the guidance goes with it.
• Forgetting vs. summarising. Truncation drops context; it doesn't summarise it. For long-running chats, pair the window with a CompactionPart summary of the dropped prefix.
• Token budgets. For real budgeting use the model's tokenizer (pydantic_ai's count_tokens), not a per-message heuristic.

pya-memory currently loads the full conversation before ProcessHistory trims it — fine for hundreds of messages, wasteful at tens of thousands. A future load_limit= on store.capability(...) would push the limit down into the store.

Manual mode (no capability)

If you prefer explicit control, the capability is just sugar over two methods:

store = SQLiteMessageStore("chat.sqlite")
history = await store.load_messages(conversation_id="thread-1")
result = await agent.run("Hello", message_history=history)
await store.append_messages(
    conversation_id="thread-1",
    messages=result.new_messages(),
    owner_id="user-42",
)
await store.aclose()  # optional, at app shutdown

The store opens its connection lazily — no await .open() needed. async with store: ... is supported but not required.

Backends

Class	Driver	Use when
InMemoryMessageStore()	none	tests, ephemeral chats
SQLiteMessageStore(path)	aiosqlite	single-process, file-backed
PostgresMessageStore(conninfo)	psycopg + AsyncConnectionPool	multi-process / production
DynamoDBMessageStore(table_name)	aioboto3	AWS-native, scale-to-zero pricing

All four implement the same async MessageStore ABC.

DynamoDB

from pya_memory import DynamoDBMessageStore

store = DynamoDBMessageStore(
    table_name="pya_memory",
    region_name="us-east-1",          # or via AWS_REGION env
)

On first method call the store creates a single table (PK conversation_id

• range key, one GSI on owner_id + updated_at). On-demand billing by default — at idle, you pay only for storage (~$0.25/GB-month).

Image / large-message offload to S3. DynamoDB caps each item at 400 KB, which BinaryContent images can quickly exceed. Configure a bucket and oversize messages are transparently written to S3, with a tiny pointer kept in DynamoDB:

store = DynamoDBMessageStore(
    table_name="pya_memory",
    large_message_bucket="my-app-pya-memory",
    large_message_prefix="messages/",   # default
    large_message_threshold=300_000,    # bytes; default leaves headroom
)

Without a bucket configured, an oversize message raises SerializationError.

Documented divergences from SQL backends:

Behaviour	Why
Max 99 messages per append_messages() call	TransactWriteItems caps at 100 ops; one slot reserved for the aggregate update. Typical agent runs produce 2–10 messages, so this only bites bulk imports.
list_conversations(owner_id=None) is a Scan	DynamoDB can't sort an unbounded keyspace; pass owner_id whenever possible.
Per-message inline cap ~398 KB	DynamoDB item limit. Configure large_message_bucket=... for transparent S3 offload, or use ImageUrl references.
GSI reads are eventually consistent	A freshly-created conversation may take milliseconds to appear in list_conversations.
delete_conversation does NOT delete S3 blobs	Forked conversations share blob references. Use a bucket lifecycle policy for cleanup, or refcount manually if you need stricter behaviour.

Querying

# Filter the slice you want — every column is indexed
await store.load_messages(
    conversation_id="thread-1",
    kinds=["response"],                          # only model responses
    models=["openai-chat:gpt-5.4-nano"],              # only this model
    since=datetime(2026, 5, 1, tzinfo=UTC),     # only recent
    limit=50, offset=0,
)

# Cheap aggregates — precomputed, no GROUP BY needed
info = await store.get_conversation(conversation_id="thread-1")
print(info.message_count, info.total_input_tokens, info.total_output_tokens)

# List per owner, newest first
for c in await store.list_conversations(owner_id="user-42", limit=20):
    print(c.conversation_id, c.updated_at, c.message_count)

Forking conversations

# Copy a conversation; both can diverge independently from this point.
await store.fork_conversation(source="thread-1", new_id="thread-1-branch-a")

Owner scoping

owner_id is optional. Pass None (default) for single-tenant use. The first write sets the owner; subsequent appends with a different owner_id raise OwnerMismatchError.

Errors

All raised exceptions subclass pya_memory.StoreError:

Exception	When
StoreClosedError	Method called after aclose()
OwnerMismatchError	Appending to a conversation owned by someone else
ConversationNotFoundError	fork_conversation(source=X) where X doesn't exist
ConversationExistsError	fork_conversation(new_id=X) where X already exists
SerializationError	ModelMessage failed to round-trip via TypeAdapter

Migrations

The schema is owned by numbered SQL files under pya_memory/backends/{sqlite,postgres}/migrations/. On first connect the store applies any unapplied migrations and records them in a pya_memory_schema_version table — serialised across processes via BEGIN IMMEDIATE (SQLite) or pg_advisory_lock (Postgres). Databases that pre-date this migration system (i.e. created by hand before pya-memory existed) are adopted as-is: the initial migration is idempotent, so existing tables are left alone and v1 is recorded as applied.

If you'd rather drive migrations from a deploy step than on first request, pass auto_migrate=False and call await store.migrate() explicitly:

store = PostgresMessageStore(conninfo, auto_migrate=False)
await store.migrate()        # run pending migrations now
result = await agent.run(...)

Schema

Two tables — aggregates separated from the event log so reads stay cheap:

CREATE TABLE conversations (
    conversation_id TEXT PRIMARY KEY,
    owner_id TEXT,
    created_at TIMESTAMPTZ NOT NULL,
    updated_at TIMESTAMPTZ NOT NULL,
    message_count INTEGER NOT NULL DEFAULT 0,
    total_input_tokens BIGINT NOT NULL DEFAULT 0,
    total_output_tokens BIGINT NOT NULL DEFAULT 0,
    last_model_name TEXT,
    title TEXT
);

CREATE TABLE messages (
    id BIGSERIAL PRIMARY KEY,                          -- INTEGER AUTOINCREMENT on sqlite
    conversation_id TEXT NOT NULL REFERENCES conversations(conversation_id) ON DELETE CASCADE,
    message_index INTEGER NOT NULL,
    message_json JSONB NOT NULL,                       -- TEXT on sqlite
    kind TEXT NOT NULL,                                -- 'request' | 'response'
    run_id TEXT,
    model_name TEXT,
    has_tool_calls BOOLEAN NOT NULL,
    has_thinking BOOLEAN NOT NULL,
    has_files BOOLEAN NOT NULL,
    has_compaction BOOLEAN NOT NULL,
    has_native_tools BOOLEAN NOT NULL,
    input_tokens INTEGER,
    output_tokens INTEGER,
    timestamp TIMESTAMPTZ,
    UNIQUE(conversation_id, message_index)
);

append_messages runs as a single transaction (insert messages + update aggregates), so concurrent writers see a consistent view.

Logging

The library uses stdlib logging and is silent by default (a NullHandler is attached at import time). To see what it's doing, configure logging in your host application:

import logging
logging.basicConfig(level=logging.INFO)
logging.getLogger("pya_memory").setLevel("INFO")   # or DEBUG for per-call traces

Logger namespace:

pya_memory                              # root (NullHandler)
├── pya_memory.backends.memory
├── pya_memory.backends.sqlite.store
├── pya_memory.backends.postgres.store
├── pya_memory.backends.dynamodb.store
├── pya_memory.backends.dynamodb._migrations
├── pya_memory.capability
└── pya_memory._serde

Level policy:

Level	What you see
DEBUG	Per-call entry/exit with duration_ms, count, bytes
INFO	Store open/close, migration applied, DynamoDB→S3 offload
WARNING	SQLite BEGIN IMMEDIATE retry, DynamoDB transient retry, owner-mismatch
ERROR	Retry exhausted, migration failed, transaction rolled back

Structured fields are attached to each LogRecord via extra=, so a JSON or logfire handler can read them directly:

for record in caplog.records:
    print(record.duration_ms, record.conv_id, record.backend)

PII discipline: the library never logs message content, tool arguments, or owner_id values. Only stable identifiers (conv_id), sizes (bytes, count), durations (duration_ms), and error classes (error_class) appear in records. The contract is pinned by tests/test_logging.py.

Extending

Subclass MessageStore to add a new backend (Redis, DuckDB, SQLAlchemy session, ...). The capability and all helpers will work unchanged — they only depend on the ABC.

Per-backend customisation hooks (overridable metadata extraction, extra indexed columns, external connection injection) are intentionally deferred — open an issue if you have a concrete use case.

Development

make sync           # install with dev deps
make test           # memory + sqlite + postgres + dynamodb (auto-provisioned via testcontainers)
make postgres-up    # start a persistent local Postgres container
make dynamodb-up    # start DynamoDB Local + MinIO (S3-compatible) containers
make test-live      # run the 7 live-model tests against gpt-5.4-nano
                    #   (needs OPENAI_API_KEY in .env; ~$0.005 / run)
                    # also runs the DynamoDB live smoke test against real AWS
                    #   when AWS_ACCESS_KEY_ID is in .env (~$0.0001 / run)

CI runs the full non-live suite on every push and PR via GitHub Actions (.github/workflows/test.yml) against Python 3.11–3.13 and both the pinned and latest 1.x pydantic-ai. Releases publish to PyPI on tags matching v* via .github/workflows/publish.yml.

License

MIT