Add PostgreSQL-based session management for agent memory

examples/postgresql_session/main.py

@@ -0,0 +1,38 @@
+from agents import Agent, Runner,set_tracing_disabled,OpenAIChatCompletionsModel
+from openai import AsyncOpenAI
+from pg_session import PostgreSQLSession
+from dotenv import load_dotenv
+import asyncio
+import os
+
+load_dotenv(override=True)
+set_tracing_disabled(disabled=True)
+openai_client = AsyncOpenAI(
+    api_key=os.getenv("OPEN_API_KEY",""),
+    base_url="https://api.openai.com/v1",
+)
+
+async def main():
+    # Create agent
+    agent = Agent(
+        name="Assistant",
+        instructions="Reply very concisely.",
+        model=OpenAIChatCompletionsModel(
+            model="gpt-4o-mini",
+            openai_client=openai_client,
+        )
+    )
+
+    # Create a session instance with a session ID
+    session = PostgreSQLSession("conversation_123",neon_url=os.getenv("NEON_DB_URL"))
+
+    # await session.clear_session()
+
+    result = await Runner.run(
+        agent,
+        "what is my name",
+        session=session
+    )
+    print(result.final_output)
+
+asyncio.run(main())

examples/postgresql_session/pg_session.py

@@ -0,0 +1,74 @@
+import asyncpg
+import json
+
+class PostgreSQLSession:
+    def __init__(self, session_id: str, neon_url: str):
+        self.session_id = session_id
+        self.__neon_url = neon_url
+        self.__conn = None
+
+    async def __connect(self):
+        if not self.__conn:
+            self.__conn = await asyncpg.connect(self.__neon_url)
+            await self.__create_tables()
+
+    async def __create_tables(self):
+        await self.__conn.execute("""
+            CREATE TABLE IF NOT EXISTS sessions (
+                session_id TEXT PRIMARY KEY,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            );
+        """)
+        await self.__conn.execute("""
+            CREATE TABLE IF NOT EXISTS messages (
+                id SERIAL PRIMARY KEY,
+                session_id TEXT,
+                message_data TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            );
+        """)
+
+    async def add_items(self, items: list[dict])->None:
+        await self.__connect()
+        # Ensure session row exists
+        await self.__conn.execute("""
+            INSERT INTO sessions (session_id)
+            VALUES ($1)
+            ON CONFLICT DO NOTHING;
+        """, self.session_id)
+
+        # Insert each message
+        for item in items:
+            await self.__conn.execute("""
+                INSERT INTO messages (session_id, message_data)
+                VALUES ($1, $2);
+            """, self.session_id, json.dumps(item))
+
+    async def get_items(self) -> list[dict]:
+        await self.__connect()
+        rows = await self.__conn.fetch("""
+            SELECT message_data FROM messages
+            WHERE session_id = $1
+            ORDER BY created_at ASC;
+        """, self.session_id)
+        return [json.loads(row['message_data']) for row in rows]
+
+    async def pop_item(self) -> dict | None:
+        await self.__connect()
+        row = await self.__conn.fetchrow("""
+            SELECT id, message_data FROM messages
+            WHERE session_id = $1
+            ORDER BY created_at DESC
+            LIMIT 1;
+        """, self.session_id)
+
+        if row:
+            await self.__conn.execute("DELETE FROM messages WHERE id = $1;", row["id"])
+            return json.loads(row["message_data"])
+        return None
+
+    async def clear_session(self)->None:
+        await self.__connect()
+        await self.__conn.execute("""
+            DELETE FROM messages WHERE session_id = $1;
+        """, self.session_id)

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "requests>=2.0, <3",
     "types-requests>=2.0, <3",
     "mcp>=1.11.0, <2; python_version >= '3.10'",
+    "asyncpg>=0.30.0",
 ]
 classifiers = [
     "Typing :: Typed",