database.py

import sqlite3
from typing import Optional

class Database:
    def __init__(self, db_path="chat_history.db"):
        self.db_path = db_path
        self.conn = None

    def connect(self):
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self.create_tables()

    def create_tables(self):
        c = self.conn.cursor()
        c.execute('''CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE,
            password TEXT
        )''')
        c.execute('''CREATE TABLE IF NOT EXISTS conversations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER,
            conversation TEXT,
            FOREIGN KEY(user_id) REFERENCES users(id)
        )''')
        self.conn.commit()

    def add_user(self, username, password):
        c = self.conn.cursor()
        try:
            c.execute("INSERT INTO users (username, password) VALUES (_, _)", (username, password))
            self.conn.commit()
            return True
        except sqlite3.IntegrityError:
            return False

    def get_user(self, username, password) -> Optional[int]:
        c = self.conn.cursor()
        c.execute("SELECT id FROM users WHERE username=? AND password=?", (username, password))
        row = c.fetchone()
        return row[0] if row else None

    def add_conversation(self, user_id, conversation):
        c = self.conn.cursor()
        c.execute("INSERT INTO conversations (user_id, conversation) VALUES (?, ?)", (user_id, conversation))
        self.conn.commit()

    def get_conversations(self, user_id):
        c = self.conn.cursor()
        c.execute("SELECT conversation FROM conversations WHERE user_id=? ORDER BY id DESC LIMIT 1", (user_id,))
        row = c.fetchone()
        return row[0] if row else None

    def clear_conversation(self, user_id):
        c = self.conn.cursor()
        c.execute("DELETE FROM conversations WHERE user_id=?", (user_id,))
        self.conn.commit()