Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import random
-from app_config import SYSTEM_PROMPT, NLP_MODEL_NAME, NUMBER_OF_VECTORS_FOR_RAG, NLP_MODEL_TEMPERATURE, NLP_MODEL_MAX_TOKENS, VECTOR_MAX_TOKENS,my_vector_store,chat,tiktoken_len
+from app_config import SYSTEM_PROMPT, NLP_MODEL_NAME, NUMBER_OF_VECTORS_FOR_RAG, NLP_MODEL_TEMPERATURE, NLP_MODEL_MAX_TOKENS, VECTOR_MAX_TOKENS, my_vector_store, chat, tiktoken_len
 from langchain.memory import ConversationSummaryBufferMemory
 from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain.chains.summarize import load_summarize_chain
@@ -9,37 +9,29 @@ from langchain_groq import ChatGroq
 from dotenv import load_dotenv
 from pathlib import Path
 import os
+
 env_path = Path('.') / '.env'
 load_dotenv(dotenv_path=env_path)
 
-def response_generator(prompt: str) -> str:
-    """this function can be used for general quetion answers which are related to tyrex and tyre recycling 
-
-    Args:
-        prompt (string): user query
-
-    Returns:
-        string: answer of the query
-    """
+# Initialize vector store and LLM outside session state
+retriever = my_vector_store.as_retriever(k=NUMBER_OF_VECTORS_FOR_RAG)
+llm = ChatGroq(temperature=NLP_MODEL_TEMPERATURE, groq_api_key=str(os.getenv('GROQ_API_KEY')), model_name=NLP_MODEL_NAME)
 
+def response_generator(prompt: str) -> str:
     try:
-        retriever = st.session_state.retriever
         docs = retriever.invoke(prompt)
         my_context = [doc.page_content for doc in docs]
         my_context = '\n\n'.join(my_context)
-        system_message = SystemMessage(content = SYSTEM_PROMPT.format(context=my_context, previous_message_summary=st.session_state.rag_memory.moving_summary_buffer))
+        system_message = SystemMessage(content=SYSTEM_PROMPT.format(context=my_context, previous_message_summary=st.session_state.rag_memory.moving_summary_buffer))
         print(system_message)
         chat_messages = (system_message + st.session_state.rag_memory.chat_memory.messages + HumanMessage(content=prompt)).messages
         print("total tokens: ", tiktoken_len(str(chat_messages)))
-        # print("my_context*********",my_context)
-        response = st.session_state.llm.invoke(chat_messages)
+        response = llm.invoke(chat_messages)
         return response.content
-    
     except Exception as error:
         print(error, "ERROR")
         return "Oops! something went wrong, please try again."
 
-
 st.markdown(
     """
 <style>
@@ -52,30 +44,16 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
-# When user gives input
-
-print("SYSTEM MESSAGE")
+# Initialize session state
 if "messages" not in st.session_state:
-    st.session_state.messages=[{"role": "system", "content": SYSTEM_PROMPT}]
-    
-print("SYSTEM MODEL")
-if "llm" not in st.session_state:
-    st.session_state.llm = ChatGroq(temperature=NLP_MODEL_TEMPERATURE, groq_api_key=str(os.getenv('GROQ_API_KEY')), model_name=NLP_MODEL_NAME)
-
-print("rag")
+    st.session_state.messages = [{"role": "system", "content": SYSTEM_PROMPT}]
 if "rag_memory" not in st.session_state:
-    st.session_state.rag_memory = ConversationSummaryBufferMemory(llm=st.session_state.llm, max_token_limit= 5000)
-
-print("retrival")
+    st.session_state.rag_memory = ConversationSummaryBufferMemory(llm=llm, max_token_limit=5000)
 if "retriever" not in st.session_state:
-    # vector_store = get_vectorstore_with_doc_from_pdf('GPT OUTPUT.pdf')
-    st.session_state.retriever = my_vector_store.as_retriever(k=NUMBER_OF_VECTORS_FOR_RAG)
-
+    st.session_state.retriever = retriever
 
 st.title("Insurance Bot")
-print("container")
-# Display chat messages from history
-container =  st.container(height=600)
+container = st.container(height=600)
 for message in st.session_state.messages:
     if message["role"] != "system":
         with container.chat_message(message["role"]):
@@ -84,7 +62,7 @@ for message in st.session_state.messages:
 if prompt := st.chat_input("Enter your query here... "):
     with container.chat_message("user"):
         st.write(prompt)
-    st.session_state.messages.append({"role":"user" , "content":prompt})
+    st.session_state.messages.append({"role": "user", "content": prompt})
     
     with container.chat_message("assistant"):  
         response = response_generator(prompt=prompt)
@@ -94,4 +72,4 @@ if prompt := st.chat_input("Enter your query here... "):
     
     print("Response is:", response)
     st.session_state.rag_memory.save_context({'input': prompt}, {'output': response})
-    st.session_state.messages.append({"role":"assistant" , "content":response})
+    st.session_state.messages.append({"role": "assistant", "content": response})