import os
from langchain.chat_models import ChatOpenAI
from langchain.schema import AIMessage, HumanMessage
import gradio as gr

os.environ["OPENAI_API_KEY"] = "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"
llm = ChatOpenAI(model='gpt-3.5-turbo-0613', streaming=True)


def stream_resp(message, history, flag1, flag2):
    history_langchain_format = []
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))
    history_langchain_format.append(HumanMessage(content=message))

    partial_message = ""
    for chunk in llm.stream(history_langchain_format):
        partial_message = partial_message + chunk.content
        yield partial_message
    

demo = gr.ChatInterface(
    stream_resp,
    chatbot=gr.Chatbot(height=430, label="ChatReport"),
    textbox=gr.Textbox(placeholder="请输入问题", container=False, scale=7),
    title="研报助手",
    description="清芬院研报助手",
    theme="soft",
    examples=["你好", "你是谁"],
    cache_examples=True,
    retry_btn="retry",
    undo_btn="清空输入框",
    clear_btn="清空聊天记录",
    additional_inputs=[
            gr.Checkbox(label = "研报问答"),
            gr.Checkbox(label = "研报生成"),
        ],
).queue()


if __name__ == "__main__":
    demo.launch(server_port=8080, share=True)