import os

import dotenv
from langchain.memory import ConversationBufferMemory, ConversationTokenBufferMemory, ConversationSummaryMemory
from langchain_openai import ChatOpenAI

memory = ConversationBufferMemory()
memory.chat_memory.add_user_message("我叫小明")
memory.chat_memory.add_ai_message("噢，好的，你叫小明")
print(memory)
print(memory.memory_key)
print(memory.load_memory_variables({}))
print('---------')

memory2 = ConversationBufferMemory(memory_key="memory2")
memory2.chat_memory.add_user_message("我叫小明")
memory2.chat_memory.add_ai_message("噢，好的，你叫小明")
print(memory2)
print(memory2.memory_key)
print(memory2.load_memory_variables({}))
print('---------')

memory3 = ConversationBufferMemory(return_messages=True)
memory3.chat_memory.add_user_message("我叫小明")
memory3.chat_memory.add_ai_message("噢，好的，你叫小明")
print(memory3)
print(memory3.memory_key)
print(memory3.load_memory_variables({}))
print('---------')
dotenv.load_dotenv()

## 设置环境变量
os.environ['OPENAI_API_KEY'] = os.getenv("SILICONFLOW_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("SILICONFLOW_BASE_URL")

# 默认的 'model_name': 'deepseek-ai/DeepSeek-V3.1',
llm = ChatOpenAI(model="moonshotai/Kimi-Dev-72B")

memory_test = ConversationSummaryMemory(llm=llm)
memory_test.save_context({"input": "我叫小明"}, {"output": "噢，好的，你叫小明"})
memory_test.save_context({"input": "那么你是谁呢"}, {"output": "我是一个无所不能的AI聊天助手，可以帮你解答任何问题。"})
print(memory_test.load_memory_variables({}))
# {'history': '\n\nThe human introduces themselves as xiaoming. The AI confirms the name and responds. The human then asks who the AI is. The AI introduces itself as an all-powerful AI chat assistant who can answer any questions.'}