cleanup

Author: jssmith

openai_agents/memory/postgres_session.py

@@ -56,42 +56,6 @@ class PostgresSessionConfig(BaseModel):
     idempotence_table: str = "activity_idempotence"
 
 
-async def init_schema(conn: asyncpg.Connection, config: PostgresSessionConfig) -> None:
-    """Initialize the PostgreSQL schema."""
-    async with conn.transaction():
-        # Create sessions table
-        sessions_ddl = f"""
-            CREATE TABLE IF NOT EXISTS {config.sessions_table} (
-                session_id TEXT NOT NULL,
-                created_at TIMESTAMP DEFAULT NOW(),
-                updated_at TIMESTAMP DEFAULT NOW(),
-                PRIMARY KEY (session_id)
-            )
-        """
-        await conn.execute(sessions_ddl)
-
-        # Create operation_id sequence
-        operation_id_ddl = f"""
-            CREATE SEQUENCE IF NOT EXISTS {config.operation_id_sequence} START 1
-        """
-        await conn.execute(operation_id_ddl)
-
-        # Create messages table
-        messages_ddl = f"""
-            CREATE TABLE IF NOT EXISTS {config.messages_table} (
-                session_id TEXT NOT NULL,
-                operation_id INTEGER NOT NULL DEFAULT nextval('{config.operation_id_sequence}'),
-                message_data TEXT NOT NULL,
-                created_at TIMESTAMP DEFAULT NOW(),
-                deleted_at TIMESTAMP NULL,
-                PRIMARY KEY (session_id, operation_id),
-                FOREIGN KEY (session_id) 
-                REFERENCES {config.sessions_table} (session_id)
-                ON DELETE CASCADE
-            )
-        """
-        await conn.execute(messages_ddl)
-
 
 class PostgresSessionGetItemsRequest(BaseModel):
     config: PostgresSessionConfig
@@ -284,3 +248,50 @@ def _get_connection():
         if _connection_factory is None:
             raise ValueError("Connection factory not set")
         return _connection_factory()
+
+    @staticmethod
+    async def init_schema(config: PostgresSessionConfig) -> None:
+        conn = PostgresSession._get_connection()
+        """Initialize the PostgreSQL schema."""
+        async with conn.transaction():
+            # Create sessions table
+            sessions_ddl = f"""
+                CREATE TABLE IF NOT EXISTS {config.sessions_table} (
+                    session_id TEXT NOT NULL,
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    updated_at TIMESTAMP DEFAULT NOW(),
+                    PRIMARY KEY (session_id)
+                )
+            """
+            await conn.execute(sessions_ddl)
+
+            # Create operation_id sequence
+            operation_id_ddl = f"""
+                CREATE SEQUENCE IF NOT EXISTS {config.operation_id_sequence} START 1
+            """
+            await conn.execute(operation_id_ddl)
+
+            # Create messages table
+            messages_ddl = f"""
+                CREATE TABLE IF NOT EXISTS {config.messages_table} (
+                    session_id TEXT NOT NULL,
+                    operation_id INTEGER NOT NULL DEFAULT nextval('{config.operation_id_sequence}'),
+                    message_data TEXT NOT NULL,
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    deleted_at TIMESTAMP NULL,
+                    PRIMARY KEY (session_id, operation_id),
+                    FOREIGN KEY (session_id) 
+                    REFERENCES {config.sessions_table} (session_id)
+                    ON DELETE CASCADE
+                )
+            """
+            await conn.execute(messages_ddl)
+
+    @staticmethod
+    def get_activities() -> list[Callable[[], activity.Activity]]:
+        return [
+            postgres_session_get_items_activity,
+            postgres_session_add_items_activity,
+            postgres_session_pop_item_activity,
+            postgres_session_clear_session_activity,
+        ]

openai_agents/memory/run_postgres_session_worker.py

@@ -14,12 +14,7 @@
 )
 from openai_agents.memory.postgres_session import (
     PostgresSessionConfig,
-    init_schema,
     PostgresSession,
-    postgres_session_get_items_activity,
-    postgres_session_add_items_activity,
-    postgres_session_pop_item_activity,
-    postgres_session_clear_session_activity,
 )
 from openai_agents.memory.db_utils import IdempotenceHelper
 
@@ -32,12 +27,18 @@ async def main():
         messages_table="session_messages",
         sessions_table="session",
         operation_id_sequence="session_operation_id_sequence",
+        idempotence_table="activity_idempotence",
     )
-    PostgresSession.set_connection_factory(lambda: db_connection)
-    await init_schema(db_connection, config=postgres_session_config)
-    idempotence_helper = IdempotenceHelper(table_name="activity_idempotence")
+
+    # Create the idempotence table. This is used to ensure that activities are idempotent with
+    # respect to database modifications.
+    idempotence_helper = IdempotenceHelper(table_name=postgres_session_config.idempotence_table)
     await idempotence_helper.create_table(db_connection)
+
+    # Configure the Postgres Session with the database connection.
+    # Initialize the schema.
     PostgresSession.set_connection_factory(lambda: db_connection)
+    await PostgresSession.init_schema(config=postgres_session_config)
 
     # Create client connected to server at the given address
     client = await Client.connect(
@@ -53,16 +54,11 @@ async def main():
 
     worker = Worker(
         client,
-        task_queue="openai-agents-memory-task-queue",
+        task_queue="openai-postgres-session-task-queue",
         workflows=[
             PostgresSessionWorkflow,
         ],
-        activities=[
-            postgres_session_get_items_activity,
-            postgres_session_add_items_activity,
-            postgres_session_pop_item_activity,
-            postgres_session_clear_session_activity,
-        ],
+        activities=[*PostgresSession.get_activities()],
     )
     await worker.run()