Update app.py

app.py

@@ -1,12 +1,15 @@
 import streamlit as st
 from langchain_openai import ChatOpenAI
+import os
+import dotenv
 from langchain_community.document_loaders import WebBaseLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
 from langchain_openai import OpenAIEmbeddings
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain.memory import ConversationBufferMemory
 
 # Set page config
 st.set_page_config(page_title="Tbank Assistant", layout="wide")
@@ -18,11 +21,16 @@ st.title("Tbank Customer Support Chatbot")
 with st.sidebar:
     st.header("Configuration")
     api_key = st.text_input("Enter your OpenAI API Key:", type="password")
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
 
 # Main app logic
-@st.cache_resource
-def initialize_components():
-        chat = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0.2,api_key=api_key)
+if "OPENAI_API_KEY" in os.environ:
+    # Initialize components
+    @st.cache_resource
+    def initialize_components():
+        dotenv.load_dotenv()
+        chat = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0.2)
 
         loader = WebBaseLoader("https://www.tbankltd.com/about-us")
         data = loader.load()
@@ -51,6 +59,9 @@ def initialize_components():
         <context>
         {context}
         </context>
+
+        Chat History:
+        {chat_history}
         """
 
         question_answering_prompt = ChatPromptTemplate.from_messages(
@@ -59,63 +70,68 @@ def initialize_components():
                     "system",
                     SYSTEM_TEMPLATE,
                 ),
+                MessagesPlaceholder(variable_name="chat_history"),
                 MessagesPlaceholder(variable_name="messages"),
             ]
         )
 
+        memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
         document_chain = create_stuff_documents_chain(chat, question_answering_prompt)
 
-        return retriever, document_chain
-
+        return retriever, document_chain, memory
 
     # Load components
-with st.spinner("Initializing Tbank Assistant..."):
-    retriever, document_chain = initialize_components()
-
-# Chat interface
-st.subheader("Chat with Tbank Assistant")
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# Display chat messages from history on app rerun
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# React to user input
-# React to user input
-if prompt := st.chat_input("What would you like to know about Tbank?"):
-    # Display user message in chat message container
-    st.chat_message("user").markdown(prompt)
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-
-    with st.chat_message("assistant"):
-        message_placeholder = st.empty()
-
-        # Retrieve relevant documents
-        docs = retriever.get_relevant_documents(prompt)
-
-        # Generate response
-        response = document_chain.invoke(
-            {
-                "context": docs,
-                "messages": [
-                    HumanMessage(content=prompt)
-                ],
-            }
-        )
-
-        # The response is already a string, so we can use it directly
-        full_response = response
-        message_placeholder.markdown(full_response)
-
-    # Add assistant response to chat history
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
-
-
+    with st.spinner("Initializing Tbank Assistant..."):
+        retriever, document_chain, memory = initialize_components()
+
+    # Chat interface
+    st.subheader("Chat with Tbank Assistant")
+
+    # Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display chat messages from history on app rerun
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # React to user input
+    if prompt := st.chat_input("What would you like to know about Tbank?"):
+        # Display user message in chat message container
+        st.chat_message("user").markdown(prompt)
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        with st.chat_message("assistant"):
+            message_placeholder = st.empty()
+
+            # Retrieve relevant documents
+            docs = retriever.get_relevant_documents(prompt)
+
+            # Generate response
+            response = document_chain.invoke(
+                {
+                    "context": docs,
+                    "chat_history": memory.load_memory_variables({})["chat_history"],
+                    "messages": [
+                        HumanMessage(content=prompt)
+                    ],
+                }
+            )
+
+            # The response is already a string, so we can use it directly
+            full_response = response
+            message_placeholder.markdown(full_response)
+
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+        # Update memory
+        memory.save_context({"input": prompt}, {"output": full_response})
+
+else:
+    st.warning("Please enter your OpenAI API Key in the sidebar to start the chatbot.")
 
 # Add a footer
 st.markdown("---")