import pickle

from langchain.chains.conversation.base import ConversationChain
from langchain_core.messages import messages_to_dict, messages_from_dict

from ChatGLM_new import xinghuo_llm
from langchain.memory import ChatMessageHistory, ConversationBufferMemory, ConversationSummaryMemory

history = ChatMessageHistory()
history.add_user_message("hi!")
history.add_ai_message("whats up?")
print(history.messages)

print("##################对话缓冲区内存（ConversationBufferMemory）#######################")

conversation = ConversationChain(
    llm=xinghuo_llm,
    verbose=True,
    memory=ConversationBufferMemory()
)
# conversation.predict(input="I'm doing well!")
# conversation.predict(input="你是谁")
# conversation.predict(input="你的背后原理是什么")
#
# print(conversation.memory.chat_memory.messages)

print("##################保存历史消息#######################")
# dicts = messages_to_dict(conversation.memory.chat_memory.messages)
# print(dicts)
#
# f = open("./memory", 'wb')
# pickle.dump(dicts,f)
# f.close()

print("##################读取历史消息#######################")
# 读取
# dicts_load = pickle.load(open("./memory", "rb"))
#
# new_messages = messages_from_dict(dicts_load)
# retrieved_chat_history = ChatMessageHistory(messages=new_messages)
# retrieved_memory = ConversationBufferMemory(chat_memory=retrieved_chat_history)
# conversation_reload = ConversationChain(
#     llm=xinghuo_llm,
#     verbose=True,
#     memory=retrieved_memory
# )
# conversation_reload.predict(input="我回来了")

print("##################会话摘要（ConversationSummaryMemory）#######################")
# conversation_with_summary = ConversationChain(
#     llm=xinghuo_llm,
#     memory=ConversationSummaryMemory(llm=xinghuo_llm),
#     verbose=True
# )
# print(conversation_with_summary.predict(input="嗨，怎么了?"))
# print(conversation_with_summary.predict(input="告诉我更多关于它的信息！"))
# print(conversation_with_summary.predict(input="很酷 -- 项目的范围是什么?"))