42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
|
import sqlite3
|
||
|
|
import json
|
||
|
|
from typing import List, Dict, Any
|
||
|
|
|
||
|
|
class Database:
|
||
|
|
def __init__(self, db_file='telegram_bot.db'):
|
||
|
|
self.db_file = db_file
|
||
|
|
self.conn = sqlite3.connect(self.db_file)
|
||
|
|
self.create_tables()
|
||
|
|
|
||
|
|
def create_tables(self):
|
||
|
|
with self.conn:
|
||
|
|
self.conn.execute('''
|
||
|
|
CREATE TABLE IF NOT EXISTS conversations (
|
||
|
|
user_id INTEGER PRIMARY KEY,
|
||
|
|
history TEXT
|
||
|
|
)
|
||
|
|
''')
|
||
|
|
|
||
|
|
def save_conversation(self, user_id: int, history: List[Dict[str, Any]]):
|
||
|
|
with self.conn:
|
||
|
|
self.conn.execute('''
|
||
|
|
INSERT OR REPLACE INTO conversations (user_id, history)
|
||
|
|
VALUES (?, ?)
|
||
|
|
''', (user_id, json.dumps(history)))
|
||
|
|
|
||
|
|
def get_conversation(self, user_id: int) -> List[Dict[str, Any]]:
|
||
|
|
cursor = self.conn.execute('SELECT history FROM conversations WHERE user_id = ?', (user_id,))
|
||
|
|
result = cursor.fetchone()
|
||
|
|
if result:
|
||
|
|
return json.loads(result[0])
|
||
|
|
return []
|
||
|
|
|
||
|
|
def clear_conversation(self, user_id: int):
|
||
|
|
with self.conn:
|
||
|
|
self.conn.execute('DELETE FROM conversations WHERE user_id = ?', (user_id,))
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
self.conn.close()
|
||
|
|
|
||
|
|
# Create a global instance of the Database class
|
||
|
|
db = Database()
|