# 大模型记忆功能
# 使用消息历史组件ChatMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, \
    MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_redis import RedisChatMessageHistory

from models import get_lc_model_client
client = get_lc_model_client()
parser = StrOutputParser()

# 手动版
# chat_prompt = ChatPromptTemplate.from_messages(
#     [
#         SystemMessagePromptTemplate.from_template("你是人工智能助手"),
#         HumanMessagePromptTemplate.from_template("{input}"),
#         MessagesPlaceholder(variable_name="messages"),
#     ]
# )
#
# chain = chat_prompt | client | parser

# # 使用组件InMemoryChatMessageHistory(存在内存中)
chat_history = InMemoryChatMessageHistory()
#
# while True:
#     user_input = input("请输入:")
#     if user_input == "exit":
#         print("退出聊天,再见!!!")
#         break
#     chat_history.add_user_message(user_input)
#     response = chain.invoke({"messages":chat_history.messages,'input':user_input})
#     print(f'来自大模型的回复:{response}')
#     print("*"*10,"A"*10,"*"*10)
#     chat_history.add_ai_message(response)

# 自动版
chat_prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate.from_template("你是人工智能助手,用中文回答所有的问题"),
        ("human","{input}"),
    ]
)

chain = chat_prompt | client | parser

# 串联历史记录(记忆)和链
# 将对话历史自动集成到模型调用链中，解决了聊天机器人在上下文连续性和多用户支持中的核心问题
chat_with_human = RunnableWithMessageHistory(
    chain,
    # 根据session_id获取历史记录
    get_session_history = lambda session_id: RedisChatMessageHistory(
        session_id = session_id,
        redis_url="redis://localhost:6379"
    ),
    input_message_key = "input"
)
# 多用户
user_A = chat_with_human.invoke({
    "input":"我是用户A,平时喜欢下地干活,今天发现地里有块不规则的金黄色的石头，这是什么呢？"},
    config = {
        "configurable":{
            "session_id":"user_A"
        },
        # 注册跟踪器和控制台输出
        # "callbacks": [tracer, console_handler]
    }
)
print("user_A的回复1:",user_A)
print("*"*10,"A"*10,"*"*10)

user_B = chat_with_human.invoke({
    "input":"我是用户B,平时喜欢打篮球，最近篮球扁了需要重买一个，请给我推荐一个？"},
    config = {
        "configurable":{
            "session_id":"user_B"
        },
        # 注册跟踪器和控制台输出
        # "callbacks": [tracer, console_handler]
    }
)
print("user_B的回复1:",user_B)
print("*"*10,"B"*10,"*"*10)