Spaces:
Sleeping
Sleeping
File size: 5,757 Bytes
3f47213 43384a3 3f47213 43384a3 3f47213 43384a3 3f47213 43384a3 3f47213 43384a3 5e12eac 43384a3 3f47213 5e12eac 3f47213 43384a3 5e12eac 3f47213 5e12eac 3f47213 43384a3 3f47213 5e12eac 3f47213 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | # memory_db.py - TiDB Integration
import pymysql
import os
from datetime import datetime
class TiDBMemory:
def __init__(self):
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:
# Create chat_history table
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
)
""")
# 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,
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_chat(self, user_input, ai_response, language_mode="en"):
if not self.connection:
return False
try:
with self.connection.cursor() as cursor:
# 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:
print(f"❌ Save chat error: {e}")
return False
def get_recent_chats(self, limit=10):
if not self.connection:
return []
try:
with self.connection.cursor() as cursor:
# 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 []
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() |