app.py

import streamlit as st
from langchain_groq import ChatGroq
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage
import numexpr

# ================== Streamlit UI ==================
st.set_page_config(
    page_title="Text To Math Problem Solver And Data Search Assistant",
    page_icon="🧮",
)

st.title("Text To Math Problem Solver Using Google Gemma 2")
st.write("✅ App loaded. If you see this, the script is running.")

groq_api_key = st.sidebar.text_input(label="Groq API Key", type="password")

if not groq_api_key:
    st.info("Please add your Groq API key to continue")
    st.stop()

# ================== LLM ==================
llm = ChatGroq(model="openai/gpt-oss-120b", groq_api_key=groq_api_key)

# ================== Tools ==================

wikipedia_wrapper = WikipediaAPIWrapper()

@tool
def wikipedia_tool(query: str) -> str:
    """Search Wikipedia and return a summary for the given topic."""
    return wikipedia_wrapper.run(query)

@tool
def calculator(expression: str) -> str:
    """Evaluate a pure math expression like '2+2*3' or 'sqrt(49)'."""
    try:
        return str(numexpr.evaluate(expression))
    except Exception as e:
        return f"Error evaluating expression: {e}"

prompt_text = """
You are an agent tasked with solving users' mathematical questions.
Think step by step, logically arrive at the solution, and provide a detailed explanation.
Display your explanation point-wise.

Question: {question}
Answer:
"""

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

reasoning_chain = prompt_template | llm | StrOutputParser()

@tool
def reasoning_tool(question: str) -> str:
    """Answer math/logic questions with step-by-step reasoning and explanation."""
    return reasoning_chain.invoke({"question": question})

tools = [wikipedia_tool, calculator, reasoning_tool]

# Bind tools so LLM can call them
llm_with_tools = llm.bind_tools(tools)

# ================== Simple tool-calling "agent" ==================

def call_agent(user_input: str) -> str:
    # First call: let the LLM decide whether to call tools
    msg = llm_with_tools.invoke([HumanMessage(content=user_input)])

    # If there are no tool calls, just return the answer
    if not getattr(msg, "tool_calls", None):
        return msg.content

    tool_calls = msg.tool_calls
    tool_results = []

    # Execute tools requested by the LLM
    for tc in tool_calls:
        name = tc["name"]
        args = tc["args"]

        for t in tools:
            if t.name == name:
                result = t.invoke(args)
                tool_results.append(f"Tool `{name}` result:\n{result}")
                break

    # Second call: send tool results back and ask for final answer
    followup_messages = [
        HumanMessage(content=user_input),
        HumanMessage(
            content=(
                "Here are the tool results. Use them to give the final, clear answer:\n\n"
                + "\n\n".join(tool_results)
            )
        ),
    ]

    final_msg = llm.invoke(followup_messages)
    return final_msg.content

# ================== Chat State ==================
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {
            "role": "assistant",
            "content": "Hi, I'm a Math chatbot who can answer all your maths questions!",
        }
    ]

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

# ================== Input & Response ==================
question = st.text_area(
    "Enter your math question or topic to search:",
    key="question_input",
)

if st.button("Submit"):
    if question.strip():
        with st.spinner("Thinking..."):
            st.session_state.messages.append({"role": "user", "content": question})
            st.chat_message("user").write(question)

            response = call_agent(question)

            st.session_state.messages.append(
                {"role": "assistant", "content": response}
            )
            
            st.write("---")
            st.success(response)
            
    else:
        st.warning("Please enter a question or topic to proceed.")