from langchain.prompts import PromptTemplate
from config.llm import llm,llmThink
import db.milvus as milvus
import db.postgres as pg
import json

memPathPrompt = PromptTemplate(
    input_variables=["ai_role", "CHAT_RECORD"],
    template = """
你是一个记忆筛选器，负责判断最近对话的信息中，用户的回复内容是否对业务具有长期价值或潜在价值，或者可以帮助形成用户画像。
首先，请仔细阅读以下关于你业务的描述： 
<ai_role> 
{ai_role} 
</ai_role> 
现在，请仔细阅读以下你与用户的聊天记录： 
<聊天记录> 
{CHAT_RECORD} 
</聊天记录>

请仔细考虑以下标准：

1. 长期价值：用户最新回复信息是否能为你的业务提供知识积累、经验总结、数据支持。

2. 相关性：用户最新回复是否与业务核心需求、目标、流程或潜在业务场景相关。

3. 潜在可用性：用户最新回复是否可能在未来的业务场景中被重复使用、参考或触发进一步操作。

你需要根据以上标准给出判断并得出"yes"或"no"
yes：用户最新回复具有直接或潜在长期价值，值得保留。
no：用户最新回复价值有限或几乎不会在未来业务中使用。

回复不要带任何标点符号以及空格、换行符。
请给出你的判断结果：
    """
)
memPathChain = memPathPrompt | llmThink
memPrompt = PromptTemplate(
    input_variables=["CHAT_RECORD"],
    template = """
你的任务是对给定的聊天记录进行关键信息的记忆总结。请仔细阅读以下聊天记录，并按照要求进行总结：
<聊天记录>
{CHAT_RECORD}
</聊天记录>
在总结时，请遵循以下指南：
1. 提取聊天记录中的用户所说的关键信息，包括主要话题、重要观点、达成的共识或决定等。
2. 用简洁明了的语言进行总结有价值的信息，避免冗长和复杂的表述。
3. 确保总结内容准确反映聊天记录中用户的核心内容，并尽可能简短。
4. 总结内容应包含时间，并确保时间是准确的。
5. 你需要针对你的业务场景{ai_role}，展开对用户最后回复的总结。
请生成你的总结，以用户、时间开头：
    """
)
memChain = memPrompt | llmThink
def take_memory(ai_id:str,sessionId: str,user_id:str, max_retry=3):
    """根据用户输入选择数据来源"""
    history = pg.get_history_with_time(sessionId,10)
    print("获取的历史记录:",history)
    ai_service = pg.get_description(ai_id)
    if ai_service == "":
        # AI描述没有描述，则取业务字段
        json = pg.get_ai_personality(ai_id)
        if json.get("业务", "") == "":
            # AI没有任何描述，无法对记忆价值进行判断
            print("AI没有任何描述，无法对记忆价值进行判断")
            return
        else:
            ai_service = json["业务"]
    print("获取的描述是：", ai_service)
    choice = memPathChain.invoke({
        "ai_role": ai_service,
        "CHAT_RECORD": history,
    }).content.strip().lower()
    print("记忆判断器判断的结果是：", choice)
    if choice == "yes":
        # 对对话进行总结
        memory = memChain.invoke({
            "CHAT_RECORD": history,
            "ai_role": ai_service,
        }).content.strip().lower()
        print("记忆生成结果是：", memory)
        milvus.add_memory(mem = memory,user_id = user_id, is_active = True, ai_id = ai_id)
        return