117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Memory Bank 简化测试脚本
|
||
"""
|
||
|
||
import os
|
||
import asyncio
|
||
from datetime import datetime
|
||
import json
|
||
|
||
# Google GenAI 导入
|
||
import google.genai as genai
|
||
|
||
class MemoryBankTest:
|
||
"""Memory Bank 测试类"""
|
||
|
||
def __init__(self):
|
||
self.api_key = os.getenv('GOOGLE_API_KEY')
|
||
if not self.api_key:
|
||
raise ValueError("请设置 GOOGLE_API_KEY 环境变量")
|
||
|
||
self.client = genai.Client(api_key=self.api_key)
|
||
|
||
# 八仙人格基线
|
||
self.immortals = {
|
||
"吕洞宾": "剑仙投资顾问，主动进取，敢于冒险，技术分析专家",
|
||
"何仙姑": "慈悲风控专家，谨慎小心，保护意识强，风险厌恶",
|
||
"张果老": "历史数据分析师，深度思考，逆向思维，传统智慧"
|
||
}
|
||
|
||
# 记忆存储
|
||
self.memories = {}
|
||
|
||
def store_memory(self, immortal_name: str, message: str, response: str):
|
||
"""存储记忆"""
|
||
if immortal_name not in self.memories:
|
||
self.memories[immortal_name] = []
|
||
|
||
self.memories[immortal_name].append({
|
||
"message": message,
|
||
"response": response,
|
||
"timestamp": datetime.now().isoformat()
|
||
})
|
||
|
||
def chat_with_immortal(self, immortal_name: str, message: str) -> str:
|
||
"""与仙人对话"""
|
||
# 构建上下文
|
||
context = f"你是{immortal_name}，{self.immortals[immortal_name]}。"
|
||
|
||
# 添加记忆
|
||
if immortal_name in self.memories and self.memories[immortal_name]:
|
||
context += "\n\n你的历史对话："
|
||
for memory in self.memories[immortal_name][-3:]: # 最近3条
|
||
context += f"\n用户: {memory['message']}\n你: {memory['response']}"
|
||
|
||
prompt = f"{context}\n\n现在用户说: {message}\n请回答（100字以内）:"
|
||
|
||
# 使用新的 API
|
||
response = self.client.models.generate_content(
|
||
model="gemini-2.0-flash-exp",
|
||
contents=[{"parts": [{"text": prompt}]}]
|
||
)
|
||
return response.candidates[0].content.parts[0].text
|
||
|
||
def test_memory_continuity(self):
|
||
"""测试记忆连续性"""
|
||
print("🧪 测试记忆连续性")
|
||
print("=" * 50)
|
||
|
||
# 测试吕洞宾
|
||
print("\n🎭 测试吕洞宾:")
|
||
messages = [
|
||
"你觉得特斯拉股票怎么样？",
|
||
"现在市场波动很大，你怎么看？",
|
||
"你之前不是看好特斯拉吗？现在还是这个观点吗？"
|
||
]
|
||
|
||
for message in messages:
|
||
print(f"\n用户: {message}")
|
||
response = self.chat_with_immortal("吕洞宾", message)
|
||
print(f"吕洞宾: {response}")
|
||
self.store_memory("吕洞宾", message, response)
|
||
|
||
# 测试何仙姑
|
||
print("\n🎭 测试何仙姑:")
|
||
messages = [
|
||
"特斯拉股票风险大吗？",
|
||
"现在适合投资吗？",
|
||
"你一直很谨慎，现在还是建议观望吗？"
|
||
]
|
||
|
||
for message in messages:
|
||
print(f"\n用户: {message}")
|
||
response = self.chat_with_immortal("何仙姑", message)
|
||
print(f"何仙姑: {response}")
|
||
self.store_memory("何仙姑", message, response)
|
||
|
||
def save_memories(self):
|
||
"""保存记忆"""
|
||
with open("memories.json", "w", encoding="utf-8") as f:
|
||
json.dump(self.memories, f, ensure_ascii=False, indent=2)
|
||
print("💾 记忆已保存到 memories.json")
|
||
|
||
def main():
|
||
"""主函数"""
|
||
print("🚀 Memory Bank 测试开始")
|
||
|
||
test = MemoryBankTest()
|
||
test.test_memory_continuity()
|
||
test.save_memories()
|
||
|
||
print("\n✅ 测试完成！")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|