from langchain.agents import initialize_agent, AgentType
from langchain_openai import OpenAI
from langchain_community.agent_toolkits.load_tools import load_tools
import gradio as gr
import os

# 🔐 Get API keys from environment
openai_api_key = os.getenv("OPENAI_API_KEY")
serpapi_api_key = os.getenv("SERPAPI_API_KEY")

# ✅ Load LLM
llm = OpenAI(temperature=0, api_key=openai_api_key)

# ✅ Load tools
tools = load_tools(
    ["wikipedia", "serpapi", "ddg-search", "pubmed"],
    llm=llm,
    serpapi_api_key=serpapi_api_key
)

# ✅ Initialize agent
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True
)

# 🧠 Agent logic
def ask_agent(question):
    try:
        return agent.run(question)
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

def clear_input():
    return "",""

# 🎨 Gradio UI
with gr.Blocks(css="""
.main-title {
    text-align: center;
    font-size: 30px;
    font-weight: bold;
    margin-bottom: 10px;
}
.gr-button {
    font-weight: bold;
}
.gr-textbox {
    border-radius: 10px;
}
textarea {
    font-family: monospace;
}
""") as demo:

    gr.Markdown("<div class='main-title'>🤖 Smart AI Research Agent</div>")

    gr.Markdown(
        "This agent can:\n"
        "- 🔍 Search Wikipedia\n"
        "- 🌐 Perform web search (SerpAPI + DuckDuckGo)\n"
        "- 🧬 Find scientific papers (PubMed)\n\n"
        "Type your question below and click **Ask Agent**."
    )

    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(label="Your Question", placeholder="e.g. What is the role of serotonin?")
            with gr.Row():
                submit = gr.Button("🚀 Ask Agent")
                clear = gr.Button("🧹 Clear")
        with gr.Column(scale=5):
            output = gr.Textbox(label="Agent Response", lines=15, interactive=False)

    submit.click(fn=ask_agent, inputs=question, outputs=output)
    clear.click(fn=clear_input, inputs=[], outputs=[question,output])

demo.launch()
if __name__ == "__main__":
    demo.launch()