import requests
import json

SYS_MSG_4_CLEANING = "你是一个AI助手, 能将我给你的文章去除与主题无关的句子, 并尽量保留所有与主题相关的句子."

class GlmHelper(object):
    def clean_raw_content(self, content: str):
        history = []
        rply = self.bot_message_handler(message=content, history=history, sys_msg=SYS_MSG_4_CLEANING)
        return rply

    # 携带知识库文本询问LLM
    def bot_message_handler(self, message: str, history: [list], sys_msg: str):
        request_body = {
            "prompt": f"""
            <s>[INST] <<SYS>>\n{sys_msg}\n<</SYS>>\n\n{message} [/INST]
            """,
            "knowledge": """
            """,
            "history": history,
            "max_length": 2048 * 4,
        }
        rply = requests.post("http://region-9.autodl.pro:19567/gradio", data=json.dumps(request_body))
        try:
            reply_from_GLM = rply.json()["response"]
        except:
            reply_from_GLM = "GLM Api返回了坏的请求..."

        return reply_from_GLM