78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
import os
|
||
|
||
import dotenv
|
||
from langchain.memory import ConversationBufferMemory, ConversationTokenBufferMemory, ConversationSummaryMemory
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain.memory import ConversationBufferWindowMemory
|
||
from langchain_core.messages import HumanMessage, AIMessage, trim_messages
|
||
|
||
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("OLLAMA_API_KEY")
|
||
os.environ['OPENAI_BASE_URL'] = os.getenv("OLLAMA_BASE_URL")
|
||
|
||
# 默认的 'model_name': 'deepseek-ai/DeepSeek-V3.1',
|
||
llm = ChatOpenAI(model="gemma4:e2b")
|
||
|
||
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.'}
|
||
print('---------')
|
||
|
||
# 这是openai专用的OpenAI套件
|
||
# memory_token = ConversationTokenBufferMemory(llm=llm,max_token_limit=7)
|
||
# memory_token.save_context({"input": "我叫小明"}, {"output": "噢，好的，你叫小明"})
|
||
# memory_token.save_context({"input": "那么你是谁呢"}, {"output": "我是一个无所不能的AI聊天助手，可以帮你解答任何问题。"})
|
||
# print(memory_token.load_memory_variables({}))
|
||
|
||
# k=1 代表只記住「最新的一問一答」（1輪）
|
||
memory_window = ConversationBufferWindowMemory(k=1)
|
||
memory_window.save_context({"input": "我叫小明"}, {"output": "噢，好的，你叫小明"})
|
||
memory_window.save_context({"input": "那么你是谁呢"}, {"output": "我是一个无所不能的AI聊天助手。"})
|
||
print(memory_window.load_memory_variables({}))
|
||
# 輸出結果只會保留最後一輪，因為 k=1
|
||
print('---------')
|
||
|
||
messages = [
|
||
HumanMessage("我叫小明"),
|
||
AIMessage("噢，好的，你叫小明"),
|
||
HumanMessage("那么你是谁呢"),
|
||
AIMessage("我是一个无所不能的AI聊天助手。")
|
||
]
|
||
|
||
# 現代化做法：直接裁剪訊息列表
|
||
# 這裡我們用最簡單的 "只保留最後 2 條訊息" 策略
|
||
trimmed_messages = trim_messages(
|
||
messages,
|
||
max_tokens=2,
|
||
token_counter=len, # 關鍵：我們直接告訴它「算訊息條數」當作 token，避開底層 tiktoken 崩潰
|
||
strategy="last",
|
||
)
|
||
|
||
print(trimmed_messages) |