add pydantic ai memory examples

Author: samuelcolvin

pai-memory/README.md

@@ -0,0 +1,7 @@
+# Pydantic AI memory examples
+
+Run postgres in docker with
+
+```bash
+docker run -e POSTGRES_HOST_AUTH_METHOD=trust --rm -it --name pg -p 5432:5432 -d postgres
+```

pai-memory/with_messages.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+import asyncpg
+from pydantic_ai import Agent
+from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
+
+# hack to get around asyncpg's poor typing support
+if TYPE_CHECKING:
+    DbConn = asyncpg.Connection[asyncpg.Record]
+else:
+    DbConn = asyncpg.Connection
+
+
+import logfire
+
+logfire.configure(service_name='mem-msgs')
+logfire.instrument_pydantic_ai()
+logfire.instrument_asyncpg()
+
+
+@asynccontextmanager
+async def db() -> AsyncIterator[DbConn]:
+    conn = await asyncpg.connect('postgresql://postgres@localhost:5432')
+    await conn.execute("""
+        create table if not exists messages(
+            id serial primary key,
+            ts timestamp not null default now(),
+            user_id integer not null,
+            messages json not null
+        )
+    """)
+
+    try:
+        yield conn
+    finally:
+        await conn.close()
+
+
+agent = Agent(
+    'openai:gpt-4o',
+    instructions='You are a helpful assistant.',
+)
+
+
+@logfire.instrument
+async def run_agent(prompt: str, user_id: int):
+    async with db() as conn:
+        with logfire.span('retrieve messages'):
+            messages: list[ModelMessage] = []
+            for row in await conn.fetch('SELECT messages FROM messages WHERE user_id = $1 order by ts', user_id):
+                messages += ModelMessagesTypeAdapter.validate_json(row[0])
+
+        result = await agent.run(prompt, message_history=messages)
+        print(result.output)
+
+        with logfire.span('record messages'):
+            msgs = result.new_messages_json().decode()
+            await conn.execute('INSERT INTO messages(user_id, messages) VALUES($1, $2)', user_id, msgs)
+
+
+async def memory_messages():
+    await run_agent('My name is Samuel.', 123)
+
+    await run_agent('What is my name?', 123)
+
+
+if __name__ == '__main__':
+    asyncio.run(memory_messages())

pai-memory/with_tools.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import asyncpg
+from pydantic_ai import Agent, RunContext
+
+# hack to get around asyncpg's poor typing support
+if TYPE_CHECKING:
+    DbConn = asyncpg.Connection[asyncpg.Record]
+else:
+    DbConn = asyncpg.Connection
+
+
+import logfire
+
+logfire.configure(service_name='mem-tool')
+logfire.instrument_pydantic_ai()
+logfire.instrument_asyncpg()
+
+
+@asynccontextmanager
+async def db(reset: bool = False) -> AsyncIterator[DbConn]:
+    conn = await asyncpg.connect('postgresql://postgres@localhost:5432')
+    if reset:
+        await conn.execute('drop table if exists memory')
+    await conn.execute("""
+        create table if not exists memory(
+            id serial primary key,
+            user_id integer not null,
+            value text not null,
+            unique(user_id, value)
+        )
+    """)
+
+    try:
+        yield conn
+    finally:
+        await conn.close()
+
+
+@dataclass
+class Deps:
+    user_id: int
+    conn: DbConn
+
+
+agent = Agent(
+    'openai:gpt-4o',
+    deps_type=Deps,
+    instructions='You are a helpful assistant.',
+)
+
+
+@agent.tool
+async def record_memory(ctx: RunContext[Deps], value: str) -> str:
+    """Use this tool to store information in memory."""
+    await ctx.deps.conn.execute(
+        'insert into memory(user_id, value) values($1, $2) on conflict do nothing',
+        ctx.deps.user_id,
+        value,
+    )
+    return 'Value added to memory.'
+
+
+@agent.tool
+async def retrieve_memories(ctx: RunContext[Deps], memory_contains: str) -> str:
+    """Get all memories about the user."""
+    rows = await ctx.deps.conn.fetch(
+        'select value from memory where user_id = $1 and value ilike $2',
+        ctx.deps.user_id,
+        f'%{memory_contains}%',
+    )
+    return '\n'.join(row[0] for row in rows)
+
+
+async def memory_tools():
+    async with db(True) as conn:
+        deps = Deps(123, conn)
+        result = await agent.run('My name is Samuel.', deps=deps)
+        print(result.output)
+
+    # time goes by...
+
+    async with db() as conn:
+        deps = Deps(123, conn)
+        result = await agent.run('What is my name?', deps=deps)
+        print(result.output)
+
+
+if __name__ == '__main__':
+    asyncio.run(memory_tools())

pyproject.toml

@@ -5,9 +5,11 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
+    "asyncpg>=0.30.0",
+    "asyncpg-stubs>=0.30.2",
     "devtools>=0.12.2",
     "fastapi>=0.115.14",
-    "logfire[fastapi,httpx]>=3.21.1",
+    "logfire[asyncpg,fastapi,httpx]>=3.21.1",
     "pydantic-ai>=0.3.4",
 ]