Spaces:
Sleeping
Sleeping
File size: 3,162 Bytes
1d9f240 0b5d973 1d9f240 0b5d973 1d9f240 aea1c12 1d9f240 0b5d973 1d9f240 0b5d973 1d9f240 aea1c12 1d9f240 aea1c12 1d9f240 aea1c12 1d9f240 aea1c12 1d9f240 0b5d973 1d9f240 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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()
|