Update memory_db.py

memory_db.py

@@ -1,43 +1,102 @@
-# memory_db.py - TiDB Database Operations
+# memory_db.py - TiDB Integration
 import pymysql
-from config import settings
+import os
+from datetime import datetime
 
-class MemoryDB:
+class TiDBMemory:
     def __init__(self):
-        self.connection = pymysql.connect(
-            host=settings.TIDB_HOST,
-            port=settings.TIDB_PORT,
-            user=settings.TIDB_USER,
-            password=settings.TIDB_PASSWORD,
-            database=settings.TIDB_DATABASE,
-            ssl={'ssl': {'ca': ''}}  # TiDB SSL
-        )
-        self.create_tables()
+        self.connection = None
+        self.connect()
+    
+    def connect(self):
+        try:
+            self.connection = pymysql.connect(
+                host=os.getenv("TIDB_HOST", "gateway01.ap-southeast-1.prod.aws.tidbcloud.com"),
+                port=int(os.getenv("TIDB_PORT", 4000)),
+                user=os.getenv("TIDB_USER", "2Rg6kfo2rNEB3PN.root"),
+                password=os.getenv("TIDB_PASSWORD", "9JJabiRfo0WpH9FP"),
+                database=os.getenv("TIDB_DATABASE", "test"),
+                ssl={'ssl': {'ca': ''}}
+            )
+            self.create_tables()
+            print("✅ TiDB connected successfully")
+        except Exception as e:
+            print(f"❌ TiDB connection failed: {e}")
+            self.connection = None
     
     def create_tables(self):
+        if not self.connection:
+            return
+            
         with self.connection.cursor() as cursor:
             cursor.execute("""
                 CREATE TABLE IF NOT EXISTS chat_history (
                     id INT AUTO_INCREMENT PRIMARY KEY,
                     user_input TEXT,
                     ai_response MEDIUMTEXT,
+                    language_mode VARCHAR(20),
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS code_snippets (
+                    id INT AUTO_INCREMENT PRIMARY KEY,
+                    code_type VARCHAR(100),
+                    code_content LONGTEXT,
+                    description TEXT,
+                    usage_count INT DEFAULT 0,
                     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                 )
             """)
             self.connection.commit()
     
-    def save_conversation(self, user_input, ai_response):
-        with self.connection.cursor() as cursor:
-            cursor.execute(
-                "INSERT INTO chat_history (user_input, ai_response) VALUES (%s, %s)",
-                (user_input, ai_response)
-            )
-            self.connection.commit()
+    def save_chat(self, user_input, ai_response, language_mode):
+        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)
+                )
+                self.connection.commit()
+                return True
+        except Exception as e:
+            print(f"❌ Save chat error: {e}")
+            return False
     
-    def get_recent_conversations(self, limit=10):
-        with self.connection.cursor() as cursor:
-            cursor.execute(
-                "SELECT user_input, ai_response FROM chat_history ORDER BY created_at DESC LIMIT %s",
-                (limit,)
-            )
-            return cursor.fetchall()
+    def get_recent_chats(self, limit=10):
+        if not self.connection:
+            return []
+            
+        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()
+        except Exception as e:
+            print(f"❌ Get chats error: {e}")
+            return []
+    
+    def save_code_snippet(self, code_type, code_content, description):
+        if not self.connection:
+            return False
+            
+        try:
+            with self.connection.cursor() as cursor:
+                cursor.execute(
+                    """INSERT INTO code_snippets (code_type, code_content, description) 
+                       VALUES (%s, %s, %s)""",
+                    (code_type, code_content, description)
+                )
+                self.connection.commit()
+                return True
+        except Exception as e:
+            print(f"❌ Save code error: {e}")
+            return False
+
+# Global instance
+tidb_memory = TiDBMemory()