backend/db_utils.py

import sqlite3
from datetime import datetime
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage


DB_NAME = "research_assistant.db"

def get_db_connection():
    conn = sqlite3.connect(DB_NAME)
    conn.row_factory = sqlite3.Row
    return conn

def create_application_logs():
    conn = get_db_connection()
    conn.execute('''CREATE TABLE IF NOT EXISTS application_logs
                    (id INTEGER PRIMARY KEY AUTOINCREMENT,
                     session_id TEXT,
                     user_query TEXT,
                     gpt_response TEXT,
                     model TEXT,
                     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
    conn.close()

def create_document_store():
    conn = get_db_connection()
    conn.execute('''CREATE TABLE IF NOT EXISTS document_store
                    (id INTEGER PRIMARY KEY AUTOINCREMENT,
                     session_id TEXT,
                     filename TEXT,
                     content TEXT,
                     upload_timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
    conn.close()

def insert_application_logs(session_id, user_query, gpt_response, model):
    conn = get_db_connection()
    conn.execute('INSERT INTO application_logs (session_id, user_query, gpt_response, model) VALUES (?, ?, ?, ?)',
                 (session_id, user_query, gpt_response, model))
    conn.commit()
    conn.close()

def get_chat_history(session_id):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('SELECT user_query, gpt_response FROM application_logs WHERE session_id = ? ORDER BY created_at', (session_id,))
    messages = []
    for row in cursor.fetchall():
        messages.append(HumanMessage(content=row['user_query']))
        messages.append(AIMessage(content=row['gpt_response']))

    conn.close()
    return messages
def insert_document_record(session_id, filename, content):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('INSERT INTO document_store (session_id, filename, content) VALUES (?, ?, ?)', 
                   (session_id, filename, content))
    file_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return file_id

def delete_document_record(file_id):
    conn = get_db_connection()
    conn.execute('DELETE FROM document_store WHERE id = ?', (file_id,))
    conn.commit()
    conn.close()
    return True

def get_all_documents(session_id):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('SELECT id, filename, upload_timestamp FROM document_store WHERE session_id = ? ORDER BY upload_timestamp DESC', (session_id,))
    documents = cursor.fetchall()
    conn.close()
    return [dict(doc) for doc in documents]


def get_file_content(file_id: int) -> str | None:
    conn = get_db_connection()
    try:
        cursor = conn.cursor()
        cursor.execute('SELECT content FROM document_store WHERE id = ?', (file_id,))
        row = cursor.fetchone()
        if row is not None:
            return row[0]  
        return None
    finally:
        conn.close()


# Initialize the database tables
create_application_logs()
create_document_store()