Update app.py

app.py

@@ -215,7 +215,12 @@ if "chat_history" not in st.session_state:
 for msg in st.session_state.chat_history:
     role = "user" if isinstance(msg, HumanMessage) else "bot"
     st.markdown(f"<div class='stChatMessage {role}'>{msg.content}</div>", unsafe_allow_html=True)
- 
+
+# Show explanation button if code exists
+if "latest_explanation" in st.session_state:
+    with st.expander("🔍 Show Code Explanation"):
+        st.markdown(f"<div class='stChatMessage bot'>**🔍 Code Explanation:**<br><br>{st.session_state.latest_explanation}</div>", unsafe_allow_html=True)
+
 # Custom Chat Input Box UI
 with st.container():
     with st.form("chat_form", clear_on_submit=True):
@@ -235,10 +240,19 @@ with st.container():
             }
         
             result = graph.invoke(state_input, config=config)
+            # if result.get("code"):
+            #     st.session_state.chat_history.append(AIMessage(content=result["code"]))
+            # elif result.get("explanation"):
+            #     st.session_state.chat_history.append(AIMessage(content=result["explanation"]))
+            # elif result.get("messages"):
+            #     st.session_state.chat_history.append(result["messages"][-1])
             if result.get("code"):
+                st.session_state.latest_code = result["code"]
                 st.session_state.chat_history.append(AIMessage(content=result["code"]))
-            elif result.get("explanation"):
-                st.session_state.chat_history.append(AIMessage(content=result["explanation"]))
+            
+            if result.get("explanation"):
+                st.session_state.latest_explanation = result["explanation"]
+            
             elif result.get("messages"):
                 st.session_state.chat_history.append(result["messages"][-1])