Update core/memory_db.py

core/memory_db.py

@@ -29,6 +29,7 @@ class TiDBMemory:
             return
             
         with self.connection.cursor() as cursor:
+            # Create chat_history table
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS chat_history (
                     id INT AUTO_INCREMENT PRIMARY KEY,
@@ -38,6 +39,16 @@ class TiDBMemory:
                     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                 )
             """)
+            
+            # Check if language_mode column exists, if not add it
+            try:
+                cursor.execute("SHOW COLUMNS FROM chat_history LIKE 'language_mode'")
+                if not cursor.fetchone():
+                    print("⚠️ Adding missing 'language_mode' column to chat_history")
+                    cursor.execute("ALTER TABLE chat_history ADD COLUMN language_mode VARCHAR(20) DEFAULT 'en'")
+            except Exception as e:
+                print(f"⚠️ Column check failed: {e}")
+            
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS code_snippets (
                     id INT AUTO_INCREMENT PRIMARY KEY,
@@ -50,16 +61,29 @@ class TiDBMemory:
             """)
             self.connection.commit()
     
-    def save_chat(self, user_input, ai_response, language_mode):
+    def save_chat(self, user_input, ai_response, language_mode="en"):
         if not self.connection:
             return False
             
         try:
             with self.connection.cursor() as cursor:
-                cursor.execute(
-                    "INSERT INTO chat_history (user_input, ai_response, language_mode) VALUES (%s, %s, %s)",
-                    (user_input, ai_response, language_mode)
-                )
+                # Try with language_mode column
+                try:
+                    cursor.execute(
+                        "INSERT INTO chat_history (user_input, ai_response, language_mode) VALUES (%s, %s, %s)",
+                        (user_input, ai_response, language_mode)
+                    )
+                except pymysql.err.OperationalError as e:
+                    # If language_mode column doesn't exist, try without it
+                    if "Unknown column 'language_mode'" in str(e):
+                        print("⚠️ language_mode column missing, inserting without it")
+                        cursor.execute(
+                            "INSERT INTO chat_history (user_input, ai_response) VALUES (%s, %s)",
+                            (user_input, ai_response)
+                        )
+                    else:
+                        raise e
+                
                 self.connection.commit()
                 return True
         except Exception as e:
@@ -72,11 +96,26 @@ class TiDBMemory:
             
         try:
             with self.connection.cursor() as cursor:
-                cursor.execute(
-                    "SELECT user_input, ai_response, language_mode FROM chat_history ORDER BY created_at DESC LIMIT %s",
-                    (limit,)
-                )
-                return cursor.fetchall()
+                # Try with language_mode column
+                try:
+                    cursor.execute(
+                        "SELECT user_input, ai_response, language_mode FROM chat_history ORDER BY created_at DESC LIMIT %s",
+                        (limit,)
+                    )
+                    return cursor.fetchall()
+                except pymysql.err.OperationalError as e:
+                    # If language_mode column doesn't exist, select without it
+                    if "Unknown column 'language_mode'" in str(e):
+                        print("⚠️ language_mode column missing, selecting without it")
+                        cursor.execute(
+                            "SELECT user_input, ai_response FROM chat_history ORDER BY created_at DESC LIMIT %s",
+                            (limit,)
+                        )
+                        # Add default language_mode for compatibility
+                        rows = cursor.fetchall()
+                        return [(row[0], row[1], "en") for row in rows]
+                    else:
+                        raise e
         except Exception as e:
             print(f"❌ Get chats error: {e}")
             return []