Update app.py

app.py

@@ -6,7 +6,7 @@ from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTem
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain.chat_models import ChatOpenAI
-from langchain.schema import SystemMessage, HumanMessage
+from langchain.schema import SystemMessage, HumanMessage, AIMessage
 from PyPDF2 import PdfReader
 import aiohttp
 from io import BytesIO
@@ -15,7 +15,7 @@ from io import BytesIO
 os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
 
 # Set up prompts
-system_template = "Use the following context to answer a user's question. If you cannot find the answer in the context, say you don't know the answer."
+system_template = "You are an AI assistant answering questions about AI. Use the following context to answer the user's question. If you cannot find the answer in the context, say you don't know the answer but you can try to help with related information."
 system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)
 
 human_template = "Context:\n{context}\n\nQuestion:\n{question}"
@@ -29,7 +29,7 @@ class RetrievalAugmentedQAPipeline:
         self.llm = llm
         self.vector_db = vector_db
 
-    async def arun_pipeline(self, user_query: str):
+    async def arun_pipeline(self, user_query: str, chat_history: list):
         context_docs = self.vector_db.similarity_search(user_query, k=2)
         context_list = [doc.page_content for doc in context_docs]
         context_prompt = "\n".join(context_list)
@@ -38,7 +38,9 @@ class RetrievalAugmentedQAPipeline:
         if len(context_prompt) > max_context_length:
             context_prompt = context_prompt[:max_context_length]
         
-        messages = chat_prompt.format_prompt(context=context_prompt, question=user_query).to_messages()
+        messages = [SystemMessage(content=system_template)]
+        messages.extend(chat_history)
+        messages.append(HumanMessage(content=human_template.format(context=context_prompt, question=user_query)))
 
         response = await self.llm.agenerate([messages])
         return {"response": response.generations[0][0].text}
@@ -86,13 +88,36 @@ async def main():
 # Streamlit UI
 st.title("Ask About AI!")
 
+# Initialize session state for chat history
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
 pipeline = initialize_pipeline()
 
+# Display chat history
+for message in st.session_state.chat_history:
+    if isinstance(message, HumanMessage):
+        st.write("You:", message.content)
+    elif isinstance(message, AIMessage):
+        st.write("AI:", message.content)
+
 user_query = st.text_input("Enter your question about AI:")
 
 if user_query:
+    # Add user message to chat history
+    st.session_state.chat_history.append(HumanMessage(content=user_query))
+    
     with st.spinner("Generating response..."):
-        result = asyncio.run(pipeline.arun_pipeline(user_query))
+        result = asyncio.run(pipeline.arun_pipeline(user_query, st.session_state.chat_history))
     
-    st.write("Response:")
-    st.write(result["response"])
+    # Add AI response to chat history
+    ai_message = AIMessage(content=result["response"])
+    st.session_state.chat_history.append(ai_message)
+    
+    # Display the latest response
+    st.write("AI:", result["response"])
+
+# Add a button to clear chat history
+if st.button("Clear Chat History"):
+    st.session_state.chat_history = []
+    st.experimental_rerun()