Metadata-Version: 2.4
Name: pya-memory
Version: 0.0.1
Summary: Pydantic-AI native message store: persist ModelMessage history to SQLite, Postgres, or DynamoDB with one line.
Project-URL: Homepage, https://github.com/aristide1997/pya-memory
Project-URL: Repository, https://github.com/aristide1997/pya-memory
Project-URL: Issues, https://github.com/aristide1997/pya-memory/issues
Project-URL: Changelog, https://github.com/aristide1997/pya-memory/blob/main/CHANGELOG.md
Author-email: Aristide Fumagalli <aristide1997@gmail.com>
License: MIT
License-File: LICENSE
Keywords: chat-history,dynamodb,llm,message-store,postgres,pydantic-ai,sqlite
Classifier: Development Status :: 3 - Alpha
Classifier: Framework :: AsyncIO
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Software Development :: Libraries
Requires-Python: >=3.11
Requires-Dist: aiosqlite<1,>=0.20
Requires-Dist: pydantic-ai<2,>=1.97
Provides-Extra: dynamodb
Requires-Dist: aiobotocore<4,>=3.7; extra == 'dynamodb'
Provides-Extra: postgres
Requires-Dist: psycopg[binary,pool]<4,>=3.2; extra == 'postgres'
Description-Content-Type: text/markdown

# pya-memory

[![tests](https://github.com/aristide1997/pya-memory/actions/workflows/test.yml/badge.svg)](https://github.com/aristide1997/pya-memory/actions/workflows/test.yml)
[![PyPI](https://img.shields.io/pypi/v/pya-memory.svg)](https://pypi.org/project/pya-memory/)
[![Python versions](https://img.shields.io/pypi/pyversions/pya-memory.svg)](https://pypi.org/project/pya-memory/)
[![License](https://img.shields.io/pypi/l/pya-memory.svg)](LICENSE)

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

```python
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

```bash
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:
```bash
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`](https://github.com/pydantic/pydantic-ai-harness)
("batteries for your Pydantic AI agent") package — see
[PR #176](https://github.com/pydantic/pydantic-ai-harness/pull/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.

```python
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`](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:

```python
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

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

```python
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

```python
# 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

```python
# 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:

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

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

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

```python
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`](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

```bash
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`](.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`](.github/workflows/publish.yml).

## License

MIT
