# Import necessary libraries
import streamlit as st
import autogen
from autogen.agentchat.contrib.math_user_proxy_agent import MathUserProxyAgent

# Function to run the query
def run_query(math_problem, api_key):
    config_list = [
        {
            'model': 'gpt-3.5-turbo',
            'api_key': api_key,
        },
    ]

    autogen.ChatCompletion.start_logging()

    assistant = autogen.AssistantAgent(
        name="assistant",
        system_message="You are a helpful assistant.",
        llm_config={
            "request_timeout": 600,
            "seed": 42,
            "config_list": config_list,
        }
    )

    mathproxyagent = MathUserProxyAgent(
        name="mathproxyagent",
        human_input_mode="NEVER",
        code_execution_config={"use_docker": False},
    )

    return mathproxyagent.initiate_chat(assistant, problem=math_problem)
   

# Streamlit app
st.title('Math Problem Solver')

# Input fields for API key and math problem
api_key = st.text_input("Enter your API Key:")
math_problem = st.text_input("Enter your math problem:")

# When the button is pressed, call run_query and display the result
if st.button('Solve'):
    if api_key and math_problem:
        result = run_query(math_problem, api_key)
        st.write(result)
    else:
        st.write("Please provide both API key and math problem.")