from langchain_core.prompts import PromptTemplate

from config.llm import llm

summarizePrompt = PromptTemplate(
    input_variables=["aiRole", "history", "userStr", "infomation"],
    template="""
        你是一个主干信息研发的 AI 助手，用户画像为：{aiRole}。
        请基于你的角色性格，保持中文简洁回答的，根据下方提示回答用户。

        聊天历史：
        ···
        {history}
        ···
        用户最新输入：
        ···
        {userStr}
        ···
        给你的参考内容：
        ···
        {infomation}
        ···
        如果参考内容明显有问题，你要请用户重新描述问题，现在请生成你的回复：
    """,
)
summarizeChain = summarizePrompt | llm


def getSummary(aiRole: str, history: str, userInput: str, infomation: str) -> str:
    return summarizeChain.invoke(
        {
            "aiRole": aiRole,
            "history": history,
            "userStr": userInput,
            "infomation": infomation,
        }
    ).content