276 lines
10 KiB
Python
276 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Memory Bank 实验脚本
|
||
测试八仙人格的长期记忆功能
|
||
"""
|
||
|
||
import os
|
||
import asyncio
|
||
from datetime import datetime
|
||
from typing import Dict, List, Any
|
||
import json
|
||
|
||
# Google GenAI 导入
|
||
try:
|
||
import google.genai as genai
|
||
from google.genai import types
|
||
except ImportError:
|
||
print("❌ 请安装 google-genai: pip install google-genai")
|
||
exit(1)
|
||
|
||
class MemoryBankExperiment:
|
||
"""Memory Bank 实验类"""
|
||
|
||
def __init__(self):
|
||
self.api_key = os.getenv('GOOGLE_API_KEY')
|
||
if not self.api_key:
|
||
raise ValueError("请设置 GOOGLE_API_KEY 环境变量")
|
||
|
||
# 初始化 GenAI
|
||
genai.configure(api_key=self.api_key)
|
||
|
||
# 八仙人格基线
|
||
self.immortal_baselines = {
|
||
"吕洞宾": {
|
||
"mbti_type": "ENTJ",
|
||
"core_traits": {
|
||
"assertiveness": 0.9,
|
||
"analytical": 0.8,
|
||
"risk_tolerance": 0.8,
|
||
"optimism": 0.7
|
||
},
|
||
"personality_description": "剑仙投资顾问，主动进取，敢于冒险，技术分析专家"
|
||
},
|
||
"何仙姑": {
|
||
"mbti_type": "ISFJ",
|
||
"core_traits": {
|
||
"empathy": 0.9,
|
||
"caution": 0.8,
|
||
"loyalty": 0.8,
|
||
"optimism": 0.4
|
||
},
|
||
"personality_description": "慈悲风控专家，谨慎小心，保护意识强，风险厌恶"
|
||
},
|
||
"张果老": {
|
||
"mbti_type": "INTP",
|
||
"core_traits": {
|
||
"analytical": 0.9,
|
||
"curiosity": 0.8,
|
||
"traditional": 0.7,
|
||
"caution": 0.6
|
||
},
|
||
"personality_description": "历史数据分析师，深度思考，逆向思维，传统智慧"
|
||
}
|
||
}
|
||
|
||
# 记忆存储（模拟 Memory Bank）
|
||
self.memory_bank = {}
|
||
|
||
def initialize_immortal_memory(self, immortal_name: str):
|
||
"""初始化仙人的记忆空间"""
|
||
if immortal_name not in self.memory_bank:
|
||
self.memory_bank[immortal_name] = {
|
||
"personality_baseline": self.immortal_baselines[immortal_name],
|
||
"conversation_history": [],
|
||
"viewpoint_evolution": [],
|
||
"decision_history": [],
|
||
"created_at": datetime.now().isoformat(),
|
||
"last_updated": datetime.now().isoformat()
|
||
}
|
||
print(f"🎭 初始化 {immortal_name} 的记忆空间")
|
||
|
||
def store_memory(self, immortal_name: str, memory_type: str, content: Dict[str, Any]):
|
||
"""存储记忆到 Memory Bank"""
|
||
self.initialize_immortal_memory(immortal_name)
|
||
|
||
memory_entry = {
|
||
"type": memory_type,
|
||
"content": content,
|
||
"timestamp": datetime.now().isoformat(),
|
||
"session_id": f"session_{len(self.memory_bank[immortal_name]['conversation_history'])}"
|
||
}
|
||
|
||
if memory_type == "conversation":
|
||
self.memory_bank[immortal_name]["conversation_history"].append(memory_entry)
|
||
elif memory_type == "viewpoint":
|
||
self.memory_bank[immortal_name]["viewpoint_evolution"].append(memory_entry)
|
||
elif memory_type == "decision":
|
||
self.memory_bank[immortal_name]["decision_history"].append(memory_entry)
|
||
|
||
self.memory_bank[immortal_name]["last_updated"] = datetime.now().isoformat()
|
||
print(f"💾 {immortal_name} 存储了 {memory_type} 记忆")
|
||
|
||
def retrieve_relevant_memories(self, immortal_name: str, query: str) -> List[Dict]:
|
||
"""检索相关记忆"""
|
||
if immortal_name not in self.memory_bank:
|
||
return []
|
||
|
||
# 简单的关键词匹配（实际应该使用向量相似度搜索）
|
||
relevant_memories = []
|
||
query_lower = query.lower()
|
||
|
||
for memory in self.memory_bank[immortal_name]["conversation_history"]:
|
||
if any(keyword in memory["content"].get("message", "").lower()
|
||
for keyword in query_lower.split()):
|
||
relevant_memories.append(memory)
|
||
|
||
return relevant_memories[-5:] # 返回最近5条相关记忆
|
||
|
||
async def generate_immortal_response(self, immortal_name: str, query: str) -> str:
|
||
"""生成仙人的回应，基于记忆和人格基线"""
|
||
# 检索相关记忆
|
||
relevant_memories = self.retrieve_relevant_memories(immortal_name, query)
|
||
|
||
# 构建上下文
|
||
context = self.build_context(immortal_name, relevant_memories)
|
||
|
||
# 生成回应
|
||
model = genai.GenerativeModel('gemini-2.0-flash-exp')
|
||
|
||
prompt = f"""
|
||
你是{immortal_name}，{self.immortal_baselines[immortal_name]['personality_description']}。
|
||
|
||
你的核心人格特质：
|
||
{json.dumps(self.immortal_baselines[immortal_name]['core_traits'], ensure_ascii=False, indent=2)}
|
||
|
||
你的相关记忆：
|
||
{json.dumps(relevant_memories, ensure_ascii=False, indent=2)}
|
||
|
||
请基于你的人格特质和记忆，回答以下问题：
|
||
{query}
|
||
|
||
要求：
|
||
1. 保持人格一致性
|
||
2. 参考历史记忆
|
||
3. 回答控制在100字以内
|
||
4. 体现你的独特风格
|
||
"""
|
||
|
||
response = await model.generate_content_async(prompt)
|
||
return response.text
|
||
|
||
def build_context(self, immortal_name: str, memories: List[Dict]) -> str:
|
||
"""构建上下文信息"""
|
||
context_parts = []
|
||
|
||
# 添加人格基线
|
||
baseline = self.immortal_baselines[immortal_name]
|
||
context_parts.append(f"人格类型: {baseline['mbti_type']}")
|
||
context_parts.append(f"核心特质: {json.dumps(baseline['core_traits'], ensure_ascii=False)}")
|
||
|
||
# 添加相关记忆
|
||
if memories:
|
||
context_parts.append("相关记忆:")
|
||
for memory in memories[-3:]: # 最近3条记忆
|
||
context_parts.append(f"- {memory['content'].get('message', '')}")
|
||
|
||
return "\n".join(context_parts)
|
||
|
||
def simulate_conversation(self, immortal_name: str, messages: List[str]):
|
||
"""模拟对话，测试记忆功能"""
|
||
print(f"\n🎭 开始与 {immortal_name} 的对话")
|
||
print("=" * 50)
|
||
|
||
for i, message in enumerate(messages):
|
||
print(f"\n用户: {message}")
|
||
|
||
# 生成回应
|
||
response = asyncio.run(self.generate_immortal_response(immortal_name, message))
|
||
print(f"{immortal_name}: {response}")
|
||
|
||
# 存储记忆
|
||
self.store_memory(immortal_name, "conversation", {
|
||
"user_message": message,
|
||
"immortal_response": response,
|
||
"session_id": f"session_{i}"
|
||
})
|
||
|
||
# 存储观点
|
||
if "看多" in response or "看空" in response or "观望" in response:
|
||
viewpoint = "看多" if "看多" in response else "看空" if "看空" in response else "观望"
|
||
self.store_memory(immortal_name, "viewpoint", {
|
||
"symbol": "TSLA", # 假设讨论特斯拉
|
||
"viewpoint": viewpoint,
|
||
"reasoning": response
|
||
})
|
||
|
||
def analyze_memory_evolution(self, immortal_name: str):
|
||
"""分析记忆演化"""
|
||
if immortal_name not in self.memory_bank:
|
||
print(f"❌ {immortal_name} 没有记忆数据")
|
||
return
|
||
|
||
memory_data = self.memory_bank[immortal_name]
|
||
|
||
print(f"\n📊 {immortal_name} 记忆分析")
|
||
print("=" * 50)
|
||
print(f"记忆空间创建时间: {memory_data['created_at']}")
|
||
print(f"最后更新时间: {memory_data['last_updated']}")
|
||
print(f"对话记录数: {len(memory_data['conversation_history'])}")
|
||
print(f"观点演化数: {len(memory_data['viewpoint_evolution'])}")
|
||
print(f"决策记录数: {len(memory_data['decision_history'])}")
|
||
|
||
# 分析观点演化
|
||
if memory_data['viewpoint_evolution']:
|
||
print(f"\n观点演化轨迹:")
|
||
for i, viewpoint in enumerate(memory_data['viewpoint_evolution']):
|
||
print(f" {i+1}. {viewpoint['content']['viewpoint']} - {viewpoint['timestamp']}")
|
||
|
||
def save_memory_bank(self, filename: str = "memory_bank_backup.json"):
|
||
"""保存记忆库到文件"""
|
||
with open(filename, 'w', encoding='utf-8') as f:
|
||
json.dump(self.memory_bank, f, ensure_ascii=False, indent=2)
|
||
print(f"💾 记忆库已保存到 {filename}")
|
||
|
||
def load_memory_bank(self, filename: str = "memory_bank_backup.json"):
|
||
"""从文件加载记忆库"""
|
||
try:
|
||
with open(filename, 'r', encoding='utf-8') as f:
|
||
self.memory_bank = json.load(f)
|
||
print(f"📂 记忆库已从 {filename} 加载")
|
||
except FileNotFoundError:
|
||
print(f"⚠️ 文件 {filename} 不存在，使用空记忆库")
|
||
|
||
def main():
|
||
"""主实验函数"""
|
||
print("🚀 开始 Memory Bank 实验")
|
||
print("=" * 60)
|
||
|
||
# 创建实验实例
|
||
experiment = MemoryBankExperiment()
|
||
|
||
# 测试对话场景
|
||
test_scenarios = {
|
||
"吕洞宾": [
|
||
"你觉得特斯拉股票怎么样？",
|
||
"现在市场波动很大，你怎么看？",
|
||
"你之前不是看好特斯拉吗？现在还是这个观点吗？"
|
||
],
|
||
"何仙姑": [
|
||
"特斯拉股票风险大吗？",
|
||
"现在适合投资吗？",
|
||
"你一直很谨慎，现在还是建议观望吗？"
|
||
],
|
||
"张果老": [
|
||
"从历史数据看，特斯拉表现如何？",
|
||
"现在的估值合理吗？",
|
||
"你之前分析过特斯拉的历史数据，现在有什么新发现？"
|
||
]
|
||
}
|
||
|
||
# 执行实验
|
||
for immortal_name, messages in test_scenarios.items():
|
||
experiment.simulate_conversation(immortal_name, messages)
|
||
experiment.analyze_memory_evolution(immortal_name)
|
||
|
||
# 保存记忆库
|
||
experiment.save_memory_bank()
|
||
|
||
print("\n🎉 Memory Bank 实验完成！")
|
||
print("=" * 60)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|