Update app.py

app.py

@@ -23,10 +23,28 @@ llm = ChatGroq(
     groq_api_key=api_key
 )
 
-# 4. Load Math Tools & Create Agent
-# 'llm-math' uses numexpr for advanced reasoning
+# ... (Keep your imports and LLM initialization the same)
+
+# 4. Load Math Tools
 tools = load_tools(["llm-math"], llm=llm)
-agent_executor = create_react_agent(llm, tools)
+
+# --- ADD THE NEW BLOCK HERE ---
+system_message = (
+    "You are a strict Mathematics Assistant. "
+    "Your ONLY task is to solve mathematical problems using your tools. "
+    "If a user asks about anything other than math (such as history, general knowledge, "
+    "coding, or personal opinions), you must politely refuse and explain that you "
+    "only provide assistance with mathematics."
+)
+
+agent_executor = create_react_agent(
+    llm, 
+    tools, 
+    state_modifier=system_message
+)
+# --- END OF NEW BLOCK ---
+
+# ... (Keep the Chat Interface logic below this)
 
 # 5. Chat Interface Setup
 if "messages" not in st.session_state: