Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import os
 from langchain_groq import ChatGroq
 from langchain_community.agent_toolkits.load_tools import load_tools
+from langchain_experimental.tools import PythonREPLTool
 from langgraph.prebuilt import create_react_agent
 
 # 1. Page Configuration
@@ -31,21 +32,32 @@ llm = ChatGroq(
 
 # 4. Load Math Tools
 # Load the tools FIRST
-tools = load_tools(["llm-math"], llm=llm)
+# Create the Python Tool
+python_tool = PythonREPLTool()
+
+# Load the standard math tools
+math_tools = load_tools(["llm-math"], llm=llm)
+
+# Combine them into a single list for the agent
+tools = [python_tool] + math_tools
 from langchain_core.messages import SystemMessage
 
 system_message = SystemMessage(
     content=(
-        "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."
+        "You are a specialized Mathematics Assistant. Your ONLY task is to solve problems "
+        "within these branches: Arithmetic, Algebra, Geometry, Calculus, Analysis, "
+        "Statistics & Probability, Trigonometry, Number Theory, and Topology.\n\n"
+        "GUIDELINES:\n"
+        "1. Use the 'python_repl' tool for complex formulas, symbolic algebra (SymPy), "
+        "or statistical analysis (SciPy/NumPy).\n"
+        "2. Use 'llm-math' for simple arithmetic.\n"
+        "3. If a user asks about history, coding (non-math), or any non-mathematical topic, "
+        "politely refuse.\n"
+        "4. Always show your work or explain the mathematical steps taken."
     )
 )
 
-
-# We will pass the system message in the actual chat logic instead
+# Re-create the agent with the new tools and instructions
 agent_executor = create_react_agent(llm, tools)
 
 # 5. Chat Interface Setup