99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
|
import sqlite3
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
|
||
|
|
class DBManager:
|
||
|
|
def __init__(self, db_name='chat.db'):
|
||
|
|
self.conn = sqlite3.connect(db_name)
|
||
|
|
self.cursor = self.conn.cursor()
|
||
|
|
self._initialize_database()
|
||
|
|
|
||
|
|
def _initialize_database(self):
|
||
|
|
self.execute_query('''
|
||
|
|
CREATE TABLE IF NOT EXISTS chats (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
title TEXT NOT NULL DEFAULT 'New Chat',
|
||
|
|
created_at DATETIME NOT NULL,
|
||
|
|
updated_at DATETIME NOT NULL
|
||
|
|
)
|
||
|
|
''')
|
||
|
|
self.execute_query('''
|
||
|
|
CREATE TABLE IF NOT EXISTS chat_messages (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
chat_id INTEGER,
|
||
|
|
category TEXT,
|
||
|
|
content TEXT,
|
||
|
|
created_at DATETIME NOT NULL,
|
||
|
|
updated_at DATETIME NOT NULL,
|
||
|
|
FOREIGN KEY(chat_id) REFERENCES chats(id)
|
||
|
|
)
|
||
|
|
''')
|
||
|
|
self.execute_query('''
|
||
|
|
CREATE TABLE IF NOT EXISTS generated_files (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
chat_message_id INTEGER,
|
||
|
|
name TEXT,
|
||
|
|
content BLOB,
|
||
|
|
created_at DATETIME NOT NULL,
|
||
|
|
updated_at DATETIME NOT NULL,
|
||
|
|
FOREIGN KEY(chat_message_id) REFERENCES chat_messages(id)
|
||
|
|
)
|
||
|
|
''')
|
||
|
|
|
||
|
|
def execute_query(self, query, params=()):
|
||
|
|
self.cursor.execute(query, params)
|
||
|
|
self.conn.commit()
|
||
|
|
|
||
|
|
def fetch_query(self, query, params=()):
|
||
|
|
self.cursor.execute(query, params)
|
||
|
|
return self.cursor.fetchall()
|
||
|
|
|
||
|
|
# Chat Operations
|
||
|
|
def save_chat(self, title):
|
||
|
|
now = datetime.now()
|
||
|
|
self.execute_query('''
|
||
|
|
INSERT INTO chats (title, created_at, updated_at)
|
||
|
|
VALUES (?, ?, ?)
|
||
|
|
''', (title, now, now))
|
||
|
|
return self.cursor.lastrowid
|
||
|
|
|
||
|
|
def update_chat_title(self, chat_id, title):
|
||
|
|
now = datetime.now()
|
||
|
|
self.execute_query('''
|
||
|
|
UPDATE chats
|
||
|
|
SET title = ?, updated_at = ?
|
||
|
|
WHERE id = ?
|
||
|
|
''', (title, now, chat_id))
|
||
|
|
|
||
|
|
def get_chats(self):
|
||
|
|
return self.fetch_query("SELECT * FROM chats ORDER BY updated_at DESC")
|
||
|
|
|
||
|
|
def get_chat(self, chat_id):
|
||
|
|
return self.fetch_query("SELECT * FROM chats WHERE id = ?", (chat_id,))[0]
|
||
|
|
|
||
|
|
# Chat Message Operations
|
||
|
|
def save_message(self, chat_id, category, content):
|
||
|
|
now = datetime.now()
|
||
|
|
self.execute_query('''
|
||
|
|
INSERT INTO chat_messages (chat_id, category, content, created_at, updated_at)
|
||
|
|
VALUES (?, ?, ?, ?, ?)
|
||
|
|
''', (chat_id, category, content, now, now))
|
||
|
|
return self.cursor.lastrowid
|
||
|
|
|
||
|
|
def get_chat_messages(self, chat_id):
|
||
|
|
return self.fetch_query("SELECT * FROM chat_messages WHERE chat_id = ?", (chat_id,))
|
||
|
|
|
||
|
|
def get_chat_message(self, chat_message_id):
|
||
|
|
return self.fetch_query("SELECT * FROM chat_messages WHERE id = ?", (chat_message_id,))[0]
|
||
|
|
|
||
|
|
# File Operations
|
||
|
|
def save_file(self, chat_message_id, name, content):
|
||
|
|
now = datetime.now()
|
||
|
|
self.execute_query('''
|
||
|
|
INSERT INTO generated_files (chat_message_id, name, content, created_at, updated_at)
|
||
|
|
VALUES (?, ?, ?, ?, ?)
|
||
|
|
''', (chat_message_id, name, content, now, now))
|
||
|
|
|
||
|
|
def get_generated_files(self, chat_message_id):
|
||
|
|
return self.fetch_query("SELECT * FROM generated_files WHERE chat_message_id = ?", (chat_message_id,))
|