CoolFace
Apppublic

shamiimuhammad/NexGen-Support-Hero

sourceHugging Faceupdated 8mo agoView on Hugging Face
0likes
database.py44 linesDownload Raw Back to src
1import sqlite32from langchain_core.messages import HumanMessage, AIMessage3 4DB_PATH = "chat_history.db"5 6def init_db():7    """Logic: Creates the database file if it doesn't exist."""8    conn = sqlite3.connect(DB_PATH)9    cursor = conn.cursor()10    cursor.execute('''11        CREATE TABLE IF NOT EXISTS messages (12            id INTEGER PRIMARY KEY AUTOINCREMENT,13            session_id TEXT,14            role TEXT,15            content TEXT16        )17    ''')18    conn.commit()19    conn.close()20 21def save_message(session_id, role, content):22    """Logic: Stores a new message in the local file."""23    conn = sqlite3.connect(DB_PATH)24    cursor = conn.cursor()25    cursor.execute('INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)', 26                   (session_id, role, content))27    conn.commit()28    conn.close()29 30def load_messages(session_id):31    """Logic: Re-hydrates the chat history from the file."""32    conn = sqlite3.connect(DB_PATH)33    cursor = conn.cursor()34    cursor.execute('SELECT role, content FROM messages WHERE session_id = ? ORDER BY id ASC', (session_id,))35    rows = cursor.fetchall()36    conn.close()37    38    history = []39    for role, content in rows:40        if role == "user":41            history.append(HumanMessage(content=content))42        else:43            history.append(AIMessage(content=content))44    return history