import gradio as gr
import spaces
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage

# 导入已编译的 LangGraph 应用
from graph import app

@spaces.GPU
def respond(message, history, system_message, hf_token: gr.OAuthToken = None):
    """Gradio 接口的响应函数，调用 LangGraph 应用"""
    
    # 将 Gradio 的 history 格式转换为 LangChain 消息格式
    messages = []
    if system_message:
        messages.append(SystemMessage(content=system_message))
    
    for chat_message in history:
        if chat_message['role'] == "user":
            messages.append(HumanMessage(content=chat_message['content']))
        elif chat_message['role'] == "assistant":
            messages.append(AIMessage(content=chat_message['content']))
            
    messages.append(HumanMessage(content=message))

    # 使用 invoke 方法进行一次性调用
    inputs = {"messages": messages}
    final_state = app.invoke(inputs)
    
    # 从最终状态中提取最后一条消息
    final_response = final_state["messages"][-1].content
    
    return final_response

# 重新定义 ChatInterface
chatbot = gr.ChatInterface(
    respond,
    type="messages", # 改为 messages 类型以更好地匹配 LangChain
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
    ],
)

with gr.Blocks() as demo:
    gr.Markdown("# HuggingFace Running")
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()


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