Update app.py

app.py

@@ -242,26 +242,40 @@ with st.form("chat_form", clear_on_submit=True):
 if st.session_state.chat_history and st.session_state.chat_history[-1][0] == "user":
     inputs = {"messages": st.session_state.chat_history}
 
-    # Create a container for streaming output
-    response_placeholder = st.empty()
-    assistant_message = ""
+    # Create separate placeholders for tool output and final answer.
+    tool_output_placeholder = st.empty()
+    final_answer_placeholder = st.empty()
+
+    # Accumulators for each section.
+    tool_output_text = ""
+    final_answer_text = ""
 
     # Stream the agent's response chunk-by-chunk.
     for s in graph.stream(inputs, stream_mode="values"):
-        # Extract the last message from the messages list.
+        # Get the last message in the chunk.
         message = s["messages"][-1]
+        
         if isinstance(message, tuple):
+            # This is a tool output message (e.g., retrieved docs or tool usage).
             role, text = message
+            tool_output_text += text
+            # Update the tool output area with a heading and the accumulated text.
+            tool_output_placeholder.markdown(
+                f"### Tool Output\n\n{tool_output_text}",
+                unsafe_allow_html=True
+            )
         else:
+            # This is the final answer generated by the AI.
             text = message.content
+            final_answer_text += text
+            # Update the final answer area with a heading and the accumulated text.
+            final_answer_placeholder.markdown(
+                f"### Final Answer\n\n{final_answer_text}",
+                unsafe_allow_html=True
+            )
 
-        # Append new text and update the markdown display.
-        assistant_message += text
-
-        # Update the placeholder with the current assistant message.
-        # Using unsafe_allow_html=True ensures that any markdown (or embedded HTML) is rendered.
-        response_placeholder.markdown(assistant_message, unsafe_allow_html=True)
+    # Once complete, combine both sections into the chat history.
+    combined_response = f"**Tool Output:**\n\n{tool_output_text}\n\n**Final Answer:**\n\n{final_answer_text}"
+    st.session_state.chat_history.append(("assistant", combined_response))
 
-    # Once complete, add the full assistant response to the chat history.
-    st.session_state.chat_history.append(("assistant", assistant_message))