Refactor database initialization in SQLiteSessionMemory

knowsuchagency · knowsuchagency · commit 57b4f5e8b83b · 2025-05-24T13:24:53.000-07:00
- Introduced a new method `_init_db_for_connection` to handle database schema initialization for a specific connection.
- Updated the `_init_db` method to call the new method, improving clarity and separation of concerns.
- Added a comment to indicate the initialization of the database schema for the connection.
diff --git a/src/agents/memory/session_memory.py b/src/agents/memory/session_memory.py
@@ -96,11 +96,12 @@ def _get_connection(self) -> sqlite3.Connection:
                 check_same_thread=False,
             )
             self._local.connection.execute("PRAGMA journal_mode=WAL")
+            # Initialize the database schema for this connection
+            self._init_db_for_connection(self._local.connection)
         return self._local.connection
 
-    def _init_db(self) -> None:
-        """Initialize the database schema."""
-        conn = self._get_connection()
+    def _init_db_for_connection(self, conn: sqlite3.Connection) -> None:
+        """Initialize the database schema for a specific connection."""
         conn.execute(
             f"""
             CREATE TABLE IF NOT EXISTS {self.sessions_table} (
@@ -132,6 +133,11 @@ def _init_db(self) -> None:
 
         conn.commit()
 
+    def _init_db(self) -> None:
+        """Initialize the database schema."""
+        conn = self._get_connection()
+        # The schema initialization is now handled in _init_db_for_connection
+
     async def get_messages(self, session_id: str) -> list[TResponseInputItem]:
         """Retrieve the conversation history for a given session.
 
