Spaces:
Sleeping
Sleeping
File size: 4,152 Bytes
81f3119 07768c8 81f3119 07768c8 81f3119 0e0e663 81f3119 0e0e663 81f3119 0e0e663 81f3119 | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | import sqlite3
import numpy as np
import pickle
import os
DB_NAME = os.getenv("DB_NAME", "documents.db")
def get_connection():
conn = sqlite3.connect(DB_NAME, timeout=30)
conn.execute("PRAGMA journal_mode=WAL;")
conn.execute("PRAGMA busy_timeout=30000;")
return conn
def create_tables():
conn = get_connection()
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS document_chunks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT,
page INTEGER,
chunk_id INTEGER,
content TEXT,
embedding BLOB
)
""")
cursor.execute("""
CREATE TABLE IF NOT EXISTS query_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT,
question TEXT,
answer TEXT,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
)
""")
cursor.execute("""
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT UNIQUE,
password TEXT,
role TEXT
)
""")
cursor.execute("""
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT UNIQUE,
uploaded_by TEXT,
uploaded_at DATETIME DEFAULT CURRENT_TIMESTAMP
)
""")
conn.commit()
conn.close()
def insert_chunk(chunk, embedding):
conn = get_connection()
cursor = conn.cursor()
embedding_blob = pickle.dumps(embedding)
cursor.execute("""
INSERT INTO document_chunks
(filename, page, chunk_id, content, embedding)
VALUES (?, ?, ?, ?, ?)
""", (
chunk["filename"],
chunk["page"],
chunk["chunk_id"],
chunk["content"],
embedding_blob
))
conn.commit()
conn.close()
def insert_chunks_bulk(chunks, embeddings):
conn = get_connection()
cursor = conn.cursor()
rows = []
for chunk, embedding in zip(chunks, embeddings):
rows.append(
(
chunk["filename"],
chunk["page"],
chunk["chunk_id"],
chunk["content"],
pickle.dumps(embedding),
)
)
cursor.executemany(
"""
INSERT INTO document_chunks
(filename, page, chunk_id, content, embedding)
VALUES (?, ?, ?, ?, ?)
""",
rows,
)
conn.commit()
conn.close()
def fetch_all_chunks():
conn = get_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT filename, page, chunk_id, content, embedding
FROM document_chunks
""")
rows = cursor.fetchall()
conn.close()
chunks = []
embeddings = []
for row in rows:
chunks.append({
"filename": row[0],
"page": row[1],
"chunk_id": row[2],
"content": row[3]
})
embeddings.append(pickle.loads(row[4]))
return chunks, np.array(embeddings)
def document_exists(filename):
conn = get_connection()
cursor = conn.cursor()
cursor.execute(
"SELECT COUNT(*) FROM documents WHERE filename = ?",
(filename,)
)
exists = cursor.fetchone()[0] > 0
conn.close()
return exists
def insert_query(username, question, answer):
conn = get_connection()
cursor = conn.cursor()
cursor.execute(
"INSERT INTO query_history (username, question, answer) VALUES (?, ?, ?)",
(username, question, answer[:500])
)
conn.commit()
conn.close()
def fetch_query_history(limit=10):
conn = get_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT question, answer, timestamp
FROM query_history
ORDER BY timestamp DESC
LIMIT ?
""", (limit,))
rows = cursor.fetchall()
conn.close()
return rows
def clear_query_history():
conn = get_connection()
cursor = conn.cursor()
cursor.execute("DELETE FROM query_history")
conn.commit()
conn.close()
|