Update utils/database.py

utils/database.py

@@ -3,11 +3,14 @@ import streamlit as st
 import sqlite3
 from sqlite3 import Error
 from datetime import datetime
-from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
+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)
@@ -16,6 +19,7 @@ def create_connection(db_file):
         st.error(f"Error: {e}")
     return None
 
+
 def create_tables(conn):
     try:
         sql_create_documents_table = """
@@ -56,16 +60,20 @@ def create_tables(conn):
     except Error as e:
         st.error(f"Error: {e}")
 
+
 def get_documents(conn):
     """Retrieve documents from database"""
     try:
         cursor = conn.cursor()
-        cursor.execute("SELECT id, name, upload_date FROM documents ORDER BY upload_date DESC")
+        cursor.execute(
+            "SELECT id, name, upload_date FROM documents ORDER BY upload_date DESC"
+        )
         return cursor.fetchall()
     except Exception as e:
         st.error(f"Error retrieving documents: {e}")
         return []
 
+
 def insert_document(conn, doc_name, doc_content):
     """Insert a document into database"""
     try:
@@ -74,13 +82,14 @@ def insert_document(conn, doc_name, doc_content):
         if not cursor.fetchone():
             conn.execute(
                 "INSERT INTO documents (name, content) VALUES (?, ?)",
-                (doc_name, doc_content)
+                (doc_name, doc_content),
             )
             return True
     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:
@@ -92,7 +101,7 @@ def initialize_qa_system(vector_store):
 
         memory = ConversationBufferMemory(
             memory_key="chat_history",
-            return_messages=True
+            return_messages=True,
         )
 
         qa_chain = ConversationalRetrievalChain.from_llm(
@@ -100,7 +109,7 @@ def initialize_qa_system(vector_store):
             retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
             memory=memory,
             return_source_documents=True,
-            verbose=True
+            verbose=True,
         )
 
         return qa_chain
@@ -109,25 +118,29 @@ def initialize_qa_system(vector_store):
         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]
+            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
+
+@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