Update app.py

app.py

@@ -57,8 +57,8 @@ def setup_rag_pipeline():
     return retriever
 
 def reset_conversation():
-    st.session_state.conversation = []
-    st.session_state.messages = []
+    st.session_state.visible_messages = []
+    st.session_state.full_context = []
 
 def main():
     st.header('Multi-Models with RAG')
@@ -73,9 +73,9 @@ def main():
         st.session_state.prev_option = selected_model
 
     if st.session_state.prev_option != selected_model:
-        st.session_state.messages = []
+        st.session_state.visible_messages = []
+        st.session_state.full_context = []
         st.session_state.prev_option = selected_model
-        reset_conversation()
 
     st.markdown(f'_powered_ by ***:violet[{selected_model}]***')
 
@@ -84,11 +84,13 @@ def main():
     st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
     # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
+    if "visible_messages" not in st.session_state:
+        st.session_state.visible_messages = []
+    if "full_context" not in st.session_state:
+        st.session_state.full_context = []
 
     # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
+    for message in st.session_state.visible_messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
@@ -103,27 +105,28 @@ def process_user_input(client, prompt, selected_model, temperature, retriever):
     # Display user message
     with st.chat_message("user"):
         st.markdown(prompt)
+    st.session_state.visible_messages.append({"role": "user", "content": prompt})
     
     # Retrieve relevant documents
     relevant_docs = retriever.get_relevant_documents(prompt)
     context = "\n".join([doc.page_content for doc in relevant_docs])
     
-    # Prepare messages with context
-    messages = [
+    # Prepare full context with system message and retrieved context
+    full_context = [
         {"role": "system", "content": f"You are an AI assistant. Use the following context to answer the user's question: {context}"},
+        *st.session_state.full_context,
         {"role": "user", "content": prompt}
     ]
-    st.session_state.messages.extend(messages)
+    
+    # Update full context in session state
+    st.session_state.full_context = full_context
 
     # Generate and display assistant response
     with st.chat_message("assistant"):
         try:
             stream = client.chat.completions.create(
                 model=model_links[selected_model],
-                messages=[
-                    {"role": m["role"], "content": m["content"]}
-                    for m in st.session_state.messages
-                ],
+                messages=full_context,
                 temperature=temperature,
                 stream=True,
                 max_tokens=MAX_TOKENS,
@@ -133,7 +136,9 @@ def process_user_input(client, prompt, selected_model, temperature, retriever):
             handle_error(e)
             return
 
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    # Update visible messages and full context
+    st.session_state.visible_messages.append({"role": "assistant", "content": response})
+    st.session_state.full_context.append({"role": "assistant", "content": response})
 
 def handle_error(error):
     response = """😵‍💫 Looks like someone unplugged something!