Spaces:
Sleeping
Sleeping
File size: 1,533 Bytes
43384a3 | 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 | # memory_db.py - TiDB Database Operations
import pymysql
from config import settings
class MemoryDB:
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()
def create_tables(self):
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,
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 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()
|