Update app.py

app.py

@@ -38,6 +38,17 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 # Database configuration
 DATABASE_PATH = '/tmp/chat_database.db'
 
+# Database connection manager
+@contextmanager
+def get_db_connection():
+    """Context manager for database connections"""
+    conn = sqlite3.connect(DATABASE_PATH)
+    conn.row_factory = sqlite3.Row
+    try:
+        yield conn
+    finally:
+        conn.close()
+
 def initialize_model():
     """Initialize the model with appropriate settings"""
     try:
@@ -78,57 +89,6 @@ def initialize_model():
         print(f"Error initializing model: {e}")
         raise
 
-# Initialize database
-def init_db():
-    """Initialize the database"""
-    try:
-        with get_db_connection() as conn:
-            conn.execute('''
-                CREATE TABLE IF NOT EXISTS chats (
-                    id TEXT PRIMARY KEY,
-                    title TEXT,
-                    date TEXT,
-                    last_message TEXT
-                )
-            ''')
-
-            conn.execute('''
-                CREATE TABLE IF NOT EXISTS messages (
-                    id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    chat_id TEXT,
-                    role TEXT,
-                    content TEXT,
-                    timestamp TEXT,
-                    FOREIGN KEY (chat_id) REFERENCES chats (id)
-                )
-            ''')
-
-            conn.execute('''
-                CREATE TABLE IF NOT EXISTS important_info (
-                    id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    chat_id TEXT,
-                    content TEXT,
-                    FOREIGN KEY (chat_id) REFERENCES chats (id)
-                )
-            ''')
-            conn.commit()
-    except Exception as e:
-        print(f"Error initializing database: {e}")
-        raise
-
-# Initialize the application
-try:
-    print("Initializing database...")
-    init_db()
-    print("Database initialized successfully")
-    
-    print("Starting model initialization...")
-    llm = initialize_model()
-    print("Model initialized successfully")
-except Exception as e:
-    print(f"Initialization error: {e}")
-    raise
-
 
 class ChatSession:
     def __init__(self, session_id):
@@ -235,6 +195,44 @@ class ChatSession:
         self.important_info = []
 
 
+def init_db():
+    """Initialize the database"""
+    try:
+        with get_db_connection() as conn:
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS chats (
+                    id TEXT PRIMARY KEY,
+                    title TEXT,
+                    date TEXT,
+                    last_message TEXT
+                )
+            ''')
+
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id TEXT,
+                    role TEXT,
+                    content TEXT,
+                    timestamp TEXT,
+                    FOREIGN KEY (chat_id) REFERENCES chats (id)
+                )
+            ''')
+
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS important_info (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id TEXT,
+                    content TEXT,
+                    FOREIGN KEY (chat_id) REFERENCES chats (id)
+                )
+            ''')
+            conn.commit()
+    except Exception as e:
+        print(f"Error initializing database: {e}")
+        raise
+
+        
 # Rest of the prompt template and other configurations remain the same
 prompt_template = """
 Role: You are Figr Code Assistant, specializing in providing clear, error-free Python code solutions.
@@ -660,6 +658,17 @@ def home():
     return render_template("index.html")
 
 
-if __name__ == "__main__":
+try:
+    print("Initializing database...")
     init_db()
+    print("Database initialized successfully")
+    
+    print("Starting model initialization...")
+    llm = initialize_model()
+    print("Model initialized successfully")
+except Exception as e:
+    print(f"Initialization error: {e}")
+    raise
+
+if __name__ == "__main__":
     app.run(host="0.0.0.0", port=PORT)