import os
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from mcp_use import MCPAgent, MCPClient
import gradio as gr

# Load environment variables
load_dotenv()
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
config_file = "browser_mcp.json"

client = MCPClient.from_config_file(config_file)
llm = ChatGroq(model="llama3-8b-8192")
agent = MCPAgent(llm=llm, client=client, max_steps=15, memory_enabled=True)

def chat(user_input, history=[]):
    if user_input.lower() == "clear":
        agent.clear_conversation_history()
        return "", []
    response = agent.run(user_input)
    history = history + [[user_input, response]]
    return "", history

with gr.Blocks() as demo:
    gr.Markdown("# MCP Chatbot")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Type your message")
    clear = gr.Button("Clear Conversation")

    def respond(message, chat_history):
        _, updated_history = chat(message, chat_history)
        return "", updated_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: ("", []), None, [msg, chatbot])

demo.launch()