Pydantic AI Memory and Conversation History

"""Teaching example: server-owned, versioned chat snapshots, not an auth system."""
import sqlite3
import time
from contextlib import closing
from dataclasses import dataclass
from typing import Callable
from uuid import uuid4

from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter


class ChatUnavailable(Exception):
    """Same response for unknown, wrong-owner, deleted and expired chats."""


class VersionConflict(Exception):
    """Another writer has already changed this chat."""


@dataclass(frozen=True)
class Snapshot:
    version: int
    messages: list[ModelMessage]


class ChatStore:
    def __init__(self, path: str, clock: Callable[[], float] = time.time,
                 max_bytes: int = 256_000):
        self.path, self.clock, self.max_bytes = path, clock, max_bytes
        with closing(sqlite3.connect(path)) as db, db:
            db.execute('''CREATE TABLE IF NOT EXISTS chats (
                owner TEXT NOT NULL, thread TEXT NOT NULL,
                version INTEGER NOT NULL, payload BLOB NOT NULL,
                expires REAL NOT NULL, deleted INTEGER NOT NULL DEFAULT 0,
                format INTEGER NOT NULL DEFAULT 1,
                PRIMARY KEY (owner, thread))''')

    def create(self, owner: str, ttl_seconds: int = 86_400) -> str:
        if not owner or ttl_seconds <= 0:
            raise ValueError('Owner and positive lifetime are required')
        thread = str(uuid4())
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute('INSERT INTO chats(owner,thread,version,payload,expires) VALUES (?,?,?,?,?)',
                       (owner, thread, 0, b'[]', self.clock() + ttl_seconds))
        return thread

    def load(self, owner: str, thread: str) -> Snapshot:
        with closing(sqlite3.connect(self.path)) as db:
            row = db.execute('SELECT version,payload,format FROM chats '
                             'WHERE owner=? AND thread=? AND deleted=0 AND expires>?',
                             (owner, thread, self.clock())).fetchone()
        if row is None:
            raise ChatUnavailable('Chat unavailable')
        if row[2] != 1:
            raise ValueError('Unsupported stored format')
        if len(row[1]) > self.max_bytes:
            raise ValueError('Stored history exceeds byte limit')
        return Snapshot(row[0], ModelMessagesTypeAdapter.validate_json(row[1]))

    def save(self, owner: str, thread: str, expected_version: int,
             messages: list[ModelMessage]) -> int:
        payload = ModelMessagesTypeAdapter.dump_json(messages)
        ModelMessagesTypeAdapter.validate_json(payload)
        if len(payload) > self.max_bytes:
            raise ValueError('History exceeds byte limit; do not silently truncate')
        with closing(sqlite3.connect(self.path)) as db, db:
            changed = db.execute('UPDATE chats SET payload=?,version=version+1 '
                                 'WHERE owner=? AND thread=? AND version=? '
                                 'AND deleted=0 AND expires>?',
                                 (payload, owner, thread, expected_version, self.clock())).rowcount
        if changed != 1:
            self.load(owner, thread)
            raise VersionConflict('Reload the chat before continuing')
        return expected_version + 1

    def delete(self, owner: str, thread: str) -> None:
        # A tombstone prevents an in-flight writer from resurrecting this ID.
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute('UPDATE chats SET payload=?,deleted=1,version=version+1 '
                       'WHERE owner=? AND thread=?', (b'[]', owner, thread))

To use a real model, pass a configured Agent into reply() and install its provider extra. Do not put API keys in this file. Provider calls can cost money and transmit the history you supply. Start with fictional conversations and verify the provider’s applicable data-handling settings before introducing real user information.

Not automatically. The sample clears the active payload and blocks subsequent access to that thread. Backups, logs, caches and provider-side retention need separate handling according to your actual systems and applicable policies. Do not promise complete erasure based only on one database update.