fix: chat_history export

pages/chat.py

@@ -18,40 +18,38 @@ if col2.button("Home"):
 universal_chat = st.toggle("Universal Chat")
 st.caption("Note: Universal Chat uses the complete DB to retrieve context, use it with caution")
 
-st.download_button(
-    label="Export",
-    data=st.session_state.encoded_text,
-    file_name='chat_history.md',
-    mime='text/markdown',
-)
-
 st.divider()
 
-# 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"])
 
 if prompt := st.chat_input("Ask Drake your questions"):
-    # Display user message in chat message container
+
     with st.chat_message("user"):
         st.markdown(prompt)
-    # Add user message to chat history
+
     st.session_state.messages.append({"role": "user", "content": prompt})
 
     with st.spinner("Drake is thinking..."):
         query = f"{prompt}"
-        _, drake = initialize_models()
+        _, drake = initialize_models() # Check resources in cache
         if universal_chat:
             response = drake.ask_llm(query)
         else:
             response = drake.ask_llm(query, metadata_filter=st.session_state["metadata"])
 
-        st.session_state.chat_notes += query + "\n" + response + "\n\n"
-        st.session_state.encoded_text = st.session_state.chat_notes.encode('utf-8')
-
         with st.chat_message("assistant"):
+            st.session_state.chat_notes += query + "\n" + response + "\n\n"
+            st.session_state.encoded_text = st.session_state.chat_notes.encode('utf-8')
+
             st.markdown(response)
             st.session_state.messages.append({"role": "assistant", "content": response})
 
+st.download_button(
+    label="Export",
+    data=st.session_state.encoded_text,
+    file_name='chat_history.md',
+    mime='text/markdown',
+)