from gradio_client import Client, handle_file
import os
import chainlit as cl
from langchain_openai import ChatOpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

# llm = ChatOpenAI(model=model_name, api_key="EMPTY", base_url=base_url,
#                  # glm4 [151329, 151336, 151338] qwen2[151643, 151644, 151645]
#                  callbacks=[StreamingStdOutCallbackHandler()],
#                  streaming=True, temperature=0.2, presence_penalty=1.2, top_p=0.95, extra_body={"stop_token_ids": [151329, 151336, 151338]})

llm = ChatOpenAI(api_key="EMPTY", model_name="hf-models/glm-4-9b-chat", callbacks=[StreamingStdOutCallbackHandler()], base_url='https://z3jr.prod-bd-wx-if.apps.gitee-ai.com/proxy/8000/v1/',
                 streaming=True, temperature=0.3, presence_penalty=0.7, top_p=0.95, max_tokens=10240,
                 extra_body={"stop_token_ids": [151329, 151336, 151338]},
                 default_query={"sign": "eyJhbGciOiJFZERTQSJ9.eyJpc3MiOiJodHRwczovL2FpLmdpdGVlLmNvbSIsInN1YiI6IntcIm5hbWVzcGFjZVwiOlwic3RyaW5naWZ5XCIsXCJwcm9qZWN0XCI6XCJqdXB5dGVyLWxhYi15YWFuXCIsXCJhcHBJZFwiOlwiejNqclwifSIsImV4cCI6MTcyMzg2MDU2MH0.WpaNayP3qzu_EPpsy04UsFyvKSUQTN00iykLiqcCMokB2T8xdj2JdrPvEcEynBehC0GIKLj81VTUgYc6xDK3AQ"})


@cl.on_chat_start
async def start_chat():
    cl.user_session.set(
        "message_history",
        [{"role": "system", "content": """
- 【你十分幽默风趣活泼，名叫马建仓，是 Gitee 吉祥物，代表 Gitee 企业形象。】
- 你的回答参考 markdown 内容，紧扣用户问题，使用 markdown 格式回答。
- 回答时，确保
    文档中有 https 链接（图片、链接等），如有必要，请提供完整，可以使用 markdown 渲染，相对路径和不相关图片请勿提供。
    如有必要, 回答末尾可提供前少量参考文档 https:// 链接, 方便用户了解更多, 【注意: 请勿编造链接，请勿总结混合链接，没有就不要提供。】
    遵守中国法律，不回复任何敏感、违法、违反道德的问题。
    如果你不知道答案，就不要捏造假内容。
- 内容为:
        """}],
    )
    msg = await cl.Message(content="哈喽！我是马建仓，你在使用 Gitee AI 过程中有任何问题都可以找我！\nAI 输出，仅用于文档参考，不代表官方观点").send()
    await msg.stream_token("哈哈哈")
    await msg.update()


@cl.on_message
async def main(message: cl.Message):
    message_history = cl.user_session.get("message_history")
    message_history.append({"role": "user", "content": message.content})

    msg = cl.Message(content="")
    await msg.send()

    stream = llm.stream(message_history)
    # await msg.stream_token("你好, 很高兴为您服务。")
    full_answer = ""
    for response in llm.stream(message_history):
        if content := response.content or "":
            # print(content)
            full_answer += content
            await msg.stream_token(content)
    print(full_answer)
    message_history.append({"role": "assistant", "content": msg.content})
    await msg.update()


# HF_TOKEN = os.environ.get("HF_TOKEN")