import gradio as gr
from langchain.agents import Tool, AgentType, initialize_agent
from langchain_community.llms import HuggingFacePipeline
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from transformers import pipeline
import os

# 创建搜索工具
search = DuckDuckGoSearchAPIWrapper()
search_tool = Tool(
    name="Web Search",
    func=search.run,
    description="Useful for when you need to search the web for current information."
)

# 使用 Transformers pipeline 初始化语言模型
model_name = "google/flan-t5-large"
pipe = pipeline("text2text-generation", model=model_name, max_length=512)
llm = HuggingFacePipeline(pipeline=pipe)

# 创建 agent
tools = [search_tool]
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

# 定义 Gradio 接口的主函数
def answer_question(question):
    try:
        response = agent.run(question)
        return response
    except Exception as e:
        return f"An error occurred: {str(e)}"

# 创建 Gradio 接口
iface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
    outputs="text",
    title="AI Research Assistant",
    description="Ask a question, and I'll search the web to find an answer for you.",
    examples=[
        ["What is the latest news about SpaceX?"],
        ["Who won the last FIFA World Cup?"],
        ["What are the main features of Python 3.10?"]
    ]
)

# 启动 Gradio 应用
iface.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)), debug=True)