Spaces:
Sleeping
Sleeping
File size: 3,275 Bytes
1f01f87 | 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 | import sqlite3
import json
from datetime import datetime
class DatabaseManager:
def __init__(self, db_path="chat_history.db"):
self.db_path = db_path
self.init_db()
def init_db(self):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
# Create documents table
cursor.execute("""
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT NOT NULL,
path TEXT NOT NULL,
upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
# Create chat history table
cursor.execute("""
CREATE TABLE IF NOT EXISTS chat_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
role TEXT NOT NULL,
content TEXT NOT NULL,
document_id INTEGER,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (document_id) REFERENCES documents(id)
)
""")
conn.commit()
def add_document(self, filename: str, path: str) -> int:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO documents (filename, path) VALUES (?, ?)",
(filename, path)
)
conn.commit()
return cursor.lastrowid
def get_documents(self):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("SELECT * FROM documents ORDER BY upload_date DESC")
return cursor.fetchall()
def add_chat_message(self, role: str, content: str, document_id: int):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO chat_history (role, content, document_id) VALUES (?, ?, ?)",
(role, content, document_id)
)
conn.commit()
def get_chat_history(self, document_id: int):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute(
"SELECT role, content FROM chat_history WHERE document_id = ? ORDER BY timestamp",
(document_id,)
)
return [{
"role": role,
"content": content
} for role, content in cursor.fetchall()]
def clear_chat_history(self, document_id: int):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("DELETE FROM chat_history WHERE document_id = ?", (document_id,))
conn.commit()
def clear_all_documents(self):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
# Delete all chat history first due to foreign key constraint
cursor.execute("DELETE FROM chat_history")
# Then delete all documents
cursor.execute("DELETE FROM documents")
conn.commit() |