import os
from dotenv import load_dotenv
import gradio as gr
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import initialize_agent
from langchain_community.agent_toolkits.load_tools import load_tools


def main(query, google_api_key, google_cse_id):
    # Set environment variables
    os.environ["GOOGLE_API_KEY"] = google_api_key
    os.environ["GOOGLE_CSE_ID"] = google_cse_id

    llm = ChatGoogleGenerativeAI(
        model="gemini-pro",
        temperature=0.0
        )
    tools = load_tools(["google-search"], llm=llm)
    agent = initialize_agent(tools, llm, agent="zero-shot-react-description", verbose=True)

    return agent.run(query)


if __name__ == "__main__":
    try:
        app = gr.Interface(
            fn=main,
            inputs=[
                gr.Textbox(label="Search Query (検索クエリ)"),
                gr.Textbox(label="Google API Key (https://console.cloud.google.com/apis/credentials)"),
                gr.Textbox(label="Google Programmable Search Engine ID (https://programmablesearchengine.google.com)")
                ],
            outputs=[gr.Textbox(label="Search Result (検索結果)")],
            title="Google Search enhanced by LLM"
        )
        app.launch(share=True)
    except Exception as e:
        raise e