Update app.py

app.py

@@ -69,27 +69,28 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# 6. User Input Logic
-if prompt := st.chat_input("Ask me a complex math question!"):
-    # Add user message to history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
-
-    # Generate Agent Response
-    with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
+# --- 6. The Solve Action ---
+if st.button("🚀 Solve with AI", type="primary", use_container_width=True):
+    if st.session_state.equation:
+        user_input = st.session_state.equation
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        with st.chat_message("assistant"):
+            # --- START OF TRY BLOCK ---
             try:
-                # ALL lines below 'try' must be indented 4 spaces further
-                # We pass the system_message rules directly here
+                # This is the line that actually "runs" the math
                 response = agent_executor.invoke({
-                    "messages": [system_message, ("user", prompt)]
+                    "messages": [system_message, ("user", user_input)]
                 })
-                
-                final_answer = response["messages"][-1].content
-                
-                st.markdown(final_answer)
-                st.session_state.messages.append({"role": "assistant", "content": final_answer})
+                answer = response["messages"][-1].content
+                st.markdown(answer)
+                st.session_state.messages.append({"role": "assistant", "content": answer})
+                st.session_state.equation = "" 
+
+            except ZeroDivisionError:
+                st.error("🚫 Mathematicians have determined that division by zero is undefined! Try a different calculation.")
             
             except Exception as e:
-                st.error(f"An error occurred: {e}")
+                # This catches other errors like "Unknown format" or "Timeout"
+                st.warning(f"⚠️ I ran into a bit of a snag: {e}")
+            # --- END OF TRY BLOCK ---