Update utils/database.py

utils/database.py

@@ -1,7 +1,60 @@
-# src/utils/database.py
+# utils/database.py
 import streamlit as st
 import sqlite3
+from sqlite3 import Error
 from datetime import datetime
+from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from langchain.chat_models import ChatOpenAI
+import os
+
+def create_connection(db_file):
+    try:
+        conn = sqlite3.connect(db_file)
+        return conn
+    except Error as e:
+        st.error(f"Error: {e}")
+    return None
+
+def create_tables(conn):
+    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)
+        );
+        """
+
+        c = conn.cursor()
+        c.execute(sql_create_documents_table)
+        c.execute(sql_create_queries_table)
+        c.execute(sql_create_annotations_table)
+    except Error as e:
+        st.error(f"Error: {e}")
 
 def get_documents(conn):
     """Retrieve documents from database"""
@@ -27,3 +80,57 @@ def insert_document(conn, doc_name, doc_content):
     except Exception as e:
         st.error(f"Error inserting document: {e}")
     return False
+
+def initialize_qa_system(vector_store):
+    """Initialize QA system with proper chat handling"""
+    try:
+        llm = ChatOpenAI(
+            temperature=0,
+            model_name="gpt-4",
+            api_key=os.environ.get("OPENAI_API_KEY"),
+        )
+
+        memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True
+        )
+
+        qa_chain = ConversationalRetrievalChain.from_llm(
+            llm=llm,
+            retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
+            memory=memory,
+            return_source_documents=True,
+            verbose=True
+        )
+
+        return qa_chain
+
+    except Exception as e:
+        st.error(f"Error initializing QA system: {e}")
+        return None
+
+def initialize_faiss(embeddings, documents, document_names):
+    """Initialize FAISS vector store"""
+    try:
+        from langchain.vectorstores import FAISS
+        vector_store = FAISS.from_texts(
+            documents,
+            embeddings,
+            metadatas=[{"source": name} for name in document_names]
+        )
+        return vector_store
+    except Exception as e:
+        st.error(f"Error initializing FAISS: {e}")
+        return None
+
+        @st.cache_resource
+def get_embeddings_model():
+    """Get the embeddings model"""
+    try:
+        from langchain.embeddings import HuggingFaceEmbeddings
+        model_name = "sentence-transformers/all-MiniLM-L6-v2"
+        embeddings = HuggingFaceEmbeddings(model_name=model_name)
+        return embeddings
+    except Exception as e:
+        st.error(f"Error loading embeddings model: {e}")
+        return None