Update app.py

app.py

@@ -2,8 +2,9 @@ import os
 import streamlit as st
 from langchain_chroma import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.chains.question_answering import load_qa_chain
 from langchain.memory import ConversationBufferMemory
-from langchain.prompts import PromptTemplate
+from langchain_core.prompts import PromptTemplate
 from langchain_groq import ChatGroq
 from dotenv import load_dotenv
 from sentence_transformers import SentenceTransformer
@@ -22,7 +23,8 @@ if 'initialized' not in st.session_state:
     try:
         with st.spinner("Initializing..."):
             # Initialize embeddings model
-            model_path = "sentence-transformers/all-MiniLM-L12-v2"  
+            model_path = "sentence-transformers/all-MiniLM-L12-v2"  # Use a smaller, faster model
+            
             st.session_state.embedding_function = HuggingFaceEmbeddings(
                 model_name=model_path,
                 model_kwargs={'device': 'cpu'},
@@ -59,7 +61,7 @@ if 'initialized' not in st.session_state:
             # Load QA chain
             st.session_state.qa_chain = load_qa_chain(
                 llm=st.session_state.chat_model, 
-                chain_type="question_answering", 
+                chain_type="stuff", 
                 memory=st.session_state.memory, 
                 prompt=prompt
             )
@@ -73,12 +75,12 @@ if 'initialized' not in st.session_state:
 
 # Clear chat history buttons
 if st.button("Clear Chat History"):
-    if 'memory' in st.session_state and st.session_state.memory:
+    if 'memory' in st.session_state:
         st.session_state.memory.clear()
     st.experimental_rerun()  # Refresh the app to reflect the cleared history
 
 # Display chat history if initialized
-if st.session_state.initialized and 'memory' in st.session_state and st.session_state.memory:
+if st.session_state.initialized and 'memory' in st.session_state:
     if st.session_state.memory.buffer_as_messages:
         for message in st.session_state.memory.buffer_as_messages:
             if message.type == "ai":
@@ -87,21 +89,20 @@ if st.session_state.initialized and 'memory' in st.session_state and st.session_
                 st.chat_message(name="human", avatar="👤").write(message.content)
 
 # Input for new query
-if 'initialized' in st.session_state and st.session_state.initialized:
-    query = st.chat_input("Ask something")
-    if query:
-        try:
-            with st.spinner("Answering..."):
-                # Perform similarity search and get response
-                docs = st.session_state.docsearch.similarity_search(query, k=1)  
-                response = st.session_state.qa_chain(
-                    {"input_documents": docs, "human_input": query}, 
-                    return_only_outputs=True
-                )["output_text"]
+query = st.chat_input("Ask something")
+if query:
+    try:
+        with st.spinner("Answering..."):
+            # Perform similarity search and get response
+            docs = st.session_state.docsearch.similarity_search(query, k=1)  # Reduced k for speed
+            response = st.session_state.qa_chain(
+                {"input_documents": docs, "human_input": query}, 
+                return_only_outputs=True
+            )["output_text"]
 
-                # Display new message
-                st.chat_message(name="human", avatar="👤").write(query)
-                st.chat_message(name="ai", avatar="🤖").write(response)
+            # Display new message
+            st.chat_message(name="human", avatar="👤").write(query)
+            st.chat_message(name="ai", avatar="🤖").write(response)
 
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
+    except Exception as e:
+        st.error(f"An error occurred: {e}")