Update db.py

db.py

@@ -1,7 +1,8 @@
-# db.py (Corrected)
+# db.py (Improved)
+import os
 import sqlite3
 import logging
-import json # Import json
+import json
 
 logging.basicConfig(
     level=logging.INFO,
@@ -11,11 +12,22 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 class SessionDB:
-    def __init__(self, db_path="sessions.db"):
-        self.conn = sqlite3.connect(db_path, check_same_thread=False)
-        self.cursor = self.conn.cursor()
-        self.create_table()
-    
+    def __init__(self, db_path="db/sessions.db"):
+        # Ensure parent folder exists
+        os.makedirs(os.path.dirname(db_path), exist_ok=True)
+
+        try:
+            self.conn = sqlite3.connect(db_path, check_same_thread=False)
+            self.cursor = self.conn.cursor()
+            self.create_table()
+            logger.info(f"Database initialized at {db_path}")
+        except Exception as e:
+            logger.error(f"Failed to initialize database at {db_path}: {e}")
+            logger.warning("Falling back to in-memory database")
+            self.conn = sqlite3.connect(":memory:", check_same_thread=False)
+            self.cursor = self.conn.cursor()
+            self.create_table()
+
     def create_table(self):
         try:
             self.cursor.execute('''
@@ -27,25 +39,26 @@ class SessionDB:
             self.conn.commit()
         except Exception as e:
             logger.error(f"Error creating sessions table: {e}")
-    
+
     def get_history(self, session_id, max_turns=4):
         try:
             self.cursor.execute("SELECT history FROM sessions WHERE session_id=?", (session_id,))
             result = self.cursor.fetchone()
             if result:
-                # Use json.loads for safety instead of eval()
-                history = json.loads(result[0]) 
+                history = json.loads(result[0])  # safe parsing
                 return history[-max_turns:]
             return []
         except Exception as e:
             logger.error(f"Error retrieving history for session {session_id}: {e}")
             return []
-    
+
     def save_history(self, session_id, history):
         try:
-            # Use json.dumps for safety instead of str()
-            history_str = json.dumps(history) 
-            self.cursor.execute("INSERT OR REPLACE INTO sessions (session_id, history) VALUES (?, ?)", (session_id, history_str))
+            history_str = json.dumps(history)  # safe serialization
+            self.cursor.execute(
+                "INSERT OR REPLACE INTO sessions (session_id, history) VALUES (?, ?)",
+                (session_id, history_str)
+            )
             self.conn.commit()
         except Exception as e:
-            logger.error(f"Error saving history for session {session_id}: {e}")
+            logger.error(f"Error saving history for session {session_id}: {e}")