import os import asyncpg _pool: asyncpg.Pool | None = None DDL = """ CREATE TABLE IF NOT EXISTS messages ( id BIGSERIAL PRIMARY KEY, session_id TEXT NOT NULL, role TEXT NOT NULL, -- 'user' | 'assistant' | 'tool' content TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); CREATE INDEX IF NOT EXISTS idx_messages_session ON messages (session_id, created_at); """ async def get_pool() -> asyncpg.Pool: global _pool if _pool is None: dsn = os.environ["NEON_DATABASE_URL"] # e.g. postgresql://user:pass@ep-xxx.neon.tech/db?sslmode=require # asyncpg wants sslmode stripped, ssl handled separately dsn = dsn.replace("?sslmode=require", "").replace("&sslmode=require", "") _pool = await asyncpg.create_pool(dsn=dsn, ssl="require", min_size=1, max_size=5) async with _pool.acquire() as conn: await conn.execute(DDL) return _pool async def load_history(session_id: str, limit: int = 20) -> list[dict]: pool = await get_pool() rows = await pool.fetch( """ SELECT role, content FROM messages WHERE session_id = $1 ORDER BY created_at DESC LIMIT $2 """, session_id, limit, ) return [{"role": r["role"], "content": r["content"]} for r in reversed(rows)] async def save_message(session_id: str, role: str, content: str) -> None: pool = await get_pool() await pool.execute( "INSERT INTO messages (session_id, role, content) VALUES ($1, $2, $3)", session_id, role, content, )