app.py

import streamlit as st
from langchain_groq import ChatGroq
from langchain.chains import LLMMathChain, LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.utilities import WikipediaAPIWrapper
from langchain.agents.agent_types import AgentType
from langchain.agents import Tool, initialize_agent
from langchain.callbacks import StreamlitCallbackHandler
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Streamlit page configuration
st.set_page_config(
    page_title="AI Math Problem Solver & Research Assistant",
    page_icon="🧮",
    layout="wide"
)

# Custom CSS styling
st.markdown("""
    <style>
    .main {
        background-color: #f5f5f5;
    }
    .stTitle {
        color: #1e3d59;
        font-size: 2.5rem !important;
        font-weight: 700 !important;
        padding-bottom: 1rem;
    }
    .stTextArea textarea {
        background-color: #ffffff;
        border-radius: 10px;
        border: 1px solid #e0e0e0;
        padding: 10px;
    }
    .stButton button {
        background-color: #17b794;
        color: white;
        border-radius: 20px;
        padding: 0.5rem 2rem;
        font-weight: 600;
    }
    .stButton button:hover {
        background-color: #148f77;
    }
    div.stSpinner > div {
        border-top-color: #17b794 !important;
    }
    </style>
    """, unsafe_allow_html=True)

# App Header
col1, col2, col3 = st.columns([1,6,1])
with col2:
    st.title("🧮 AI Math Problem Solver & Research Assistant")
    st.markdown("""
    <div style='background-color: #ffffff; padding: 1rem; border-radius: 10px; margin-bottom: 2rem;'>
        <p style='color: #666666; margin-bottom: 0;'>
            Powered by Google Gemma 2 AI, this assistant can help you solve math problems, 
            provide detailed explanations, and search for additional information.
        </p>
    </div>
    """, unsafe_allow_html=True)

# API Key Check
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    st.error("⚠️ Please add your Groq API key to continue")
    st.stop()

# Initialize LLM
llm = ChatGroq(model="gemma2-9b-it", groq_api_key=groq_api_key)

# Tool Setup
wikipedia_wrapper = WikipediaAPIWrapper()
wikipedia_tool = Tool(
    name="Wikipedia",
    func=wikipedia_wrapper.run,
    description="A tool for searching the Internet to find various information on the topics mentioned"
)
def safe_calculator(expression: str) -> str:
    try:
        # Clean and validate the expression
        if any(char in expression for char in ['∫', '∂', '∑']):
            return "I apologize, but I cannot directly solve calculus problems or complex mathematical expressions. I can help explain the steps to solve it though!"
        
        # Use the math chain
        result = math_chain.run(expression)
        return result
    except Exception as e:
        return f"I encountered an error trying to solve this mathematically. Let me help explain the steps to solve it instead."
    
math_chain = LLMMathChain.from_llm(llm=llm,verbose=True,input_key="question",output_key="answer")
calculator = Tool(
    name="Calculator",
    func=safe_calculator,
    description="A tool for solving basic mathematical expressions. For complex math, it will provide step-by-step explanations"
)

prompt = """
You're a helpful math tutor tasked with solving mathematical questions. For each problem:
1. First determine if it's a basic arithmetic problem or a more complex mathematical problem
2. For basic arithmetic, use the calculator tool
3. For complex math (calculus, integrals, differential equations), explain the solution steps clearly
4. Always show your work and explain each step

Question: {question}

Let me solve this step by step:
"""

prompt_template = PromptTemplate(
    input_variables=["question"],
    template=prompt
)

chain = LLMChain(llm=llm, prompt=prompt_template)
reasoning_tool = Tool(
    name="Reasoning tool",
    func=chain.run,
    description="A tool for answering logic-based and reasoning questions."
)

# Initialize Agent
assistant_agent = initialize_agent(
    tools=[wikipedia_tool, calculator, reasoning_tool],
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=False,
    handle_parsing_errors=True
)

# Chat History
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "assistant", "content": "👋 Hi! I'm your Math Assistant. I can help you solve math problems and provide detailed explanations."}
    ]

# Display Chat History
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])

# Input Section
st.markdown("### 📝 Your Question")
question = st.text_area(
    label="Enter your question:",
    value="I have 5 bananas and 7 grapes. I eat 2 bananas and give away 3 grapes. Then I buy a dozen apples and 2 packs of blueberries. Each pack of blueberries contains 25 berries. How many total pieces of fruit do I have at the end?",
    label_visibility="collapsed",
    height=100
)

# Create two columns for button centering
col1, col2, col3 = st.columns([2,1,2])
with col2:
    solve_button = st.button("🔍 Solve Problem")

if solve_button:
    if question:
        with st.spinner("🤔 Thinking..."):
            st.session_state.messages.append({"role": "user", "content": question})
            with st.chat_message("user"):
                st.write(question)
            
            st_cb = StreamlitCallbackHandler(st.container(), expand_new_thoughts=False)
            response = assistant_agent.run(st.session_state.messages, callbacks=[st_cb])
            
            st.session_state.messages.append({"role": "assistant", "content": response})
            
            st.markdown("### 💡 Solution:")
            with st.chat_message("assistant"):
                st.success(response)
    else:
        st.warning("⚠️ Please enter your question first!")

# Footer
st.markdown("""
<div style='position: fixed; bottom: 0; left: 0; width: 100%; background-color: #f0f2f6; padding: 1rem; text-align: center;'>
    <p style='color: #666666; margin-bottom: 0;'>
        Made with ❤️ using Streamlit and Google Gemma 2
    </p>
</div>
""", unsafe_allow_html=True)