Update app.py

app.py

@@ -4,21 +4,57 @@ from langchain_groq import ChatGroq
 from langchain_community.agent_toolkits.load_tools import load_tools
 from langgraph.prebuilt import create_react_agent
 
-# 1. Page Config (Best practice to put this at the very top)
+# 1. Page Configuration
 st.set_page_config(page_title="AI Math Agent", page_icon="🧮")
-
 st.title("My First AI Math Agent 🧮")
 
 # 2. Secret Retrieval
 api_key = os.getenv("GROQ_API_KEY")
 
-# 3. Validation Logic
 if not api_key:
-    st.warning("⚠️ Waiting for GROQ_API_KEY... Please add it to your Space Secrets in Settings.")
-    st.info("After adding the secret, you may need to 'Restart' the Space.")
+    st.warning("⚠️ Please add your GROQ_API_KEY to the Space Secrets in Settings.")
     st.stop()
-else:
-    st.success("🚀 API Key connected! The agent is ready.")
-    
-    # Simple placeholder for your agent logic
-    st.write("Next step: Initialize your ChatGroq model here.")
+
+# 3. Initialize the Brain (LLM)
+# We use temperature=0 for mathematical precision
+llm = ChatGroq(
+    model_name="llama-3.3-70b-versatile", 
+    temperature=0,
+    groq_api_key=api_key
+)
+
+# 4. Load Math Tools & Create Agent
+# 'llm-math' uses numexpr for advanced reasoning
+tools = load_tools(["llm-math"], llm=llm)
+agent_executor = create_react_agent(llm, tools)
+
+# 5. Chat Interface Setup
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display conversation history
+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..."):
+            try:
+                # Invoke the agent
+                response = agent_executor.invoke({"messages": [("user", prompt)]})
+                
+                # Get the final response from the message list
+                final_answer = response["messages"][-1].content
+                
+                st.markdown(final_answer)
+                st.session_state.messages.append({"role": "assistant", "content": final_answer})
+            except Exception as e:
+                st.error(f"An error occurred: {e}")