Update utils/database.py
Browse files- utils/database.py +28 -1
utils/database.py
CHANGED
|
@@ -67,7 +67,34 @@ def create_tables(conn):
|
|
| 67 |
conn.execute(sql_create_annotations_table)
|
| 68 |
except Error as e:
|
| 69 |
st.error(f"Error: {e}")
|
| 70 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 71 |
# QA system initialization function with necessary updates
|
| 72 |
def initialize_qa_system(vector_store):
|
| 73 |
"""Initialize QA system with proper chat handling."""
|
|
|
|
| 67 |
conn.execute(sql_create_annotations_table)
|
| 68 |
except Error as e:
|
| 69 |
st.error(f"Error: {e}")
|
| 70 |
+
# Add this function to your database.py file
|
| 71 |
+
|
| 72 |
+
def get_documents(conn):
|
| 73 |
+
"""Retrieve all documents from the database.
|
| 74 |
+
|
| 75 |
+
Args:
|
| 76 |
+
conn: SQLite database connection
|
| 77 |
+
|
| 78 |
+
Returns:
|
| 79 |
+
tuple: (list of document contents, list of document names)
|
| 80 |
+
"""
|
| 81 |
+
try:
|
| 82 |
+
cursor = conn.cursor()
|
| 83 |
+
cursor.execute("SELECT content, name FROM documents")
|
| 84 |
+
results = cursor.fetchall()
|
| 85 |
+
|
| 86 |
+
if not results:
|
| 87 |
+
return [], []
|
| 88 |
+
|
| 89 |
+
# Separate contents and names
|
| 90 |
+
document_contents = [row[0] for row in results]
|
| 91 |
+
document_names = [row[1] for row in results]
|
| 92 |
+
|
| 93 |
+
return document_contents, document_names
|
| 94 |
+
|
| 95 |
+
except Error as e:
|
| 96 |
+
st.error(f"Error retrieving documents: {e}")
|
| 97 |
+
return [], []
|
| 98 |
# QA system initialization function with necessary updates
|
| 99 |
def initialize_qa_system(vector_store):
|
| 100 |
"""Initialize QA system with proper chat handling."""
|