Update utils/database.py

utils/database.py

@@ -1,8 +1,8 @@
 # utils/database.py
 from langchain.memory import ConversationBufferWindowMemory
 from langchain_core.messages import (
-    HumanMessage, 
-    AIMessage, 
+    HumanMessage,
+    AIMessage,
     SystemMessage,
     BaseMessage  # Added this import
 )
@@ -15,9 +15,62 @@ from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputP
 from langchain_core.runnables import RunnablePassthrough
 import os
 import streamlit as st
+import sqlite3
+from sqlite3 import Error
 
+# Original SQL-related functions preserved
+def create_connection(db_file):
+    """Create a database connection to the SQLite database."""
+    try:
+        conn = sqlite3.connect(db_file)
+        return conn
+    except Error as e:
+        st.error(f"Error: {e}")
+    return None
+
+def create_tables(conn):
+    """Create necessary tables in the database."""
+    try:
+        sql_create_documents_table = '''
+        CREATE TABLE IF NOT EXISTS documents (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            name TEXT NOT NULL,
+            content TEXT NOT NULL,
+            upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        );
+        '''
+
+        sql_create_queries_table = '''
+        CREATE TABLE IF NOT EXISTS queries (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            query TEXT NOT NULL,
+            response TEXT NOT NULL,
+            document_id INTEGER,
+            query_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            FOREIGN KEY (document_id) REFERENCES documents (id)
+        );
+        '''
+
+        sql_create_annotations_table = '''
+        CREATE TABLE IF NOT EXISTS annotations (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            document_id INTEGER NOT NULL,
+            annotation TEXT NOT NULL,
+            page_number INTEGER,
+            annotation_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            FOREIGN KEY (document_id) REFERENCES documents (id)
+        );
+        '''
+
+        conn.execute(sql_create_documents_table)
+        conn.execute(sql_create_queries_table)
+        conn.execute(sql_create_annotations_table)
+    except Error as e:
+        st.error(f"Error: {e}")
+
+# QA system initialization function with necessary updates
 def initialize_qa_system(vector_store):
-    """Initialize QA system with proper chat handling"""
+    """Initialize QA system with proper chat handling."""
     try:
         llm = ChatOpenAI(
             temperature=0.5,
@@ -77,9 +130,10 @@ def initialize_qa_system(vector_store):
     except Exception as e:
         st.error(f"Error initializing QA system: {e}")
         return None
-        
+
+# FAISS vector store initialization
 def initialize_faiss(embeddings, documents, document_names):
-    """Initialize FAISS vector store"""
+    """Initialize FAISS vector store."""
     try:
         from langchain.vectorstores import FAISS
 
@@ -93,10 +147,10 @@ def initialize_faiss(embeddings, documents, document_names):
         st.error(f"Error initializing FAISS: {e}")
         return None
 
-
+# Embeddings model retrieval
 @st.cache_resource
 def get_embeddings_model():
-    """Get the embeddings model"""
+    """Get the embeddings model."""
     try:
         from langchain.embeddings import HuggingFaceEmbeddings
 
@@ -105,4 +159,4 @@ def get_embeddings_model():
         return embeddings
     except Exception as e:
         st.error(f"Error loading embeddings model: {e}")
-        return None
+        return None