VishnuRamDebyez commited on
Commit
df5c110
·
verified ·
1 Parent(s): 929d2ae

Create db_utils.py

Browse files
Files changed (1) hide show
  1. db_utils.py +76 -0
db_utils.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sqlite3
2
+ from datetime import datetime
3
+
4
+ DB_NAME = "rag_app.db"
5
+
6
+ def get_db_connection():
7
+ conn = sqlite3.connect(DB_NAME)
8
+ conn.row_factory = sqlite3.Row
9
+ return conn
10
+
11
+ def create_application_logs():
12
+ conn = get_db_connection()
13
+ conn.execute('''CREATE TABLE IF NOT EXISTS application_logs
14
+ (id INTEGER PRIMARY KEY AUTOINCREMENT,
15
+ session_id TEXT,
16
+ user_query TEXT,
17
+ gpt_response TEXT,
18
+ model TEXT,
19
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
20
+ conn.close()
21
+
22
+ def insert_application_logs(session_id, user_query, gpt_response, model):
23
+ conn = get_db_connection()
24
+ conn.execute('INSERT INTO application_logs (session_id, user_query, gpt_response, model) VALUES (?, ?, ?, ?)',
25
+ (session_id, user_query, gpt_response, model))
26
+ conn.commit()
27
+ conn.close()
28
+
29
+ def get_chat_history(session_id):
30
+ conn = get_db_connection()
31
+ cursor = conn.cursor()
32
+ cursor.execute('SELECT user_query, gpt_response FROM application_logs WHERE session_id = ? ORDER BY created_at', (session_id,))
33
+ messages = []
34
+ for row in cursor.fetchall():
35
+ messages.extend([
36
+ {"role": "human", "content": row['user_query']},
37
+ {"role": "ai", "content": row['gpt_response']}
38
+ ])
39
+ conn.close()
40
+ return messages
41
+
42
+ def create_document_store():
43
+ conn = get_db_connection()
44
+ conn.execute('''CREATE TABLE IF NOT EXISTS document_store
45
+ (id INTEGER PRIMARY KEY AUTOINCREMENT,
46
+ filename TEXT,
47
+ upload_timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
48
+ conn.close()
49
+
50
+ def insert_document_record(filename):
51
+ conn = get_db_connection()
52
+ cursor = conn.cursor()
53
+ cursor.execute('INSERT INTO document_store (filename) VALUES (?)', (filename,))
54
+ file_id = cursor.lastrowid
55
+ conn.commit()
56
+ conn.close()
57
+ return file_id
58
+
59
+ def delete_document_record(file_id):
60
+ conn = get_db_connection()
61
+ conn.execute('DELETE FROM document_store WHERE id = ?', (file_id,))
62
+ conn.commit()
63
+ conn.close()
64
+ return True
65
+
66
+ def get_all_documents():
67
+ conn = get_db_connection()
68
+ cursor = conn.cursor()
69
+ cursor.execute('SELECT id, filename, upload_timestamp FROM document_store ORDER BY upload_timestamp DESC')
70
+ documents = cursor.fetchall()
71
+ conn.close()
72
+ return [dict(doc) for doc in documents]
73
+
74
+ # Initialize the database tables
75
+ create_application_logs()
76
+ create_document_store()